from __future__ import print_function,division
import numpy as np
import pylab as plt
import os
import mcutils as mc
import mcutils.xray as xray
from mcutils.xray import id9
id9 = xray.id9

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

def readCalc():
  fold = "../tiox/calculated_patterns/"
  names = "alpha beta lam".split()
  fnames = "alpha500K beta290K lambda".split()
  calc = dict()
  for name,fname in zip(names,fnames):
    q,i=np.loadtxt(fold + "%s.xye.q"%fname,unpack=True)
    calc[name] = dict( q=q, i=i ) 
  return xray.storage.DataStorage(calc)

calc = readCalc()

def azav(folder,nQ=1500,force=False,saveChi=True,mask='y>470'):
  return id9.doFolder_azav(folder,nQ=nQ,force=force,mask=mask,saveChi=saveChi)

def datared(folder,monitor=(1,5),forceDatared=False,showPlot=True,**kw):
  azavFile = folder + "/pyfai_1d" + g_default_ext
  diffFile = folder + "/diffs" + g_default_ext
  if not os.path.isfile(diffFile) or forceDatared:
    data,diffs = id9.doFolder_dataRed(folder,monitor=monitor,**kw)
  else:
    data = xray.storage.read(azavFile)
    diffs = xray.storage.read(diffFile)
  if showPlot:
    xray.utils.plotdiffs(diffs.q,diffs.data,t=diffs.scan,
         absSignal=diffs.dataAbsAvAll,absSignalScale=30)
    plt.plot(calc.lam.q,calc.lam.i/1000+0.1,label='calc')
    plt.title(folder + " norm %s" % str(monitor))
  return data,diffs

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

def anaAmplitue(run=6):
  fname = "../tiox/tiox1/run%d/diffs.npz" % run
  data  = xray.storage.DataStorage(fname)
  ranges = ( (1.75,1.85), (2.2,2.4), (3.25,3.4) )
  nPlot = len(ranges)
  fig,ax = plt.subplots(nPlot,1,sharex=True)
  for r,a in zip(ranges,ax):
    idx = (data.q>r[0]) & (data.q<r[1])
    amplitude = np.abs(data.data[:,idx]).mean(axis=1)
    a.plot(data.scan,amplitude,'-o')
    a.set_title("Range %s"%(str(r)))

def plotCalc(scale=1):
  fold = "../tiox/calculated_patterns/"
  q,i=readtxt(fold + "alpha500K.xye.q")
  plt.plot(q,i*scale,label="alpha")
  q,i=readtxt(fold + "beta290K.xye.q")
  plt.plot(q,i*scale,label="beta")
  q,i=readtxt(fold + "lambda.xye.q")
  plt.plot(q,i*scale,label="lambda")