import os, gc
from pylab import *
import nmr, flib, scopelog, timelog, massfft, nmrfit
import datetime

def fix_remote_filename(filename, indices = None, \
            map = [('d:/exp', '/data/archive/phpc329/exp'), \
                   ('d:/exp', '/data/archive2/phpc329/exp'), \
                   ('//Phpc329/exp', '/data/archive2/phpc329/exp')]):
    """
    Convert Windows filenames of the NMR traces into ones
    pointing to corresponding directories in the archive on PHAULT.
    
    'filename' - filename or filename pattern to convert
    'indices' - an array of indices to substitute to 'filename' or 'None',
        if 'filename' is a complete filename itself.
    
    The function checks that 'function' or 'function % indices[0]' and
    'function % indices[-1]' exist in the potential redirected folders.
    
    'map' - a dictionary of redirections. Applied after forcing all '\' into '/'.
    """
    filename =  filename.replace('\\', '/')
    for _from, _to in map:
        new_filename = flib.ireplace(filename, _from, _to)
        if (indices is not None and os.path.exists(filename % indices[0]) \
            and os.path.exists(filename % indices[-1]))\
           or (indices is None and os.path.exists(filename)):
            return new_filename
    raise IOError('a local version of "%s" not found' % filename)

class NMRDataDir(object):
    """
    NMR Dataset - signals to be interpreted together
    """
    def __init__(self, pressure, demag, name, signals, first_signal, last_signal, backgrounds, bg_indices):
        self.pressure = pressure
        self.demag = demag
        self.name = name
        self.signals = fix_remote_filename(signals, [first_signal, last_signal])
        self.first_signal = first_signal
        self.last_signal = last_signal
        self.backgrounds = fix_remove_filename(backgrounds, bg_indices)
        self.bg_indices = bg_indices
        
        self.backgrounds = self.backgrounds.replace('\\', '/')
        for _from, _to in map:
            backgrounds = flib.ireplace(self.backgrounds, _from, _to)
            if os.path.exists(backgrounds % self.bg_indices[0]) and os.path.exists(backgrounds % self.bg_indices[-1]):
                self.backgrounds = backgrounds
                break
    
    def __str__(self):
        return "# %.3f bar, demag %d%s\n" % (self.pressure, self.demag, self.name) +\
               "# Signals: %s  (%d-%d)\n" % (self.signals, self.first_signal, self.last_signal) +\
               "# Backgrounds: %s %s\n" % (self.backgrounds, self.bg_indices)

    def logdir(self): return "%s_demag%d%s" % ("%.3fbar" if self.pressure > 0 else "Pramp", self.demag)
    def logfile(self, fft_settings):
        #D:\exp\RUN20\Superfluid Cell\5.5bar\260708\a\signal%d.dat
        if self.signals.count('bar') > 0:
            parts = self.signals.split('bar/')[1]
        elif self.signals.count('pressureRamp') > 0:
            parts = self.signals.split('pressureRamp/')[1]
        else:
            parts = parts.split('/')[-3]
        return parts.split('/')[0] + '.dat'

