from __future__ import print_function,division
import os
import collections

# prepare logging
import logging
logfname = os.path.splitext(__file__)[0] + ".log"
logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s',
                    datefmt='%y-%m-%d %H:%M:%S',
                    filename=logfname,
                    filemode='w')
import numpy as np
import pylab as plt
import mcutils as mc
import mcutils.xray as xray
from mcutils.xray import id9
#id9 = xray.id9

# use npz files (they can handle more stuff (list of arrays,unicode) than h5py)
g_default_ext = '.npz'

id9.default_extension = g_default_ext
g_default_mask = '../masks/fesalen1_run8_careful.edf'
g_default_folder = "../fesalen/fesalen1/run%d"

runsE = collections.OrderedDict()
runsT = collections.OrderedDict()
runsE[7]  =  60; runsT[7]  = 135
runsE[8]  =  60; runsT[8]  = 135
runsE[9]  = 120; runsT[9]  = 135
runsE[10] = 240; runsT[10] = 135
runsE[11] =  30; runsT[11] = 135
runsE[12] =  45; runsT[12] = 135
runsE[13] =  60; runsT[13] = 120
runsE[14] = 120; runsT[14] = 120
runsE[15] =  30; runsT[15] = 120


runs135=(11,12,8,9,10)
runs120=(15,13,14)

def prepareLog():
  """ It allows printing to terminal on top of logfile """
  # define a Handler which writes INFO messages or higher to the sys.stderr
  console = logging.StreamHandler()
  console.setLevel(logging.INFO)
  # set a format which is simpler for console use
  formatter = logging.Formatter('%(message)s')
  # tell the handler to use this format
  console.setFormatter(formatter)
  # add the handler to the root logger (if needed)
  if len(logging.getLogger('').handlers)==1:
    logging.getLogger('').addHandler(console)

prepareLog()

def defineCalibrants(dmin=3):
  c120 = xray.cell.Triclinic(a=10.297,b=10.716,c=13.955,alpha=72.22,beta=70.45,gamma=74.13)
  p120,_ = c120.reflection_list()
  c300 = xray.cell.Triclinic(a=10.485,b=11.015,c=14.280,alpha=74.61,beta=69.29,gamma=75.29)
  p300,_ = c300.reflection_list()
  return xray.storage.DataStorage( p120=p120, c120=c120,p300=p300,c300=c300) 

def readCalculations(folder="../fesalen/cif",force=False,plot=False):
  npz = folder + "/fesalen_calc.npz"
  if not os.path.isfile(npz) or force:
    ret = dict()
    for T in ("120LS","200HS","300HS"):
      # calculated with lambda = 1 Ang
      theta,i = np.loadtxt( folder+"/fesalen_%s_fwhm0.1.tsv"%T,unpack=True )
      theta = theta/180*np.pi
      q = 4*np.pi/1*np.sin(theta/2)
      key = "T%s"%T 
      ret[ key ] = dict( q = q, i = i )
    ret = xray.storage.DataStorage(ret)
    ret.save(folder + "/fesalen_calc.npz" )
  ret = xray.storage.read(npz)
  if plot:
    T = list(ret.keys())
    T.sort()
    for i,t in enumerate(T): 
      plt.plot( ret[t].q,ret[t].i, label=t)#,color=plt.cm.Blues((0.3+0.7*i/2)) )
    plt.grid()
    plt.legend()
    plt.xlabel(r"q ($\AA^{-1}$)")
  return ret
    

def findCenter(cen=(484,479.5),dist=0.251):
  files = xray.utils.getEdfFiles("../fesalen/fesalen1/run17/",nFiles=100)
  img   = xray.azav.read(files).mean(axis=0) - 10.
  wavelength=12.398/18.*1e-10
  # 1,0,0 peak are the one close to azimuth 0 and are at 0.6596Ang-1 at 120K
  # and 0.65098 at 300K (run 17 is @ RT)
  xray.azav.find_center(img,dist=dist,psize=0.0001770834,center=cen,
                        reference=0.65098,wavelength=wavelength)
  



