import numpy as np
import matplotlib.pyplot as plt
import copy
import argparse
import collections
import xanes_analyzeRun

parser = argparse.ArgumentParser(description='Process argv')

parser.add_argument('--run', type=int,default=82,help='which run to analyze')
parser.add_argument('--force', action="store_true",help='force calculation')
args = parser.parse_args()


profile_ret = collections.namedtuple("profile_ret",["run","p1","p2","calibs"])


nice_colors = ["#1b9e77", "#d95f02", "#7570b3"]

def calcProfilesForRun(run=82,calibs="all",realign=False,init="auto",alignCalib=0):
  """ init and alignCalib are used only if realignment is performed:
      init = run for initial alignment, use auto is you want to use same run
      alignCalib = calibcycle for alignment
  """
  if init=="auto": init=run
  if isinstance(run,int):
    r = xanes_analyzeRun.AnalyzeRun(run,initAlign=init)
  else:
    r = run
  if realign:
    r.doShots(slice(20),calib=refCalib,doFit=True)
    r.saveTransform()
  # next line is used to force calculations in case of realignment
  fname = 'auto' if not realign else "thisfiledoesnotexists"
  if len(r.results) == 0:
    try:
      r.load(fname)
      print("Loading previously saved results")
    except FileNotFoundError:
      r.analyzeScan(shots=slice(600),calibs=calibs,nImagesToFit=0,nSaveImg=4)
      r.save(overwrite=True)
  # cannot take the output from r.results because it might have been calculated for
  # a bigger range than asked for.
  if isinstance(calibs,int):
    calibsForOut = (calibs,)
  elif isinstance(calibs,slice):
    calibsForOut = list(range(r.nCalib))[calibs]
  elif calibs == "all":
    calibsForOut = list(range(r.nCalib))
  else:
    calibsForOut = calibs
  p1 = [r.results[calib].p1 for calib in calibsForOut]
  p2 = [r.results[calib].p2 for calib in calibsForOut]
  return profile_ret( run =r, p1 =p1, p2=p2,calibs=calibsForOut)


def calcProfilesForRefAndSample(run=82,refCalibs=0,force=False):
  if isinstance(run,int):
    refRun    = xanes_analyzeRun.AnalyzeRun(run)
    sampleRun = refRun
    if isinstance(refCalibs,slice): refCalibs    = list(range(refRun.nCalib))[refCalibs]
    if isinstance(refCalibs,int):   refCalibs    = (refCalibs,)
    sampleCalibs = [c+1 for c in refCalibs]
  elif isinstance(run,(list,tuple)):
    refRun     = xanes_analyzeRun.AnalyzeRun(run[0])
    sampleRun  = xanes_analyzeRun.AnalyzeRun(run[1],initAlign=run[0])
    refCalibs     = [0,]
    sampleCalibs  = [0,]
  if refRun == sampleRun:
    # otherwise same does not save them both
    temp = calcProfilesForRun(refRun,calibs=sampleCalibs+refCalibs);
  ref = calcProfilesForRun(refRun,calibs=refCalibs)
  sample  = calcProfilesForRun(sampleRun,calibs=sampleCalibs)
  return ref,sample

def calcRef(r1,r2,calibs=None,threshold=0.05):
  """ r1 and r2 are list of 2d arrays (nShots,nPixels) for each calibcycle """
  if calibs is None: calibs = list(range(len(r1)))
  out = collections.OrderedDict()
  out["ratioOfAverage"] = dict()
  out["medianOfRatios"] = dict()
  for p1,p2,n in zip(r1,r2,calibs):
    out["ratioOfAverage"][n] = xanes_analyzeRun.ratioOfAverage(p1,p2,threshold=threshold)
    out["medianOfRatios"][n] = xanes_analyzeRun.medianRatio(p1,p2,threshold=threshold)
  # add curves with all calib together
  p1 = np.vstack(r1)
  p2 = np.vstack(r2)
  n = ",".join(map(str,calibs))
  ref1 = xanes_analyzeRun.ratioOfAverage(p1,p2,threshold=threshold)
  ref2 = xanes_analyzeRun.medianRatio(p1,p2,threshold=threshold)
  out["ratioOfAverage"][n] = xanes_analyzeRun.ratioOfAverage(p1,p2,threshold=threshold)
  out["medianOfRatios"][n] = xanes_analyzeRun.medianRatio(p1,p2,threshold=threshold)
  out["ratioOfAverage"]['all'] = out["ratioOfAverage"][n]
  out["medianOfRatios"]['all'] = out["medianOfRatios"][n]
  return out

def showDifferentRefs(run=82,refCalibs=slice(None,None,2),threshold=0.05):
  """ example plots showing how stable are the different ways of taking spectra """
  prof = calcProfilesForRun(run,calibs=refCalibs)
  refs = calcRef(prof.p1,prof.p2,calibs=prof.calibs)
  kind_of_av = list(refs.keys())
  fig,ax=plt.subplots(len(kind_of_av)+1,1,sharex=True,sharey=True)
  E = prof.run.E
  calibs = list(refs[kind_of_av[0]].keys())
  for ikind,kind in enumerate(kind_of_av):
    for calib in calibs:
      if isinstance(calib,int):
        ax[ikind].plot(E,refs[kind][calib],label="calib %s"%calib)
      else:
        if calibs == 'all': continue
        ax[ikind].plot(E,refs[kind][calib],label="calib %s"%calib,lw=2,color='k',alpha=0.7)
        ax[-1].plot(E,refs[kind][calib],label="calib all, %s"%kind,lw=1.5,color=nice_colors[ikind],alpha=0.8)
  for ikind,kind in enumerate(kind_of_av): ax[ikind].set_title("Run %d, %s"%(run,kind)) 
  ax[0].set_ylim(0.88,1.12)
  ax[0].set_ylim(0.88,1.12)
  ax[-2].legend()
  ax[-1].legend()
  for a in ax: a.grid()

def calcSampleAbs(run=82,refCalibs=slice(None,None,2),threshold=0.05,refKind="medianOfRatios"):
  """ example of use
      ratio = calcSampleAbs(82)
      ratio = calcSampleAbs( (155,156) )
  """
  ref,sample = calcProfilesForRefAndSample(run,refCalibs=refCalibs)
  temp  = calcRef(ref.p1,ref.p2,calibs=ref.calibs,threshold=threshold)
  ref   = temp[refKind]['all']
  p1 = np.vstack( sample.p1 )
  p2 = np.vstack( sample.p2 )
  print(p1.shape)
  p1,p2 = xanes_analyzeRun.maskLowIntensity(p1,p2,threshold=0.1)
  ratio = p2/p1
  ratio = ratio/ref
  return ratio
  
  
  

def main(run,refCalib=0,force=False):
  pass
  #r = calcProfiles(run,refCalibs=refCalib,force=force)


if __name__ == "__main__":
  main(args.run,force=args.force)