"""
Manipulating the ULT fridge Platinum Thermometer Lorentzian fit logs.

Changelog (partial)
7 Oct 2009, Lev
initial draft
"""
# ...
#
# 12 Jul 2010, Lev
# add 'set_T' method

import os
import os.path
import time
import datetime
import numpy
import matplotlib.dates
import flib
import timelog
import ascii2numpy

class PLMFitLog(timelog.TimeLog):
    """
    Represents a ULT PLM4 thermometry Lorentzian fit log over a certain time.
    Two logs can be concatenated with a '&' sign

    log format:
    # [0] Timestamp in PLM log (sec since Unix epoch)
    # [1] Scope acquisition start time (sec since Unix epoch)
    # [2] Scope acquisition final time (sec since Unix epoch)
    # [3] truncation time [usec]
    # [4] f0 [Hz]
    # [5] A [V]
    # [6] T2* [usec]
    # [7] phase [degrees]
    # [8] PLM4 XMIT
    # [9] PLM4 gain
    # [10] PLM4 int. delay [a.u.]
    # [11] PLM4 int. time [a.u.]
    # [12] PLM4 M0 [a.u.]
    # [13] PLM4 T [mK]
    """

    @staticmethod
    def fieldcount(): "return number of columns in the self.data"; return 20

    # timing
    def t_start(self): "scope acquisition start time"; return self.data[:,1]
    def t_end(self): "scope acquisition end time"; return self.data[:,2]
    def sampling_rate(self): "sampling rate [Hz] (possibly downsampled to match backgrounds)"; return self.data[:,3]
    
    # fit settings
    def fit_fmin(self): "bottom of fit frequency range [Hz]"; return self.data[:,4]
    def fit_fmax(self): "top of fit frequency range [Hz]"; return self.data[:,5]
    def fit_truncate_start(self): "start of FID portion used for FFT [usec]";  return self.data[:,6]
    def fit_truncate_end(self): "end of FID portion used for FFT [usec]";  return self.data[:,7]
    def fit_Npad(self): "FFT length (truncation/zero padding)";  return self.data[:,8]
    
    # fit results
    def f0(self): "fit resonance frequency [Hz]"; return self.data[:,9]
    def A(self): "fit time domain amplitude [A]"; return self.data[:,10]
    def T2(self): "fit T2* [usec]"; return self.data[:,11]
    def phase_deg(self): "fit phase [deg]"; return self.data[:,12]
    def fit_res(self): "fit residualts [V^2*sec]"; return self.data[:,13]
    
    # PLM log entries
    def xmit(self): "PLM4 transmitter pulse length"; return self.data[:,14].astype('i')
    def gain(self): "PLM4 gain"; return self.data[:,15].astype('i')
    def idelay(self): "PLM4 integration delay"; return self.data[:,16].astype('i')
    def itime(self): "PLM4 integration time"; return self.data[:,17].astype('i')
    def M0(self): "PLM4 M0"; return self.data[:,18]
    def T(self): "PLM4 temperature"; return self.data[:,19]
    def set_T(self, T): self.data[:,19] = T
    
    def A_extrapolate(self, PLM_STU=1./284780):
        """FID size extrapolated to the middle of the transmitter pulse
           the 'PLM_STU' parameter is the PLM time unit (XMIT pulse cycle)"""
        return self.A() * numpy.exp((self.fit_truncate_start()*1e-6 - 0.5*PLM_STU*self.xmit()) / (self.T2()*1e-6))

    def M0_correction(self, PLM_STU=1./284780):
        """assuming a Lorentzian free induction decay,
           calculate the ratio of Pt magnetisation and M0 measured in an integration routine
           the 'PLM_STU' parameter is the PLM time unit (XMIT pulse cycle)"""
        t_start = PLM_STU * (self.idelay() - 0.5*self.xmit()) * 1e6
        t_end = PLM_STU * (self.idelay() + self.itime() - 0.5*self.xmit()) * 1e6        
        return 2./numpy.pi * self.T2() * (numpy.exp(-t_start/self.T2()) - numpy.exp(-t_end/self.T2()))

    def at(self, xmit=None, gain=None, idelay=None, itime=None, overload=None):
        """
        Return a subset of the log obtained for at given settings. Any of
        the paramameters, xmit, gain, idelay, itime, overload is used for
        filtering points if it is specified (is not 'None').
        """
        
        subset = self.M0() == self.M0()
        
        if xmit is not None:
            subset &= (self.xmit() == xmit)
        if gain is not None:
            subset &= (self.gain() == gain)
        if idelay is not None:
            subset &= (self.idelay() == idelay)
        if itime is not None:
            subset &= (self.itime() == itime)
        if overload is not None:
            subset &= (self.overload() == overload)
        
        return self.subset(subset)

    def atxmit(self, xmit): return self.at(xmit=xmit)

    def xmits(self):
        """
        Return a list of transmitter pulse lengths used present in the log
        """
        return sorted(set(self.xmit()))

    def __repr__(self): return "PLMFitLog %s, xmit=%s" % (self.datespan(), self.xmits())

    @staticmethod
    def load(filename):
        """
        Load PLM FID fit log from a given file
        The function determines file format from the name:
            1. '????????_PLM_fit.dat'
            no other formats supported
        """
        
        basename = os.path.basename(filename)
        if basename[8:] == '_PLM_fit.dat': 
            # [0] Timestamp in PLM log (sec since Unix epoch)
            # [1] Scope acquisition start time (sec since Unix epoch)
            # [2] Scope acquisition final time (sec since Unix epoch)
            # [3] Sampling rate [Hz] (possibly downsampled to match backgrounds)
            #
            # Lorentzian fit settings:
            # [4,5] fit frequency range [Hz]
            # [6,7] truncation time start/end [usec]
            # [8] fit FFT length [truncation/zero padding]
            # Lorentzian fit results:
            # [9] f0 [Hz]
            # [10] A [V]
            # [11] T2* [usec]
            # [12] phase [degrees]
            # [13] residualts [V^2*sec]
            #
            # PLM settings and measurement:
            # [14] PLM4 XMIT
            # [15] PLM4 gain
            # [16] PLM4 int. delay [a.u.]
            # [17] PLM4 int. time [a.u.]
            # [18] PLM4 M0 [a.u.]
            # [19] PLM4 T [mK]        
            data = ascii2numpy.loadascii(filename, cols=PLMFitLog.fieldcount())
            return PLMFitLog(data)

##            # old format, currently unsupported
##            # [0] Timestamp in PLM log (sec since Unix epoch)
##            # [1] Scope acquisition start time (sec since Unix epoch)
##            # [2] Scope acquisition final time (sec since Unix epoch)
##            # [3] truncation time [usec]
##            # [4] f0 [Hz]
##            # [5] A [V]
##            # [6] T2* [usec]
##            # [7] phase [degrees]
##            # [8] PLM4 XMIT
##            # [9] PLM4 gain
##            # [10] PLM4 int. delay [a.u.]
##            # [11] PLM4 int. time [a.u.]
##            # [12] PLM4 M0 [a.u.]
##            # [13] PLM4 T [mK]
        else:
            # other file formats are not recognised
            return None
    
    @staticmethod
    def loadrun(run, start = None, end = None):
        """
        Return plm log for a given run, withing [start, end] period if boundaries are specified.
        """
        logdir = '/data/exp/run20/plm4/simple_fits' if run == 20 else None
        return PLMFitLog.loaddir(logdir, start=start, end=end).period(start, end)
