#!/usr/bin/env python

import os
import sys
import copy
import os
import time
import gc

import numpy as np
from .. import Group
from ..utils import read_textfile
from ..utils.strutils import bytes2str

from . import XDIFile, XDIFileException

from .xsp3_hdf5 import XSPRESS3_TAU, estimate_icr

def read_gsexdi(fname, nmca=128, bad=None, **kws):
    """Read GSE XDI Scan Data to larch group,
    summing ROI data for MCAs and apply deadtime corrections
    """

    MAX_MCAS = nmca
    group = Group()
    group.__name__ ='GSE XDI Data file %s' % fname
    xdi = XDIFile(str(fname))

    group.path = fname
    path, suffix = os.path.split(fname)
    group.filename = suffix
    group.npts = xdi.npts
    group.bad_channels = bad
    for family in ('scan', 'beamline', 'mono', 'facility'):
        for key, val in xdi.attrs.get(family, {}).items():
            if '||' in val:
                val, addr = val.split('||')
            try:
                val = float(val)
            except:
                pass
            setattr(group, "%s_%s" % (family, key), val)

    group._xdi = Group()

    for thing in ('array_addrs', 'array_labels', 'array_units', 'attrs',
                  'comments', 'status', 'user_labels', 'xdi_libversion',
                  'xdi_version', 'xdi_pyversion'):

        val = copy.deepcopy(getattr(xdi, thing, None))
        if isinstance(val, bytes):
            val = val.decode('utf-8')
        setattr(group._xdi, thing, val)

    scanparams = xdi.attrs.get('scanparameters', None)
    if scanparams is not None:
        scan_e0 = scanparams.get('e0', None)
        if scan_e0 is not None:
            group.scan_e0 = float(scan_e0)
        scan_elem = scanparams.get('element', None)
        if scan_elem is not None:
            group.element = scan_elem
        scan_edge = scanparams.get('edge', None)
        if scan_edge is not None:
            group.edge = scan_edge
        scan_type = scanparams.get('scantype', None)
        if scan_type is not None:
            group.scan_type = scan_type

    ocrs, icrs = [], []
    ctime = None
    for attrname in dir(xdi):
        if attrname.lower() == 'counttime':
            ctime = getattr(xdi, attrname)
    if ctime is None:
        try:
            ctime = xdi.TSCALER / 5.e7
        except AttributeError:
            ctime = 1.0

    is_old_xsp3 = any(['13QX4' in a[1] for a in xdi.attrs['column'].items()])

    dtc_mode = 'icr/ocr'
    for i in range(MAX_MCAS):
        mca = "mca%i" % (i+1)
        ocr    = getattr(xdi, 'OutputCounts_%s' % mca, None)
        clock  = getattr(xdi, 'Clock_%s'        % mca, None)
        icr    = getattr(xdi, 'InputCounts_%s'  % mca, None)
        dtfact = getattr(xdi, 'DTFactor_%s'     % mca, None)
        resets = getattr(xdi, 'ResetTicks_%s'   % mca, None)
        allevt = getattr(xdi, 'AllEvent_%s'     % mca, None)
        if (ocr is None and icr is None and clock is None and
            dtfact is None and resets is None):
            nmca = i
            break
        if ocr is None:
            ocr = ctime
        ocr = ocr/ctime
        if icr is not None:
            icr = icr/ctime

        # Get ICR from one of several alternatives:
        # 1. InputCounts_mca was given
        # 2. DTFactor_mca was given
        # 3. ResetTicks_mca and AllEvets_mca were given
        # 4. Use "known values" for tau
        if icr is None:
            if dtfact is not None:
                icr = ocr * dtfact
                dtc_mode = 'dtfactor'
            elif (clock is not None and
                  resets is not None and
                  allevt is not None):
                dtfact = clock/(clock - (6*allevt + resets))
                icr = ocr * dtfact
                dtc_mode = 'resets'
        # finally estimate from measured values of tau:
        if icr is None:
            icr = 1.0*ocr
            if is_old_xsp3:
                icr = estimate_icr(ocr*1.00, XSPRESS3_TAU, niter=7)
        ocrs.append(ocr)
        icrs.append(icr)

    group.dtc_mode =  dtc_mode
    labels = []
    sums = {}
    for i, arrname in enumerate(xdi.array_labels):
        dat = getattr(xdi, arrname)
        if arrname.lower() == 'data':
            arrname = '_data'
        aname = sumname = rawname = arrname.lower()
        if ('_mca' in aname and
            'outputcounts' not in aname and
            'dtfactor' not in aname and
            'clock' not in aname):
            sumname, imca = sumname.split('_mca')
            if bad is not None and int(imca) in bad:
                sumname = sumname + '_bad'
            datraw = dat*1.0
            rawname = sumname + '_nodtc'
            imca = int(imca) - 1
            dat   = dat * icrs[imca]/ ocrs[imca]
            if any(np.isnan(dat)):
                nan_pts = np.where(np.isnan(dat))[0]
                dat[nan_pts] = datraw[nan_pts]

        setattr(group, aname, dat)
        if sumname not in labels:
            labels.append(sumname)
            sums[sumname] = dat
            if rawname != sumname:
                sums[rawname] = datraw
                if rawname not in labels:
                    labels.append(rawname)

        else:
            sums[sumname] = sums[sumname] + dat
            if rawname != sumname:
                sums[rawname] = sums[rawname] + datraw

    for name, dat in sums.items():
        if not hasattr(group, name):
            setattr(group, name, dat)

    for arrname in xdi.array_labels:
        sname = arrname.lower()
        if sname not in labels:
            labels.append(sname)

    data = []
    for name in labels:
        data.append(getattr(group, name))
    group.data = np.array(data)
    for imca in range(nmca):
        setattr(group, 'ocr_mca%i' % (imca+1), ocrs[imca])
        setattr(group, 'icr_mca%i' % (imca+1), icrs[imca])
    group.array_labels = labels
    return group