def removeBaseline(folder,qlims=(0.5,2.5),max_iter=30,force=False):
  if isinstance(folder,int): folder = g_default_folder%folder
  fname    = folder + "/pyfai_1d_nobaseline"+g_default_ext
  if not os.path.isfile(fname) or force:
    logging.info("Calling azav with default parameters")
    data = azav(folder,withBaseline=True)
    # first peak (around 0.6) is weak and requires special care ...
    data.q,data.data = xray.utils.removeBackground(data.q,data.data,
      xlims=qlims,max_iter=max_iter,background_regions=[ (0,0.58) ])
    data.save(fname)
  else:
    data = xray.storage.read(fname)
  return data

def azav(folder,nQ=3000,force=False,saveChi=True,mask=g_default_mask,
         withBaseline=True):
  if isinstance(folder,int): folder = g_default_folder%folder
  if withBaseline:
    data = id9.doFolder_azav(folder,nQ=nQ,force=force,mask=mask,
      saveChi=saveChi)
  else:
    data=removeBaseline(folder)
  return data

def datared(folder,monitor=(0.5,4),showPlot=True,errMask=5,chi2Mask=1.5,
            qlims=(0.5,3),withBaseline=False,forceDatared=False,
            select=slice(None),**kw):
  if isinstance(folder,int): folder = g_default_folder%folder
  data = azav(folder,withBaseline=withBaseline)
  if withBaseline:
    diffFile = folder + "/diffs" + g_default_ext
  else:
    diffFile = folder + "/diffs_nobaseline" + g_default_ext

  if not os.path.isfile(diffFile) or forceDatared:
    data,diffs = id9.doFolder_dataRed(data,monitor=monitor,
                 errMask=errMask,chi2Mask=chi2Mask,qlims=qlims,
                 outStorageFile=diffFile)
  else:
    diffs = xray.storage.read(diffFile)

  if showPlot:
    xray.utils.plotdiffs(diffs.q,diffs.data,t=diffs.scan,select=select,
         absSignal=diffs.dataAbsAvNeg,absSignalScale=10)
    plt.title(folder + " norm %s" % str(monitor))
    plt.figure()
    xray.utils.plotdiffs(diffs.q,diffs.dataAsAbs,t=diffs.scan,
         select=select)
    plt.title(folder + " norm %s" % str(monitor))
  return data,diffs


def doall(folder,force=False,removeBaseline=True):
  azav(folder,force=force)
  return datared(folder)

def _getAmplitudeFromSVD(i,comp=0):
  u,s,v = np.linalg.svd(i)
  return u[:,comp]*s[comp]

def _getAmplitudeFromLinFit(i,ref=1):
  amp = [mc.linFit(i[ref],ii) for ii in i]
  return np.asarray(amp)

def _getAmplitudeFromAbs(i,e):
  amp = np.abs(i).mean(axis=1)
  baseline = np.median(e)
  return amp-baseline

def _getAmplitudeFromPeak(i):
  """ get max-min (to be less sensitive to baseline change """
  amp = np.max(i,axis=1)-np.min(i,axis=1)
  return amp

def anaAmplitue(run=12,plot=False,scaleWithE=False,withBaseline=False):
  folder = g_default_folder % run
  pars = dict( chi2Mask=1.5,errMask=3,monitor=(0.4,2.5),showPlot=False )
  data,diffs=datared(folder,**pars,withBaseline=withBaseline)
  ranges = ( (0.66,0.72), (0.74,0.8), (0.82,0.88) ) #, (0.96,1.02 ) )
  nPlot = len(ranges)
  if plot: fig,ax = plt.subplots(nPlot,1,sharex=True)
  amp = collections.OrderedDict()
  for num,r in enumerate(ranges):
    idx = (diffs.q>r[0]) & (diffs.q<r[1])
    #amplitude = _getAmplitudeFromLinFit( diffs.data[:,idx] )
    amplitude = _getAmplitudeFromAbs( diffs.data[:,idx],diffs.err[:,idx] )
    if scaleWithE: amplitude = amplitude/runsE[run]
    #amplitude = _getAmplitudeFromPeak( diffs.data[:,idx]  )
    amp[r] = amplitude
    if plot:
      ax[num].plot(diffs.scan,amplitude,'-o')
      ax[num].set_title("Range %s"%(str(r)))
  return diffs.scan,amp

