import os
import sys
import numpy as np
np.warnings.simplefilter('ignore')
import time
import matplotlib.pyplot as plt
import h5py
import re

from x3py import x3py
import alignment
import mcutils as mc

kw_2dplot = dict( 
  interpolation = "none",
  aspect        = "auto",
  cmap          = plt.cm.viridis
  )


g_exp   = "mecl3616"
g_exp   = "xppl3716"
g_bml   = g_exp[:3]

x3py.config.updateBeamline(g_bml)

g_folder_init = g_exp+"_init_pars/"
g_folder_out  = g_exp+"_output/"
g_folder_data = "/reg/d/psdm/"+g_bml+"/"+ g_exp +"/hdf5/"
import socket
hostname = socket.gethostname()
if hostname == "x1":
  g_folder_data = "/home/marco/temp"

# set defaults based on experiment
if g_bml == "xpp":
  g_roi_height = 200
  g_swapx = False
  g_swapy = False
else:
  g_roi_height = 100
  g_swapx = True
  g_swapy = False

print("Working on experiment",g_exp,"(beamline %s)"%g_bml)
print(" folder data      →",g_folder_data)
print(" folder init_pars →",g_folder_init)
print(" folder outout    →",g_folder_out)

#g_folder = "/reg/d/psdm/xpp/xppl3716/ftc/hdf5/"

def readDataset(fnameOrRun=7,
  force=False,
  doBkgSub=False):
  if isinstance(fnameOrRun,str) and (fnameOrRun[-3:]=="npz"):
    d = x3py.toolsVarious.DropObject()
    temp = np.load(fnameOrRun)
    spec1 = temp["spec1"]
    spec2 = temp["spec2"]
    nS    = spec1.shape[0]
    d.spec1 = x3py.toolsDetectors.wrapArray("spec1",spec1,time=np.arange(nS))
    d.spec2 = x3py.toolsDetectors.wrapArray("spec2",spec2,time=np.arange(nS))
  else:
    if isinstance(fnameOrRun,int):
      fnameOrRun=g_folder_data+"/"+g_exp+"-r%04d.h5" % fnameOrRun
    d = x3py.Dataset(fnameOrRun,detectors=["opal0","opal1","fee_spec","opal2"])
    if g_bml == "xpp":
      d.spec1 = d.opal0
      d.spec2 = d.opal1
    else:
      d.spec1 = d.fee_spec
      d.spec2 = d.opal2
  if not hasattr(d,"scan"):
    d.scan = x3py.toolsVarious.DropObject()
    d.scan.scanmotor0_values = [0,]
  return d


def getCenter(img,axis=0,threshold=0.05):
  img = img.copy()
  img[img<img.max()*threshold] = 0
  if axis == 1: img=img.T
  p   = img.mean(1)
  x   = np.arange(img.shape[0])
  return int(np.sum(x*p)/np.sum(p))

def showShots(im1,im2):
  nS = im1.shape[0]
  fig,ax = plt.subplots(2,nS,sharex=True,sharey=True)
  if im1.ndim == 3:
    for a,i1,i2 in zip(ax.T,im1,im2):
      a[0].imshow(i1.T,**kw_2dplot)
      a[1].imshow(i2.T,**kw_2dplot)
  else:
    for a,p1,p2 in zip(ax.T,im1,im2):
      a[0].plot(p1)
      a[1].plot(p2)
  
