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()

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

def calcProfiles(run=82,refCalib=0,force=False):
  r = xanes_analyzeRun.AnalyzeRun(run)
  fname = "thisfiledoesnotexists" if force else "auto"
  try:
    r.load(fname)
    print("Loading previously saved results")
  except FileNotFoundError:
    r.doShots(shots=slice(20),calib=refCalib,doFit=True)
    r.saveTransform()
    ret  = r.analyzeScan(shots="all",nImagesToFit=0,nSaveImg=4)
    r.save(overwrite=True)
  return r

def calcRef(run,refCalibs=(0,2,4,6,8),threshold=0.05):
  r = calcProfiles(run=run,refCalib=refCalibs[0])
  r1 = [r.results[calib].p1 for calib in refCalibs]
  r2 = [r.results[calib].p2 for calib in refCalibs]
  out = collections.OrderedDict()
  out["ratioOfAverage"] = dict()
  out["medianOfRatios"] = dict()
  for p1,p2,n in zip(r1,r2,refCalibs):
    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,refCalibs))
  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)
  return r.E,out

def showDifferentRefs(run=82,refCalibs=(0,2,4,6,8),threshold=0.05):
  """ example plots showing how stable are the different ways of taking spectra """
  E,refs = calcRef(run=run,refCalibs=refCalibs,threshold=threshold)
  kind_of_av = list(refs.keys())
  fig,ax=plt.subplots(len(kind_of_av)+1,1,sharex=True,sharey=True)
  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:
        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 main(run,refCalib=0,force=False):
  r = calcProfiles(run,refCalib=refCalib,force=force)


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