class NMRTask(object):
    """
    NMR Task - what to do with signals in a specified folder, a "virtual" parent class
    """
    def __init__(self, fft_settings, Naver):
        self.fft_settings = massfft.FFTSettings.preset(fft_settings)
        self.Naver = Naver
    
    def report_basic_settings(self):
        return "# Rolling average of %d signals per point.\n" % (self.Naver) +\
                self.fft_settings.report_settings() + "\n"

    def column_labels(self):
        """
        Return a list of labels of columns that are provided by a child class.
        Should be reimplemented.
        """
        raise NotImplementedError("'NMR Task' is an abstract class. Child classes should redefine 'column_labels'")
    
    def process_one(self, spectrum):
        """
        Process a given 'spectrum' and return a list of values (strings) to be written to the log file.
        The first 14 generic columns are written by this base class.
        This function should be reimplemented in child classes.
        """
        raise NotImplementedError("'NMR Task' is an abstract class. Child classes should redefine 'process_one'")
    
    def process(data_dir, toplogdir = '/data/exp/run20/sf/intgr', rewrite=False, update=True):
        """
        'data_dir' describes the data to process
        'toplogdir' is a top directory for the hyerarchical tree of log files.
        Depending on 'rewrite', the log file is erased or appended.
        With 'rewrite' false and 'update' true, the existing log file is
        examined and only signals not present there are processed.
        """
        fft_settings = massfft.FFTSettings.preset(fft_settings)
        
        #Choose output path and filename and create a header for the output file
        
        logdir = toplogdir
        if not os.path.exists(logdir):
            os.mkdir(logdir)
        
        logdir = os.path.join(logdir, data_dir.logdir())
        if not os.path.exists(logdir):
            os.mkdir(logdir)
        
        logdir = os.path.join(logdir, self.logdir())
        if not os.path.exists(logdir):
            os.mkdir(logdir)
            
        output = os.path.join(logdir, data_dir.logfile(fft_settings))
        
        out = open(output, 'w' if rewrite else 'a')
        try:
            out.write(str(data_dir))
            out.write(self.report_basic_settings())
            out.write('#\n')
            out.write('# [0] First Signal.\n')
            out.write('# [1] Last Signal.\n')
            out.write('# [2] Start time of first signal averaged (seconds since 1st Jan 1970).\n')
            out.write('# [3] End time of last signal averaged (seconds since 1st Jan 1970).\n')
            out.write('# [4] Number of traces averaged together.\n')
            out.write('# [5] Average PLM temperature, XMIT4, Gain 10 (mK).\n')
            out.write('# [6] Half PLM temperature span XMIT4, Gain 10 (mK).\n')
            out.write('# [7] Average PLM temperature, XMIT4, Gain 20 (mK).\n')
            out.write('# [8] Half PLM temperature span XMIT4, Gain 20 (mK).\n')
            out.write('# [9] Average PLM temperature, XMIT8, Gain 10 (mK).\n')
            out.write('# [10] Half PLM temperature span XMIT8, Gain 10 (mK).\n')
            out.write('# [11] Mean pressure (mbar)\n')
            out.write('# [12] Min pressure (mbar)\n')
            out.write('# [13] Max pressure (mbar)\n')
            n = 14
            for l in self.column_labels():
                if l[0] == '"':
                    out.write('#\n')
                    out.write('# %s\n' % l[1:])
                else:
                    out.write('# [%d] %s\n' % (n, l))
                    n += 1
            out.write('#\n')
            out.close()
            
            # get start and final time for each line in the log
            num_files = data_dir.last_signal - data_dir.first_signal + 1
            start = zeros(num_files)
            end = zeros(num_files)
            
            for i in range(data_dir.first_signal, data_dir.last_signal+1):
                info=scopelog.getinfos(data_dir.signals, [i])
                start[i-data_dir.first_signal] = info.starttime
                end[i-data_dir.first_signal] = info.finaltime
                
            #load PLM and p31k logs for a suitable time span to cover all signals. PLM data is for XMIT4, Gain 10
            #mct = timelog.MCTLog.loadrun(20, start = min(start) - 1800, end = max(end) + 3600).at(xmit=4, gain=10)
            p31k = timelog.ParoLog.loadHighP(20, start = min(start) - 1800, end = max(end) + 3600)
            plm = timelog.PLMLog.loadrun(20, start = min(start) - 1800, end = max(end) + 3600)
            plm4_10 = plm.at(xmit=4, gain=10)
            plm4_20 = plm.at(xmit=4, gain=20)
            plm8_10 = plm.at(xmit=8, gain=10)
            
            #load backgrounds
            b = massfft.getfft(self.fft_settings, data_dir.backgrounds, data_dir.bg_indices)
            gc.collect()
            #b = scopelog.loadbinaries(Bgs, NumBgs, 10e6)
            
            for first in range(data_dir.first_signal, (data_dir.last_signal + 1)-(self.Naver - 1)):
                last = first + self.Naver - 1
                try:
                    #load signals
                    s = massfft.getfft(self.fft_settings, data_dir.signals, frange(first, last))
                    
                    #Subtract backgrounds and compute FFT
                    w = 1e6 * (s - b)
                    
                    #Calculate PLM reading at the start and end of the signal(s) and then compute the mean reading and error bar.
                    T_start = plm4_10.interpolate(s.info.starttime)
                    T_end = plm4_10.interpolate(s.info.finaltime)
                    T4_10 = 0.5*(T_start + T_end)
                    dT4_10 = 0.5*abs(T_start - T_end)
                    
                    T_start = plm4_20.interpolate(s.info.starttime)
                    T_end = plm4_20.interpolate(s.info.finaltime)
                    T4_20 = 0.5*(T_start + T_end)
                    dT4_20 = 0.5*abs(T_start - T_end)
                    
                    T_start = plm8_10.interpolate(s.info.starttime)
                    T_end = plm8_10.interpolate(s.info.finaltime)
                    T8_10 = 0.5*(T_start + T_end)
                    dT8_10 = 0.5*abs(T_start - T_end)
                    
                    #Calculate mean pressure between start and end of signals
                    P = p31k.period(s.info.starttime, s.info.finaltime).P()
                    Pmean = mean(P)
                    #Pstd = std(P)
                    Pmin = min(P)
                    Pmax = max(P)
                    
                    #Write data to the output file
                    out = open(output, 'a')
                    out.write('\t'.join([
                            '%d' % first, # [0] First Signal.
                            '%d' % last, # [1] Last Signal.
                            '%.2f' % s.info.starttime, # [2] Start time of first signal averaged (seconds since 1st Jan 1970).
                            '%.2f' % s.info.finaltime, # [3] End time of last signal averaged (seconds since 1st Jan 1970).
                            '%d' % s.info.Nav, # [4] Number of traces averaged together.
                            '%.4f' % T4_10, # [5] Average PLM temperature, XMIT4, Gain 10 (mK).
                            '%.4f' % dT4_10, # [6] Half PLM temperature span XMIT4, Gain 10 (mK).
                            '%.4f' % T4_20, # [7] Average PLM temperature, XMIT4, Gain 20 (mK).
                            '%.4f' % dT4_20, # [8] Half PLM temperature span XMIT4, Gain 20 (mK).
                            '%.4f' % T8_10, # [9] Average PLM temperature, XMIT8, Gain 10 (mK).
                            '%.4f' % dT8_10, # [10] Half PLM temperature span XMIT8, Gain 10 (mK).
                            '%.1f' % Pmean, # [11] Mean pressure (mbar)
                            '%.1f' % Pmin, # [12] Min pressure (mbar)
                            '%.1f' % Pmax, # [13] Max pressure (mbar)
                        ] + list(self.process_one())) + '\n')
                    out.close()
                except Error, e:
                    out = open(output, 'a')
                    out.write('# %d\t%d ERROR: %s\n' % (first, last, e))
                    out.close()
                gc.collect()
        finally:
            out.close()

