'''
# Linear Stark Analysis

Perform a linear stark analysis using the nudie library.

## Functional Requirements

- Callable from command-line
- Possible to run in batch mode in a distributed manor
- Ideally, scalable (not sure if needed for linear stark)

## Design notes
## To-Do
- use Sphinx conventions for documenting functions
'''
from __future__ import division, unicode_literals
import nudie
import itertools as it
from scipy.signal import argrelmax
from pathlib import Path
from collections import deque
import re
from numpy.fft import fft, ifft, fftshift, fftfreq
from scipy.signal import get_window
from scipy.interpolate import interp1d
from scipy.io import loadmat
import arrow
import numpy as np
import h5py
import numpy.ma as ma

import matplotlib.pyplot as pyplot

import argparse

# turn on printing of errors
nudie.show_errors(nudie.logging.INFO)

def make_parser():
    parser = argparse.ArgumentParser(
        description='Analyze a single linear Stark dataset taken in ' +\
                'the pump-probe geometery. Output a hdf5 file with ' +\
                'the analysis.')
    parser.add_argument('when', type=str, action='store', 
        help='the date folder that contains the batch')
    parser.add_argument('jobname', type=str, action='store', 
        help='the name of the job set in LabVIEW')
    parser.add_argument('batch', type=int, action='store',
        help='batch number') 
    parser.add_argument('--wavelengths', metavar='wl', type=str, action='store',
        help='calibration data file for the spectrometer, including .mat', default=None)
    parser.add_argument('--plot', action='store_const', 
            const=True, default=False, help='pop up diagnostic plots')
    return parser

def trim_to_shortest(*iterable):
    # determine shortest element
    lengths = [len(x) for x in iterable]
    m = min(lengths)
    return tuple(x[:m] for x in iterable)

def load_wavelengths(path):
    '''load a pre-calibrated wavengths file generated with the
    ```manual_calibration``` script in MATLAB.
    '''
    calib_file = Path(path)
    with calib_file.open('rb') as f:
        calib = loadmat(f)
        wl = np.squeeze(calib['saved_wavelengths'])
    return wl