def _getPeakPos(x,y,e=1):
  if y.ndim == 1: y = (y,)
  ret = dict()
  for yy,ee in zip(y,e):
    res = xray.peaks.fitPeak(x,yy,err=ee)
    #plt.figure(); res.plot_fit(); input("ok")
    pars = res.params.keys()
    for p in pars:
      if not p in ret: ret[p] = dict( value = [], err = [] )
      ret[p]['value'].append( res.params[p].value )
      ret[p]['err'].append(   res.params[p].stderr )
  for p in pars:
    ret[p]['value'] = np.asarray(ret[p]['value'])
    ret[p]['err'] = np.asarray(ret[p]['err'])
  return xray.storage.DataStorage(ret)

def anaPeakPos(run=12,plot=False,withBaseline=False):
  folder = g_default_folder % run
  pars = dict( chi2Mask=1.5,errMask=3,monitor=(0.4,2.5),showPlot=False )
  data,diffs=datared(folder,**pars,withBaseline=withBaseline)
  ranges = ( (0.66,0.72), (0.74,0.8), (0.82,0.88) ) #, (0.96,1.02 ) )
  nPlot = len(ranges)
  if plot: fig,ax = plt.subplots(nPlot,1,sharex=True)
  pars = collections.OrderedDict()
  for num,r in enumerate(ranges):
    idx = (diffs.q>r[0]) & (diffs.q<r[1])
    #amplitude = _getAmplitudeFromLinFit( diffs.data[:,idx] )
    pars[r] = _getPeakPos( data.q[idx],diffs.dataAsAbs[:,idx],diffs.err[:,idx])
    if plot:
      ax[num].axhline( pars[r].center.value[0] )
      ax[num].plot(diffs.scan,pars[r].center.value,'-o')
      ax[num].set_title("Range %s"%(str(r)))
  return diffs.scan,pars


#### PLOTS #####

def compareT():
  fig,ax = plt.subplots(2,1,sharex=True)
  plt.sca( ax[0] )
  readCalculations(plot=True)
  plt.sca( ax[1] )
  t120 = removeBaseline("../fesalen/fesalen1/run16")
  t300 = removeBaseline("../fesalen/fesalen1/run17")
  plt.plot(t120.q,t120.data.mean(0),label="Data 120K")
  plt.plot( [0,.01],[0,.01] ) # just to have a matching color
  plt.plot(t300.q,t300.data.mean(0),label="Data 300K")

def showBaseline(N=2):
  data_bs = azav("../fesalen/fesalen1/run17")
  data    = removeBaseline("../fesalen/fesalen1/run17")
  idx1 = np.argwhere( data.q[0] == data_bs.q )[0][0]
  idx2 = np.argwhere( data.q[-1] == data_bs.q )[0][0]
  bs   = data_bs.data[:,idx1:idx2+1]-data.data
  fig,ax=plt.subplots(2,N,sharex=True,sharey='row',squeeze=False)
  ax = ax.T
  for i,a in enumerate(ax):
    a[0].plot(data_bs.q,data_bs.data[i])
    a[0].plot(data.q,bs[i])
    a[1].plot(data.q,data.data[i])

def showBaselineStability(N=100):
  # has to be done with cryo on
  data_bs = azav("../fesalen/fesalen1/run15",withBaseline=True)
  data    = azav("../fesalen/fesalen1/run15",withBaseline=False)
  fig,ax=plt.subplots(2,2,sharey=True)
  i = np.arange(data.data.shape[0])
  ax[0,0].pcolormesh(i,data_bs.q,data_bs.data.T)
  ax[0,0].set_title("Without baseline subtraction")
  ax[0,0].set_ylabel(r"q ($\AA^{-1}$)")
  for x in range(N):
    ax[0,1].plot(data_bs.data[x],data_bs.q,lw=0.5,color="0.5",alpha=0.1)
  ax[1,0].pcolormesh(i,data.q,data.data.T)
  ax[1,0].set_xlabel("Image number")
  ax[1,0].set_title("With baseline subtraction")
  ax[1,0].set_ylabel(r"q ($\AA^{-1}$)")
  for x in range(N):
    ax[1,1].plot(data.data[x],data.q,lw=0.5,color="0.5",alpha=0.1)
  ax[0,1].set_xlim(0,400)
  ax[1,1].set_xlim(0,180)
  ax[1,1].set_xlabel("Intensity")
  [a.set_xlim(0,400) for a in ax[:,0]]
  plt.ylim(0.4,3)

