from __future__ import print_function,division
import os


# 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

from dualtree import dualtree

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

g_default_mask = '../masks/fesalen1_run8_careful.edf'


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
  logging.getLogger('').addHandler(console)

prepareLog()

def findCenter():
  files = xray.utils.getEdfFiles("../fesalen/fesalen1/run8/",nFiles=100)
  img   = xray.azav.read(files).mean(axis=0) - 10.
  xray.azav.find_center(img)
  

def azav(folder,nQ=1500,force=False,saveChi=True,mask=g_default_mask):
  if isinstance(mask,int):
    files = xray.utils.getFiles(folder,"*.edf*")
    img   = xray.azav.pyFAIread(files[0])
    temp = np.ones_like(img,dtype=bool)
    temp[:mask] = False
    mask = temp
  return id9.doFolder_azav(folder,nQ=nQ,force=force,mask=mask,saveChi=saveChi)


def removeBaseline(folder,qlims=(0.6,3),max_iter=30):
  data = azav(folder)
  idx = (data.q>qlims[0]) & (data.q<qlims[1])
  data['q'] = data.q[idx]
  data['data'] = data.data[:,idx]
  for i in range(len(data.data)):
    data['data'][i]=data.data[i]-dualtree.baseline(data.data[i],max_iter=max_iter)
  fname = os.path.splitext(data.filename)[0] + "_nobaseline" + id9.default_extension
  data.save(fname)
  return data

def datared(folder,monitor=(0.5,4),showPlot=True,errMask=5,chi2Mask=2,
            qlims=(0.6,3),withBaseline=False,storageFile='auto',force=False,**kw):
  # ice contributes to a lot of noise, filter to q<2
  if storageFile == 'auto':
    if withBaseline:
      storageFile = folder + "/" + "pyfai_1d" + id9.default_extension
    else:
      storageFile = folder + "/" + "pyfai_1d_nobaseline" + id9.default_extension
      if not os.path.exists(storageFile) or force: removeBaseline(folder)
  data = xray.storage.DataStorage(storageFile)
  data,diffs = id9.doFolder_dataRed(folder,storageFile=data,monitor=monitor,
               errMask=errMask,chi2Mask=chi2Mask,qlims=qlims,**kw)
  if showPlot:
    xray.utils.plotdiffs(diffs.q,diffs.data,t=diffs.scan,
         absSignal=diffs.dataAbsAvAll,absSignalScale=30)
    plt.title(folder + " norm %s" % str(monitor))
    plt.figure()
    xray.utils.plotdiffs(diffs.q,diffs.dataAbsAvScanPoint,t=diffs.scan)
    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 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)))