def run(stark_name, stark_batch, when='today', wavelengths=None, plot=False, analysis_path='analyzed', datapath = None):
    '''run the batch job to analyze the linear-stark data
    
    Performs a shot-to-shot linear stark analysis. 
    
    ```wavelengths``` is a path to a wavelengths file
    ```plot``` is a Boolean parameter controlling whether to show matplotlib
    plots
    '''

    analysis_folder = Path(analysis_path)

    # create folder if it doesn't exist
    if not analysis_folder.exists():
        analysis_folder.mkdir()

    # change datapath
    if datapath is None:
        datapath = nudie.data_folder

    # load up pp data to use for linear stark
    stark_info = next(nudie.load_job(job_name=stark_name, batch_set=[stark_batch], when=when, data_path=datapath))

    # set current batch directory
    current_path = Path(stark_info['batch_path'])

    if wavelengths is not None:
        wl = load_wavelengths(current_path.parent / wavelengths)

    # tags for phase cycles
    phase_cycles = ['none1', 'zero', 'none2', 'pipi']
    nrepeat = 1 # how many times each waveform is repeated in the camera file. Assumed to be one
    nwaveforms = 1 # number of dazzler waveforms in pump probe is 1, assumed to be one

    # used to store data generated by iterations
    saved_data = deque()

    # loop over available data
    for t2,table,loop in it.product(stark_info['t2_range'], stark_info['table_range'], stark_info['loop_range']):
        # first file requires special synchronization
        # this is the rule that determines that it is the first file
        first = 'first' if all([t2 == 0, table == 0, loop == 0]) else False
        
        # Load everything for the given t2, table, and loop value
        analogs = nudie.load_analogtxt(stark_info['job_name'], current_path, t2, table, loop)
        cdata = nudie.load_camera_file(stark_info['job_name'], current_path, t2, table, loop, force_uint16=True)
        
        # Synchronize it, trimming some frames in the beginning and end
        data, (a1, a2) = nudie.trim_all(*nudie.synchronize_daq_to_camera(\
                        cdata, analog_channels=analogs, which_file=first), trim_to=slice(2, None))
        data = data.astype(float) # convert to float64 before manipulating it!
        
        start_idxs, period = nudie.detect_table_start(a1)

        # determine where the shutter is
        shutter_open, shutter_closed, duty_cycle = nudie.determine_shutter_shots(data)
        
        # tag the phases
        shutter_info = {'last open idx': shutter_open, 'first closed idx': shutter_closed}                
        tags = nudie.tag_phases(start_idxs, period, tags=phase_cycles, nframes=data.shape[1], shutter_info=shutter_info)
        
        # verify synchronization of applied field to Dazzler
        for i,k in enumerate(phase_cycles):
            # check that every second frame in analog channel 2 is either high field, or low field, but not both
            
            # Offset 0, every second frame has high field
            ao = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][0::2] > 0.2
            # Offset 0, every second frame has low field
            bo = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][0::2] <= 0.2 
            
            # Offset 1, every second frame has high field
            co = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][1::2] > 0.2
            # Offset 1, every second frame has low field
            do = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter open']][1::2] <= 0.2
            
            # Offset 0, every second frame has high field
            ac = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][0::2] > 0.2
            # Offset 0, every second frame has low field
            bc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][0::2] <= 0.2 
            
            # Offset 1, every second frame has high field
            cc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][1::2] > 0.2
            # Offset 1, every second frame has low field
            dc = a2[tags[nrepeat-1][nwaveforms-1][k]['shutter closed']][1::2] <= 0.2
            
            assert np.logical_xor(np.all(ao), np.all(bo)), 'detected stark synchronization error, offset 0, phase %s' % k
            assert np.logical_xor(np.all(co), np.all(do)), 'detected stark synchronization error, offset 1, phase %s' % k
            
            assert np.logical_xor(np.all(ac), np.all(bc)), 'detected stark synchronization error, offset 0, phase %s' % k
            assert np.logical_xor(np.all(cc), np.all(dc)), 'detected stark synchronization error, offset 1, phase %s' % k

        # determine stark-on indexes
        stark_idx, nostark_idx = np.where(a2 > 0.2)[0], np.where(a2 <= 0.2)[0]
            
        tmp = []
        # remove the probe scatter phase by phase
        for i,k in enumerate(phase_cycles):
            so = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter open'], stark_idx, assume_unique=True)
            sc = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter closed'], stark_idx, assume_unique=True)                   

            no = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter open'], nostark_idx, assume_unique=True)
            nc = np.intersect1d(tags[nrepeat-1][nwaveforms-1][k]['shutter closed'], nostark_idx, assume_unique=True)                   
                        
            # remove probe scatter for stark shots
            data[:, so] -= data[:, sc].mean(axis=-1)[:, np.newaxis]
            # and for no-stark shots
            data[:, no] -= data[:, nc].mean(axis=-1)[:, np.newaxis]        
            
            data = ma.masked_less(data, 1e-7*data.max())
            stark_spec = np.log10(np.mean(data[:, so], axis=-1)/np.mean(data[:, no], axis=-1))
            stark_spec = ma.filled(stark_spec, 0)
            tmp.append(stark_spec)
        
        # average phase cycles together to get slight improvement in S/N
        px = max([len(i) for i in tmp]) # num pixels
        m = np.zeros((px, len(phase_cycles)))
        for i,k in enumerate(tmp):
            m[:, i] = k
        avg_spectrum = m.mean(axis=1)
                       
        mean_voltage = a2[stark_idx].mean()

        saved_data.append((avg_spectrum, "loop {}".format(loop)))

        # interpret `when` parameter
        if when == 'today':
            timestamp = arrow.now()
        elif when is not None:
            # try to interpret when as a string
            timestamp = arrow.get(when, 'YY-MM-DD')
        else:
            raise NotImplemented('don\'t pass a direct path to the ' +\
                    'filename please')

        # Write out information generated
        filename = '{timestamp:s}-{name:s}-{voltage:1.1f}kV.hdf5'.format(
                timestamp=timestamp.format('YYYY-MM-DD'),
                name=stark_info['batch_name'], voltage=mean_voltage)

        with h5py.File(str(analysis_folder / filename), mode='w') as f:
            for data, j in saved_data:
                f[j] = data
            
            f.attrs['applied voltage'] = mean_voltage
            f.attrs['file path'] = stark_info['batch_path']
            f.attrs['job name'] = stark_info['job_name']
            f.attrs['batch_no'] = stark_batch # FIXME: workaround for nudie bug

            if wavelengths is not None:
                f['wavelength axis'] = wl

            if all([plot, loop == stark_info['loop_range'].stop-1]):
                for data, loop in saved_data:
                    pyplot.plot(wl, data, label=loop)

                pyplot.xlabel('wavelength / nm')
                pyplot.ylabel('OD')
                pyplot.show()




def main(config, verbosity=nudie.logging.INFO):
    nudie.show_errors(verbosity)

    try:
        try:
            val = nudie.parse_config(config, which='linear')['linear']
        except ValueError as e:
            nudie.log.error('could not validate file. Please check ' +\
                'configuration options.')
            return

        if val['stark'] != True:
            s = 'the stark flag is not set in the configuration. ' +\
                'Are you sure you are using the right configuration file?'
            nudie.log.error(s)
            return

        run(stark_name=val['jobname'],
            stark_batch=val['batch'],
            when=val['when'],
            wavelengths=val['wavelengths'],
            plot=val['plot'],
            analysis_path=val['analysis path'],
            datapath=val['data path'])
    except Exception as e:
        nudie.log.exception(e)


if __name__ == '__main__':
    parser = make_parser()
    args = vars(parser.parse_args())
    
    if args['plot']:
        import matplotlib.pyplot as pyplot

    try:
        run(stark_name=args['jobname'], stark_batch=args['batch'],
                when=args['when'], plot=args['plot'],
                wavelengths=args['wavelengths'])
    except Exception as e:
        # shut off exceptions
        raise e