def comparisonTRbaseline(folder="../fesalen/fesalen1/run9",select=slice(None)):
   pars = dict( chi2Mask=1.5,errMask=3,monitor=(0.4,2.5),showPlot=False )
   data,diffs_bs=datared(folder,**pars,withBaseline=True)
   data,diffs=datared(folder,**pars,withBaseline=False)
   fig,ax=plt.subplots(2,1,sharex=True)
   plt.sca(ax[0])
   xray.utils.plotdiffs(diffs.q,diffs.data,t=diffs.scan,
         absSignal=diffs.dataAbsAvAll,absSignalScale=10,select=select)
   plt.title(folder + " norm %s" % str(pars['monitor']))

   plt.sca(ax[1])
   xray.utils.plotdiffs(diffs_bs.q,diffs_bs.data,t=diffs_bs.scan,
         absSignal=diffs_bs.dataAbsAvAll,absSignalScale=30,select=select)
 

def compareAmplitudes(runs=(11,12,8,9,10),scaleWithE=True,withBaseline=False):
  runs=list(runs)
  nRuns = len(runs)
  fig,ax=plt.subplots(nRuns,3,sharex=True,sharey='col',figsize=[12,8])
  subplots_kw = dict( hspace=0,wspace=0 )
  plt.subplots_adjust( **subplots_kw )
  plt.rcParams['font.size']=14
  for run,arow in zip(runs,ax):
    t,amp = anaAmplitue(run,plot=False,scaleWithE=scaleWithE,withBaseline=withBaseline)
    for k,a in zip(amp.keys(),arow):
      a.axhline(0,ls='--')
      a.plot(t,amp[k],'o')
  for r,a in zip(runs,ax[:,0]):
    label = "run %d (%.0f K)\n%.0f uJ"%(r,runsT[r],runsE[r])
    a.set_ylabel(label)
  for interval,a in zip(amp.keys(),ax[0,:]): a.set_title("range %s"%(str(interval)))

def compareAmplitudes2(runs=(11,12,8,9,10),scaleWithE=False,withBaseline=False):
  runs=list(runs)
  nRuns = len(runs)
  fig,ax=plt.subplots(1,3,sharex=True,figsize=[12,8])
  subplots_kw = dict( hspace=0,wspace=0 )
  plt.subplots_adjust( **subplots_kw )
  plt.rcParams['font.size']=14
  for a in ax: a.axhline(0,ls='--')
  for run in runs:
    t,amp = anaAmplitue(run,plot=False,scaleWithE=scaleWithE,withBaseline=withBaseline)
    for k,a in zip(amp.keys(),ax):
      label = "run %d %.0f uJ"%(run,runsE[run])
      a.plot(t,amp[k],'-o',label=label)
#  for r,a in zip(runs,ax[:,0]):
#    a.set_ylabel(label)
  for interval,a in zip(amp.keys(),ax): a.set_title("range %s"%(str(interval)))


def comparePosition2(runs=(11,12,8,9,10),withBaseline=False,useRatio=True,scaleWithE=False):
  runs=list(runs)
  nRuns = len(runs)
  fig,ax=plt.subplots(1,3,sharex=True,figsize=[12,8],sharey=useRatio)
  subplots_kw = dict( hspace=0,wspace=0 )
  plt.subplots_adjust( **subplots_kw )
  plt.rcParams['font.size']=14
  if scaleWithE and not ratio:
    logging.warn("scaleWithE requires useRatio, setting scaleWithE to False")
    scaleWithE = Fale
  for run in runs:
    t,pars = anaPeakPos(run,plot=False,withBaseline=withBaseline)
    for k,a in zip(pars.keys(),ax):
      label = "run %d %.0f uJ"%(run,runsE[run])
      if useRatio:
        ratio = pars[k].center.value[1:]/pars[k].center.value[0]
        err   = pars[k].center.err[1:]/pars[k].center.value[0]
        if scaleWithE: ratio = (ratio-1)/runsE[run]+1
        line = a.errorbar(t[1:],ratio,err,fmt='-o',label=label)[0]
      else:
        line = a.plot(t[1:],pars[k].center.value[1:],'-o',label=label)[0]
        a.axhline(pars[k].center.value[0],color=line.get_color())
  for interval,a in zip(pars.keys(),ax): a.set_title("range %s"%(str(interval)))
  if useRatio:
    ax[0].set_ylabel("Peak Position / Peak Position(t=-5ns)")
  else:
    ax[0].set_ylabel("Peak Position ($\AA^{-1}$)")