DTC_header = '''# XDI/1.1  Epics StepScan File/2.0
# Beamline.name:  13-ID-E, GSECARS
# Monochromator.name:  %(mono_cut)s, LN2 Cooled
# Monochromator.dspacing:  %(mono_dspace)s
# Facility.name: APS
# Facility.xray_source: 3.6 cm undulator
# Detectors.i0:  20cm ion chamber, He
# Detectors.ifluor:  Si SDD Vortex ME-4, 4 elements
# Detectors.ifluor_electronics:  Quantum Xspress3 3.1.10'''

def is_GSEXDI(filename):
    """test if file is GSE XDI data file
    reads only the first line of file
    """
    line1 = read_textfile(filename, size=40).split('\n')[0]
    return (line1.startswith('#XDI/1') and 'Epics StepScan File' in line1)

def gsexdi_deadtime_correct(fname, channelname, subdir='DT_Corrected',
                            bad=None):
    """convert GSE XDI fluorescence XAFS scans to dead time corrected files"""
    if not is_GSEXDI(fname):
        print("'%s' is not a GSE XDI scan file\n" % fname)
        return

    out = Group()
    out.orig_filename = fname
    try:
        xdi = read_gsexdi(fname, bad=bad)
    except:
        print('Could not read XDI file ', fname)
        return

    for attr in ('energy', 'i0', 'i1', 'i2', 'tscaler',
                 'counttime',  'scan_start_time', 'scan_end_time'):
        if hasattr(xdi, attr):
            setattr(out, attr, getattr(xdi, attr))

    # some scans may not record separate counttime, but TSCALER
    # is clock ticks for a 50MHz clock
    if not hasattr(out, 'counttime'):
        out.counttime = xdi.tscaler * 2.e-8

    if hasattr(xdi, 'energy_readback'):
        out.energy = xdi.energy_readback

    arrname = None
    channelname = channelname.lower().replace(' ', '_')

    for arr in xdi.array_labels:
        if arr.lower().startswith(channelname):
            arrname = arr
            break
    if arrname is None:
        print('Cannot find Channel %s in file %s '% (channelname, fname))
        return

    out.ifluor = getattr(xdi, arrname)
    out.ifluor_raw  = getattr(xdi, arrname)
    arrname_raw = arrname + '_nodtc'
    if arrname_raw  in xdi.array_labels:
        out.ifluor_raw  = getattr(xdi, arrname_raw)

    out.mufluor = out.ifluor / out.i0
    TINY  = 2.e-20
    if hasattr(out, 'i1') or hasattr(out, 'itrans'):
        i1 = getattr(out, 'i1', None)
        if i1 is None:
            i1 = getattr(out, 'itrans', None)
        if i1 is not None:
            i1[np.isnan(i1)] = TINY
            i1 = i1 / out.i0
            i1[np.where(i1<TINY)] = TINY
            out.mutrans = -np.log(i1)


    npts   = len(out.i0)
    col0_name = xdi.array_labels[0].lower()
    col0_units = None
    col0_data  = xdi.data[0, :]
    if col0_name == 'energy':
        col0_data  = out.energy
        col0_units = 'eV'

    buff =  ['# XDI/1.0  GSE/1.0']

    header = {}
    hgroups = ['beamline', 'facility', 'mono', 'undulator', 'detectors',
               'scaler', 'detectorstage', 'samplestage', 'scan', 'scanparameters']
    hskip = ['scanparameters.end', 'scanparameters.start']
    for agroup in hgroups:
        attrs = xdi._xdi.attrs.get(agroup, {})
        if agroup == 'mono': agroup = 'monochromator'
        header[agroup] = {}
        for sname in sorted(attrs.keys()):
            if "%s.%s" %( agroup, sname) not in hskip:
                header[agroup][sname] = attrs[sname]


    header['facility']['name'] = 'APS'
    header['facility']['xray_source'] = '3.6 cm undulator'
    header['beamline']['name'] = '13-ID-E, GSECARS'

    header['detectors']['i0'] = '20cm ion chamber, He'
    header['detectors']['ifluor'] = 'Si SDD Vortex ME-4, 4 elements'
    header['detectors']['ifluor_electronics'] = 'Quantum Xspress3 3.1.10'

    mono_cut = 'Si(111)'
    if xdi.mono_dspacing < 2:
        mono_cut = 'Si(311)'
    header['monochromator']['name'] = "%s, LN2 cooled"  % mono_cut

    out_arrays = {}
    out_arrays[col0_name]  = (col0_name, col0_units)
    out_arrays['mufluor'] = ('mufluor', None)
    if hasattr(out, 'i1'):
        out_arrays['mutrans'] = ('mutrans', None)

    out_arrays['ifluor']  = ('ifluor', '# deadtime-corrected')
    out_arrays['ifluor_raw'] = ('ifluor_raw', '# not deadtime-corrected')
    out_arrays['i0'] = ('i0', None)

    if hasattr(out, 'i1'):
        out_arrays['itrans'] = ('i1', None)
    if hasattr(out, 'i2'):
        out_arrays['irefer'] = ('i2', None)

    if hasattr(out, 'counttime'):
        out_arrays['counttime'] = ('counttime', 'sec')

    arrlabel = []
    for iarr, aname in enumerate(out_arrays):
        lab = "%12s " % aname
        if iarr == 0: lab = "%11s " % aname
        arrlabel.append(lab)
        extra = out_arrays[aname][1]
        if extra is None: extra = ''
        buff.append("# Column.%i: %s %s" % (iarr+1, aname, extra))

    arrlabel = '#%s' % (' '.join(arrlabel))
    ncol = len(out_arrays)

    for family, fval in header.items():
        for attr, val in fval.items():
            buff.append("# %s.%s: %s" % (family.title(), attr, val))


    buff.append("# ///")
    for comment in bytes2str(xdi._xdi.comments).split('\n'):
        c = comment.strip()
        if len(c) > 0:
            buff.append('# %s' % c)
    buff.extend(["# summed %s fluorescence data from %s" % (channelname, fname),
                 "# Dead-time correction applied",
                 "#"+ "-"*78,   arrlabel])

    efmt = "%11.4f"
    ffmt = "%13.7f"
    gfmt = "%13.7g"
    for i in range(npts):
        dline = ["", efmt % col0_data[i], ffmt % out.mufluor[i]]

        if hasattr(out, 'i1'):
            dline.append(ffmt % out.mutrans[i])

        dline.extend([gfmt % out.ifluor[i],
                      gfmt % out.ifluor_raw[i],
                      gfmt % out.i0[i]])

        if hasattr(out, 'i1'):
            dline.append(gfmt % out.i1[i])
        if hasattr(out, 'i2'):
            dline.append(gfmt % out.i2[i])
        if hasattr(out, 'counttime'):
            dline.append(gfmt % out.counttime[i])

        buff.append(" ".join(dline))
    ofile = fname[:]
    if ofile.startswith('..'):
        ofile = ofile[3:]
    ofile = ofile.replace('.', '_') + '.dat'
    ofile = os.path.join(subdir, ofile)
    if not os.path.exists(subdir):
        os.mkdir(subdir)
    try:
       fout = open(ofile, 'w', encoding=sys.getdefaultencoding())
       fout.write("\n".join(buff))
       fout.close()
       print("wrote %s, npts=%i, channel='%s'" % (ofile, npts, channelname))
    except:
       print("could not open / write to output file %s" % ofile)

    return out