class AnalyzeRun(object):
  def __init__(self,run,initAlign="auto",swapx=g_swapx,swapy=g_swapy):
    """ swapx → swap x axis of first spectrometer
        swapy → swap y axis of first spectrometer
        initAlign: use None if you want default transformation parameters
    """
    self.d = readDataset(run)
    if isinstance(run,str):
      run = int( re.search("\d{3,4}",run).group() )
    self.run = run
    self.results = dict()
    self.swap = (swapx,swapy)
    #self.clearCache()

    d = self.d
    self.spec1 = d.spec1 ; # spec1 is the one that is moved
    self.spec2 = d.spec2 ; 

    try:
      self.loadTransform(initAlign)
    except (AttributeError,FileNotFoundError):
      if initAlign is None:
        print("Set to default transform")
        self.initAlign = self.setDefaultTransform()
      else:
        self.initAlign  = initAlign

  def getShot(self,shot=0,calib=None,bkgSub="line",roi=g_roi_height):
    # read data
    im1 = self.spec1.getShots(shot,calib=calib)
    im2 = self.spec2.getShots(shot,calib=calib)
    # subtractBkg bkg
    im1 = alignment.subtractBkg(im1,bkg_type=bkgSub)
    im2 = alignment.subtractBkg(im2,bkg_type=bkgSub)
    # rebin and swap im1 if necessary
    if im1.shape[-1] != 1024:
      im1 = mc.rebin(im1, (im1.shape[0],im1.shape[1],1024) )
    if self.swap[0]:
      im1 = im1[:,:,::-1]
    if self.swap[1]:
      im1 = im1[:,::-1,:]
    if roi is None:
      pass
    elif isinstance(roi,slice):
      im1 = im1[:,roi,:]
      im2 = im2[:,roi,:]
    elif isinstance(roi,int):
      if not hasattr(self,"roi1"): self.roi1 = alignment.findRoi(im1[0],roi)
      if not hasattr(self,"roi2"): self.roi2 = alignment.findRoi(im2[0],roi)
      im1 = im1[:,self.roi1,:]; im2 = im2[:,self.roi2,:]
    return im1,im2

  def guiAlign(self,shot=0,save="auto"):
    im1,im2 = self.getShot(shot)
    gui = alignment.GuiAlignment(im1[0],im2[0])
    input("Enter to start")
    gui.start()
    if save == "auto":
      fname = g_folder_init+"/run%04d_gui_align.npy" % self.run
    else:
      fname = save
    self.initAlign = gui.transform
    gui.save(fname)


  def analyzeScan(self,initpars=None,nShotsPerCalib=20,nC=None,doFit=False,fitEveryCalib=False):
    """ use xhift = None; in this way the fit routine does not try to automatically find the translationx parameter """
    if initpars is None: initpars= self.initAlign
    d = self.d
    if nC is None: nC = d.opal1.nCalib
    shots = slice(nShotsPerCalib)
    for i in range(nC):
      s1,s2 = self.getShot(shots,calib=i)
      if fitEveryCalib is not False:
        res = alignment.doShots(s1[:fitEveryCalib],s2[:fitEveryCalib],doFit=True,\
                initpars=initpars); #
        idx = np.argmin( [p.fom for p in res] )
        res = res[idx]
        initpars = res.final_pars; self.initAlign=res.final_pars
        print("Calib cycle %d -> %.3f aligned (best FOM: %.2f)" % (i,self.d.scan.scanmotor0_values[i],res.fom))
      ret = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit)
      res = alignment.unravel_results(ret)
      self.results[i] = res
    return self.results.values()

  def doShot(self,shot=0,calib=None,initpars=None,im1=None,im2=None,doFit=True,show=False,showInit=False,save=False,savePlot="auto"):
    if initpars is None: initpars= self.initAlign
    if (im1 is None) or (im2 is None):
      im1,im2 = self.getShot(shot,calib=calib); im1=im1[0]; im2=im2[0]
    r = alignment.doShot(im1,im2,initpars,doFit=doFit,show=showInit)
    im1 = r.im1
    im2 = r.im2
    self.initAlign = r.final_pars
    if show:
      if savePlot == "auto":
        savePlot = g_folder_out+"/run%04d_calib%s_shot%04d_fit.png" % (self.run,calib,shot)
      alignment.plotShot(im1,im2,res=r,save=savePlot)
    if save: self.saveTransform()
    return r

  def doShots(self,shots=slice(0,50),calib=None,initpars=None,doFit=False,returnBestTransform=False,unravel=True):
    if initpars is None: initpars= self.initAlign
    d = self.d
    s1,s2 = self.getShot(shots,calib=calib)
    ret,t = alignment.doShots(s1,s2,initpars=initpars,doFit=doFit,\
            returnBestTransform=True)
    if doFit:
      self.initAlign = t
    ret_unravel = alignment.unravel_results(ret)
    # keep it for later !
    self.results[calib] = ret_unravel
    if unravel: ret = ret_unravel
    if returnBestTransform:
      return ret,t
    else:
      return ret

  def save(self,fname="auto",overwrite=False):
    if fname == "auto":
      fname = g_folder_out+"/run%04d_analysis.h5" % self.run
    if os.path.exists(fname) and not overwrite:
      print("File %s exists, **NOT** saving, use overwrite=True is you want ..."%fname)
      return
    if os.path.exists(fname) and overwrite:
      os.unlink(fname)
    print("Saving results to %s"%fname)
    h = h5py.File(fname)
    h["roi1"] = (self.roi1.start,self.roi1.stop)
    h["roi2"] = (self.roi2.start,self.roi2.stop)
    #h["transform"] = self.initAlign
    for (c,v) in self.results.items():
      cname = "calib%04d/" % c if isinstance(c,int) else "calib%s/" % c
      for p,vv in v.items():
        if p == "parameters":
          for pname,parray in vv.items():
            name = cname + p + "/" + pname
            h[name] = parray
        else:
          h[cname + p] = vv
    h.close()


  def saveTransform(self,fname="auto",transform=None):
    if transform is None: transform = self.initAlign
    if fname == "auto":
      fname = g_folder_init+"/run%04d_transform.npy" % self.run
    print("Saving roi and transformation parameter to %s"%fname)
    alignment.saveAlignment(fname,self.initAlign,self.roi1,self.roi2)

  def loadTransform(self,fname="auto"):
    if fname == "auto":
      fname = g_folder_init+"/run%04d_transform.npy" % self.run
    elif isinstance(fname,int):
      fname = g_folder_init+"/run%04d_transform.npy" % fname
    temp = np.load(fname).item()
    self.initAlign = temp["transform"]
    self.roi1 = temp["roi1"]
    self.roi2 = temp["roi2"]
    print("init transform and ROIs from %s"%fname)


  def clearCache(self):
    del self.roi1 
    del self.roi2
    alignment.clearCache(); # nedded for multiprocessing can leave bad parameters in the cache

  def setDefaultTransform( self ):
    #dict( scalex=0.65,rotation=0.0,transx=90, iblur1=4.3,fix_iblur1=False )
    t = alignment.g_fit_default_kw
    self.initAlign = t
    return t

def quick_mec(run,ref=236,divideByRef=False,returnRes=False):
  """ useful to analyze the runs around 140 (done with the focusing """
  ref_run = 236
  h=h5py.File("mecl3616_output/run%04d_analysis.h5" %ref,"r")
  ref = np.nanmean(h["calibNone"]["ratio"][...],axis=0)
  r = AnalyzeRun(run,initAlign=ref,swapx=True,swapy=False)
  res=r.doShots(slice(5),doFit=False)
  ret = res["ratio"]/ref if divideByRef else res["ratio"]
  if returnRes:
    return ret,res
  else:
    return ret

def quickAndDirty(run,nShots=300,returnAll=True,doFit=False):
  """ useful to analyze the runs around 140 (done with the focusing """
  r = AnalyzeRun(run,swap=True,initAlign=g_folder_init+"/run0144_transform.npy")
  res=r.doShots(slice(nShots),doFit=doFit)
  o = alignment.unravel_results(res)
  ref = np.nanmedian(o["ratio"][:40],0)
  sam = np.nanmedian(o["ratio"][50:],0)
  if returnAll:
    return sam/ref,o["ratio"]/ref
  else:
    return sam/ref