import sys
import zmq
import numpy as np
import time
import pickle
import alignment
import matplotlib.pyplot as plt
import threading
import datetime
import copy
import pyqtgraph as pg

plt.rcParams['image.cmap'] = 'viridis'

def histVec(v,oversample=1):
  v = np.atleast_1d(v)
  v = np.unique(v)
  vd = np.diff(v)
  vd = np.hstack([vd[0],vd])
  #vv = np.hstack([v-vd/2,v[-1]+vd[-1]/2])
  vv = np.hstack([v-vd/2.,v[-1]+vd[-1]/2.])
  if oversample>1:
    vvo = []
    for i in range(len(vv)-1):
      vvo.append(np.linspace(vv[i],vv[i+1],oversample+1)[:-1])
    vvo.append(vv[-1])
    vv = np.array(np.hstack(vvo))
  return vv

def subtractBkg(imgs,nPix=100,dKtype='corners'):
  """ Opals tend to have different backgroud for every quadrant """
  if dKtype is 'corners':
    if imgs.ndim == 2: imgs = imgs[np.newaxis,:]
    imgs = imgs.astype(np.float)
    q1 = imgs[:,:nPix,:nPix].mean(-1).mean(-1)
    imgs[:,:512,:512]-=q1[:,np.newaxis,np.newaxis]
    q2 = imgs[:,:nPix,-nPix:].mean(-1).mean(-1)
    imgs[:,:512,-512:]-=q2[:,np.newaxis,np.newaxis]
    q3 = imgs[:,-nPix:,-nPix:].mean(-1).mean(-1)
    imgs[:,-512:,-512:]-=q3[:,np.newaxis,np.newaxis]
    q4 = imgs[:,-nPix:,:nPix].mean(-1).mean(-1)
    imgs[:,-512:,:512]-=q4[:,np.newaxis,np.newaxis]
  elif dKtype is 'stripes':
    if imgs.ndim == 2: imgs = imgs[np.newaxis,:]
    imgs = imgs.astype(np.float)
    s1 =  imgs[:,:nPix,:].mean(-2)




  return np.squeeze(imgs)

def getData():
  t0 = time.time()
  context = zmq.Context()
  socket = context.socket(zmq.REQ)
  socket.connect('tcp://daq-xpp-mon06:12322')
  #socket.setsockopt(zmq.SUBSCRIBE, b'')
  while True:
    socket.send(b"Request")
    ret    = socket.recv()
    ret = pickle.loads(ret, encoding='latin1')
    print('received',ret.keys(),time.time()-t0)
    t0 = time.time()

if __name__ == "__main__":
  getData()

class SpecHandler(object):
  def __init__(self,connectString='tcp://daq-xpp-mon06:12322',spec1name='opal0',spec2name='opal1',roi=[0,1024,0,1024]):
    self.connectString = connectString
    self.spec1name = spec1name
    self.spec2name = spec2name
    self.dataCollector = []
    #self.surveyplot = Surveyplot(spec1name=spec1name,spec2name=spec2name)
    self.roi = roi
    self.projsimple = ProjSimple(spec1name=spec1name,spec2name=spec2name,dataCollector=self.dataCollector)
    self._rawContinuousTime = None
    self.lastDat = None
    self.openSocket()
    #self.runningPlot = RunningPlot(self.dataCollector)


  def openSocket(self):
    context = zmq.Context()
    self.context = context
    self.socket = context.socket(zmq.REQ)
    self.socket.connect(self.connectString)
    #self.socket.setsockopt(zmq.SUBSCRIBE, b'')

  def closeSocket(self):
    del self.context
    del self.socket

  def getData(self):
    self.socket.send(b"Request")
    ret    = self.socket.recv()
    ret = pickle.loads(ret, encoding='latin1')
    for sn in [self.spec1name,self.spec2name]:
      ret[sn] = np.squeeze(alignment.subtractBkg(ret[sn], nPix=100, bkg_type='line'))
    self.lastDat = ret
    return ret

  def getRaw(self,repoenConnection=False,doAlign=False,show=False,doFit=False,updateImg=True,updateProj=True,updateFom=True):
    if doFit is True: doFit='iminuit'


    if repoenConnection:
      self.closeSocket()
      self.openSocket()
    dat = self.getData()
    im1 = dat[self.spec1name]; im2 = dat[self.spec2name]
    if doAlign:
      #t = np.load("gui_align_transform_xppl3716.npy").item()
      if hasattr(self,'transformer'):
        algn = self.transformer
      else:
        algn = alignment.loadAlignment('last_trafo.npy')
      t = algn['transform']
      roi1 = algn['roi1']
      roi2 = algn['roi2']
      
      r = alignment.doShot( im1[roi1,:],im2[roi2,:],t, show=show, doFit=doFit)
      self.transformer = dict(transform=r.final_transform,roi1=roi1,roi2=roi2)
      alignment.saveAlignment('last_trafo.npy',r.final_transform,roi1,roi2)

      im1 = r.im1; im2 = r.im2
      showDiff = True
      showRatio = True
    else:
      showDiff = False
      showRatio = False
    if doAlign:
      thres = 0.05
      #im1[im1<thres*np.max(im1.ravel())] = np.nan
      #im2[im2<thres*np.max(im2.ravel())] = np.nan
      self.dataCollector.append(\
	  dict( time=datetime.datetime.now(),
	    fom=r.fom, 
	    ratProj = np.nansum(im2/im1,axis=0),
	    im1Proj = np.nansum(im1,axis=0),
	    im2Proj = np.nansum(im2,axis=0)))
	  #if updateFom:
	#self.runningPlot.updatePlot()
    #if updateImg:
      #self.surveyplot.plotImages(im1,im2,showDiff=showDiff,showRatio=showRatio)
    #if updateProj:
    self.projsimple.plotProfiles(im1,im2)
      

  def alignFeatures(self):
    im1 = copy.copy(self.lastDat[self.spec1name])
    im2 = copy.copy(self.lastDat[self.spec2name])
    roi1 = alignment.findRoi(im1)
    roi2 = alignment.findRoi(im2)

    tra = alignment.GuiAlignment(im1[roi1,:],im2[roi2,:],autostart=False)
    self.transformer = dict(transform=tra.start(),roi1=roi1,roi2=roi2)

  
  def getRawContinuuous(self,sleepTime,**kwargs):
    self._rawContinuousTime = sleepTime
    if not hasattr(self,'_rawContinuousThread'):
      def upd():
        while not self._rawContinuousTime is None:
          self.getRaw(**kwargs)
          plt.draw()
          time.sleep(self._rawContinuousTime)

      self._rawContinuousThread = threading.Thread(target=upd)
      self._rawContinuousThread.start()