class IntegrationTask(NMRTask):
    """
    Integrate and find 
    """
    def __init__(self, fft_settings, Naver, min_f, max_f, f0):
        super(IntegrationTask,self).__init__(fft_settings, Naver)
        self.min_f = min_f
        self.max_f = max_f
        self.f0 = f0

    def column_labels(self):
        return ['Peak Frequency [Hz]',
                'Peak Frequency Shift [Hz] from %.2f Hz' % f0,
                'Peak Amplitude [uV*s]',
                'Amplitude Integral [uV]',
                'Amplitude Squared Integral [uV^2*s]',
               ]

    def process_one(self, spectrum):
        #Calculate average noise level on the low-frequency side of the slab peak, remove single point peaks from spectrum
        w = spectrum.frange(self.min_f, self.max_f)
        f, A = w.findpeak()
        
        return ['%.2f' % f,
                '%.2f' % f - f0,
                '%.8g' % A, 
                '%.6e' % flib.integrate(w.f, w.abs())[-1],
                '%.6e' % flib.integrate(w.f, w.abs()**2)[-1]
               ]

class IntgNFitTask(NMRTask):
    def __init__(self, fft_settings, Naver, min_f, max_f, f0, minA):
        super(IntegrationTask,self).__init__(fft_settings, Naver)
        self.min_f = min_f
        self.max_f = max_f
        self.f0 = f0

    def column_labels(self):
        return ['Peak Frequency [Hz]',
                'Peak Frequency Shift [Hz] from %.2f Hz' % f0,
                'Peak Amplitude [uV*s]',
                'Amplitude Integral [uV]',
                'Amplitude Squared Ingegral [uV^2*s]',
                'Fit Frequency [Hz]',
                'Fit Frequency Shift [Hz] from %.2f Hz' % f0,
                'Fit Amplitude [uV]',
                'Fit T2 [ms]',
                'Fit Phase [deg]'
                'Fit Converged? 1 or 0'
               ]

    def process_one(self, spectrum):
        #Calculate average noise level on the low-frequency side of the slab peak, remove single point peaks from spectrum
        w = spectrum.frange(self.min_f, self.max_f)
        f, A = w.findpeak()
        
        try:
            fit = nmrfit.Lorentzian.fit(w)
            goodFit = 1
        except nmrfit.FitError, ex:
            fit = ex.fit
            goodFit = 0
        
        return ['%.2f' % f,
                '%.2f' % f - f0,
                '%.8g' % A, 
                '%.6e' % flib.integrate(w.f, w.abs())[-1],
                '%.6e' % flib.integrate(w.f, w.abs()**2)[-1],
                '%.2f' % fit.f,
                '%.2f' % fit.f - f0,
                '%.6g' % fit.A,
                '%.6g' % fit.T2 * 1e3,
                '%.1f' % fit.phase_deg(),
                '%d' % goodFit
               ]