class Surveyplot(object):
  def __init __(self,spec1name='spec1',spec2name='spec2'):
    self.fig,self.axs = plt.subplots(4,1,sharex=True,sharey=True)
    self.axs[0].set_title(spec1name)
    self.axs[1].set_title(spec2name)
    self.axs[2].set_title("Difference")
    self.axs[3].set_title("Ratio")

  def plotImages(self,img1,img2,showDiff=False,showRatio=False):
    if hasattr(self,'i1'):
      self.i1.set_data(img1)
    else:
      self.i1 = self.axs[0].imshow(img1,origin='lower',interpolation='none')
    if hasattr(self,'i2'):
      self.i2.set_data(img2)
    else:
      self.i2 = self.axs[1].imshow(img2,origin='lower',interpolation='none')
    if showDiff:
      tdiff = img2-img1
      if hasattr(self,'idiff'):
        self.idiff.set_data(tdiff)
      else:
        self.idiff = self.axs[2].imshow(tdiff,interpolation='none',origin='lower')
        lms = np.percentile(tdiff,[30,70])
        self.idiff.set_clim(lms)
    if showRatio:
      tratio = img2/img1
      if hasattr(self,'iratio'):
        self.iratio.set_data(tratio)
      else:
        self.iratio = self.axs[3].imshow(tratio,interpolation='none',origin='lower')
        lms = np.percentile(tratio,[30,70])
        self.iratio.set_clim(lms)

class ProjSimple(object):
  def __init__(self,spec1name='sp ec1',spec2name='spec2',roi1=[0,1024,0,1024],roi2=[0,1024,0,1024],dataCollector=[]):
    self.fig,self.axs = plt.subplots(2,1,sharex=True)
    self.roi1 = roi1
    self.roi2 = roi2
    self.spec1name=spec1name
    self.spec2name=spec2name
    self.dataCollector = dataCollector
  
  #def getROI(self,specNo):
  def _roiit(self,img,roi):
    return img[roi[0]:roi[1],roi[2]:roi[3]]
  
  def plotProfiles(self,img1,img2):
    prof1 = np.nansum(self._roiit(img1,self.roi1),0)
    prof2 = np.nansum(self._roiit(img2,self.roi2),0)
    if hasattr(self,'l1'):
      self.l1.set_ydata(prof1)
    else:
      self.l1 = self.axs[0].plot(prof1,label=self.spec1name)[0]
    if hasattr(self,'l2'):
      self.l2.set_ydata(prof2)
    else:
      self.l2 = self.axs[0].plot(prof2,label=self.spec2name)[0]
      plt.legend()

    if hasattr(self,'lrat'):
      self.lrat.set_ydata(prof2/prof1)
    else:
      self.lrat = self.axs[1].plot(prof2/prof1,'k',label='ratio')[0]
      self.axs[1].set_ylim(0,2)
    if len(self.dataCollector) > 0 :
      im1Proj = np.asarray([i['im1Proj'] for i in self.dataCollector])
      im2Proj = np.asarray([i['im2Proj'] for i in self.dataCollector])
      #print(im1Proj.shape,im2Proj.shape)
      ratAv = np.median(im2Proj,0)/np.median(im1Proj,0)
      if hasattr(self,'lratAv'):
        self.lratAv.set_ydata(ratAv)
      else:
        self.lratAv = self.axs[1].plot(ratAv,'r',label='ratio Avg')[0]
        self.axs[1].set_ylim(0,2)

class ProjSimplePG(object):
  def __init__(self,spec1name='sp ec1',spec2name='spec2',roi1=[0,1024,0,1024],roi2=[0,1024,0,1024],dataCollector=[]):
    from pyqtgraph.Qt import QtGui, QtCore
    import numpy as np
    import pyqtgraph as pg
    self.app = QtGui.Qapplication([])
    self.win = pg.GraphicsWindow(title='Single shot XANES')
    pg.setConfigOptions(antialias=True)
    self.sp1 = win.addPlot(title='spectrometer spectra')
    self.sp2 = win.addPlot(title='Ratio')


    self.fig,self.axs = plt.subplots(2,1,sharex=True)

    self.roi1 = roi1
    self.roi2 = roi2
    self.spec1name=spec1name
    self.spec2name=spec2name
    self.dataCollector = dataCollector
  
  #def getROI(self,specNo):
  def _roiit(self,img,roi):
    return img[roi[0]:roi[1],roi[2]:roi[3]]
  
  def plotProfiles(self,img1,img2):
    prof1 = np.nansum(self._roiit(img1,self.roi1),0)
    prof2 = np.nansum(self._roiit(img2,self.roi2),0)
    if hasattr(self,'l1'):
      self.l1.set_ydata(prof1)
    else:
      self.l1 = self.axs[0].plot(prof1,label=self.spec1name)[0]
    if hasattr(self,'l2'):
      self.l2.set_ydata(prof2)
    else:
      self.l2 = self.axs[0].plot(prof2,label=self.spec2name)[0]
      plt.legend()

    if hasattr(self,'lrat'):
      self.lrat.set_ydata(prof2/prof1)
    else:
      self.lrat = self.axs[1].plot(prof2/prof1,'k',label='ratio')[0]
      self.axs[1].set_ylim(0,2)
    if len(self.dataCollector) > 0 :
      im1Proj = np.asarray([i['im1Proj'] for i in self.dataCollector])
      im2Proj = np.asarray([i['im2Proj'] for i in self.dataCollector])
      #print(im1Proj.shape,im2Proj.shape)
      ratAv = np.median(im2Proj,0)/np.median(im1Proj,0)
      if hasattr(self,'lratAv'):
        self.lratAv.set_ydata(ratAv)
      else:
        self.lratAv = self.axs[1].plot(ratAv,'r',label='ratio Avg')[0]
        self.axs[1].set_ylim(0,2)

class RunningPlot(object):
  def __init__(self,dataCollector):
    self.dataCollector = dataCollector
    self.fig,self.axs = plt.subplots(1,1)
    
  def updatePlot(self):
    if len(self.dataCollector)>0:
      times = np.asarray(([i['time'] for i in self.dataCollector]))
      foms = np.asarray(([i['fom'] for i in self.dataCollector]))
      im1Proj = np.asarray(([i['im1Proj'] for i in self.dataCollector]))
      im2Proj = np.asarray(([i['im2Proj'] for i in self.dataCollector]))

      if hasattr(self,'fomline'):
        self.fomline.set_ydata(foms)
        self.fomline.set_xdata(times)
        self.axs.set_xlim(np.min(times)-datetime.timedelta(0,10),np.max(times)+datetime.timedelta(0,10))

	#self.axs[0].autoscale(enable=True,axis='x')

      else:
        self.fomline = self.axs.plot(times,foms,'o-')[0]
        self.axs.autoscale(enable=True,axis='x')
      #if hasattr(self,'ratioimg'):
	#self.ratioimg.set_data()
	#self.ratioimg.set_ydata(foms)
      #else:
	#self.axs[0].plot(times,foms,'o-')


   

  
  #def plotOrUpdate(img1,img2):
    #if hasattr(self,i1):
      #self.i1.set_data(img1)
    #else:
      #self.i1 = self.axs.imshow(img1,interpolate='none',origin='bottom')
    #if hasattr(self,i1):
      #self.i1.set_data(img1)
    #else:
      #self.i1 = self.axs.imshow(img1,interpolate='none',origin='bottom')