from __future__ import division # in case this is called from python2 
from __future__ import print_function

import numpy as np
import matplotlib.pyplot as plt
import sys, os
import notch 
from scipy import signal
from scipy.interpolate import splprep, splev, spline # for smoothing spline FITPACK 
from scipy.interpolate import UnivariateSpline
from decay import *
import scipy.stats as stats
from pwctimeWhite import * 
# for signal slot communication with GUI 
#from PyQt4.QtCore import QObject, SIGNAL
from PySide.QtCore import QObject, SIGNAL
from pandas.io.parsers import read_csv #<<- faster than Numpy but not working right now
from timeit import timeit 


class MRProc(QObject):
    """ Class to process read and invert ORS borehole data
    """ 
    def __init__(self):
        QObject.__init__(self)
        self.burst = False           # For Schlumberger or VC burst data 

    #def __init__(self, filename):
    #    # what do we need to know that isn't in new files?       
    #    fileName, fileExtension = os.path.splitext(filename)
    #    if fileExtension == ".ors":
    #        self.loadORSFile(filename)
        #self.DistRes = { } 
        #self.DistRes = { 'env':np.array(0), 'mod':np.array(0), 'Time':np.array(0) } #[a0,b0,rt20] }
        
    def loadORSFile_OLD(self, filename):
        """ Loads ORS files with minimal header info
        """
        f = open(filename)
        header = f.readline().split() #np.loadtxt(filename, comment="#", )
        ih = 1
        while (header[0] == "#"):
            ih += 1
            header = f.readline().split() #np.loadtxt(filename, comment="#", )
        self.NR = eval(header[0])    # number of records in an echo
        self.NE = eval(header[1])    # number of echoes
        self.DT = eval(header[2])    # in s
        self.TAUE = eval(header[3])  # in s
        self.RL = self.DT * self.NR  # in s

        # pandas is about ~16 times faster than numpy for import (7 sec vs. .5 sec on tested example)  
        # Parse in datafile, TODO update for newer format and remove hardcoded 3 in pandas
        #DATA = np.loadtxt(filename, comments="#", skiprows=ih)
        DATA = read_csv(filename, comment="#", sep='\t', skiprows=3, engine='c', header=None, names=('time','inphase','quadrature')).as_matrix() 

        self.NS = int(np.shape(DATA)[0] / (self.NE * self.NR))

        # reshape DATA
        #self.DATA = np.reshape( DATA[:,1] + 1j*DATA[:,2], ( self.NS, self.NE, self.NR) )
        self.DATA = np.reshape( self.A2D(DATA[:,1]) + 1j*self.A2D(DATA[:,2]), ( self.NS, self.NE, self.NR) )
        self.WTIME = np.reshape( DATA[:,0], ( self.NS, self.NE, self.NR) )
        self.TIMES = np.arange(self.DT, self.RL + 1e-8, self.DT)           # 1 ms record, 1 us dwell 

        self.emit(SIGNAL("plotRAW()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def evaluate(self, value):
        if len(value) == 1:
            if value[0][-1].isdigit(): 
                return eval(value[0])
            else:
                if value[0][-1] == "u":   # micro
                    return eval(value[0][0:-1]) * 1e-6
                elif value[0][-1] == "M": # mega
                    return eval(value[0][0:-1]) * 1e6
                elif value[0][-1] == "s": # second
                    return eval(value[0][0:-1]) 
                else:
                    print("DOOM!")
                    exit(1)
        else:
            # TODO
            # logic here is a little lacking, but current tool does not do anything 
            # sneaky with array values, could break in the future though 
            vals = []
            for val in value:
                vals.append( eval(val) )
            return vals

    def loadORSFile(self, fname):
        """ Loads ORS files with extensive header info
        """
        parameters = False
        data = False
        pars = {}
        with open(fname) as f:
            iline = 0
            for line in f:
                if len(line.split()) == 0:
                    pass # empty do nothing
                else:
                    if line.split()[0] == "[ps]":
                        parameters = False
                        # these entries specify the chip programming 
                    if parameters:
                        # proc header info
                        split,comment = line.split("#")
                        var, val = split.split("=")
                        val = val.split(";")[0].split()
                        pars[var.strip()] = self.evaluate(val)
                    if line.split()[0] == "[parameters]":
                        parameters = True
                    if data:
                        if line[0] == "#":
                            pass
                        else:
                            header = line.split()
                            iline += 1
                            break
                    if line.split()[0] == "[data]":
                        data = True
                iline += 1

        self.NR = eval(header[0])    # number of records in an echo
        self.NE = eval(header[1])    # number of echoes
        self.DT = eval(header[2])    # in s
        self.TAUE = eval(header[3])  # in s
        self.RL = self.DT * self.NR  # in s        
        
        # Pandas is much faster than numpy
        DATA = read_csv(fname, comment="#", sep='\t', skiprows=iline, engine='c', header=None, names=('time','inphase','quadrature')).as_matrix() 
        self.NS = int(np.shape(DATA)[0] / (self.NE * self.NR))

        # reshape DATA
        #self.DATA = np.reshape( DATA[:,1] + 1j*DATA[:,2], ( self.NS, self.NE, self.NR) )
        self.DATA = np.reshape( self.A2D(DATA[:,1], pars) + 1j*self.A2D(DATA[:,2], pars), ( self.NS, self.NE, self.NR) )
        self.WTIME = np.reshape( DATA[:,0], ( self.NS, self.NE, self.NR) )
        self.TIMES = np.arange(self.DT, self.RL + 1e-8, self.DT)           # 1 ms record, 1 us dwell 

        self.emit(SIGNAL("plotRAW()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("doneStatus()"))  
    
    # analogue to digital conversion, 'don't quote me on this -- kevin'
    def A2D(self, i, pars):
            """ Actual voltage is not known. This routine will need to be expaned to account for 
                calibration datasets. Or perhaps the inversion and fitting routines instead.
                TODO, if prestacked data, we need to account for this. 
            """ 
            #return  (i/(.5 * (2.**16)))
            #cal = 7.85 # TODO query for this
            cal = 5.51 # 1/2 full
            return i / ( pars['n'] * pars['cal'] )

    def loadJavelinLog(self, filename):
        """Loads processed Javelin data"""
        #print("Javelin time")
        import scipy.io as sio
        Data =  sio.loadmat(filename)
        self.T2T = Data['time'][:,0] 
        self.TAUE =  self.T2T[3] - self.T2T[2]
        # Better sigma estimate needed
        self.sigma = 2.5 
        self.T2N = np.random.normal( 0, self.sigma, len(self.T2T) )    
        self.NS = 1
        self.T2D = self.T2N + 1j*Data['se_vector_wc'][:,56]       # depth 0 
        # What about each depth? How is modelling handled.

    def loadVCMFile(self, filename):
        #print("LOADING VCM file")
        Data = np.loadtxt(filename, comments = "#")
        self.T2T = Data[:,0]    
        self.sigma = eval(open( filename, 'r' ).readline().split()[1])       
        self.T2N = np.random.normal( 0, self.sigma, len(self.T2T) )    
        self.T2D = self.T2N + 1j*Data[:,1]
        self.NS = 1
        self.TAUE =  self.T2T[3] - self.T2T[2]
        
        #self.emit(SIGNAL("plotLOG10ENV()"))   
        #self.emit(SIGNAL("enableDSP()"))   
        #self.emit(SIGNAL("enableINV()"))   
        #self.emit(SIGNAL("doneStatus()"))  
 
    def loadTextFile(self, filename, NS, NE, NR):
        """ Loads older format files without header info """
        print ("OLD TEXT FORMAT DETECTED, not SUPPORTED RIGHT NOW")
        exit()
        pass

    def batchThread(self, pars, tag, bload=True, bwindow=True, benvelope=True, bphase=True, bgate=True):
        
        #print("batdacheing", load, window, envelope, phase, gate)
        fileName, fileExtension = os.path.splitext( str(tag) )
        if bload:
            #print ("loading", tag)
            if fileExtension == ".ors":
                self.loadORSFile(str(tag)) 
            elif fileExtension == ".vcm":
                self.loadVCMFile(str(tag)) 
        
        if bwindow:
            #print("windowing", pars['ftype'], pars['winWidth'])
            self.WindowAndStack( pars['ftype'] , pars['winTrimLeft'], pars['winTrimRight'] )

        if benvelope:
            #print("enveloping", pars['offset'])
            self.T2EnvelopeDetect( pars['offset'] ) 
        
        if bphase:
            self.CorrectedAmplitude( 0,0 ) 
        
        if bgate:
            #print("gating", pars['gpd'])
            self.gateIntegrate( pars['gpd'], pars['stackEfficiency'] ) 

        self.emit(SIGNAL("doneStatus()"))  
        

    def WindowAndStack(self, ftype="Blackman-Harris", triml=0, trimr=0):
        # TODO, consider lowpass filter and interpolate to higher DT in order to minimize window amplitude effects
        #       what would cost of such an operation be? 
        # TODO, should record be overwritten? I don't think so. I think we should make a deep copy here and then you can always fall back.
        
        # Trim ends to centre echo, this many are REMOVED        
        #triml = 18  
        #trimr = 15   
        self.TIMES = self.TIMES[triml:-(trimr)]
        self.NR -= triml + trimr
        if ftype == "Blackman-Harris":
            window = signal.blackmanharris(self.NR)
        elif ftype == "Blackman":
            window = signal.blackman(self.NR)
        elif ftype == "Hann":
            window = signal.hann(self.NR)
        elif ftype == "sinc":
            window = np.ones( (self.NR) )

        self.DATAP = self.DATA[:,:,triml:-trimr].copy()
        for istk in range(0,self.NS):
            for ie in range(0,self.NE,1):
                self.DATAP[istk, ie,:] *= window #  signal.blackmanharris(self.NR)
                #pass
                #self.DATA[istk, ie, :] *= signal.hann(self.NR)
                #self.DATA[istk, ie, :] *= signal.blackman(self.NR)
                #DATA2[istk, ie, :] *= signal.blackmanharris(self.NR)
            self.emit(SIGNAL("updateProgress(int)"), (int)(1e2*istk/self.NS))   
    
        self.DATASTACK = np.average(self.DATAP, axis=0)
        if (self.NS > 1):
            # for noise calculation, destructively average phase-cycled echoes
            self.DATAP[1::2,:,:] *= -1.
            self.NOISE = np.average(self.DATAP, axis=0)
            self.NOISE -= np.mean(self.NOISE) # The noise can be biased? Don't understand why!!
            #self.DATA[0::2,:,:] *= -1. 
        self.emit(SIGNAL("plotWIN()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def ResampleEnvelopeData(self, Q, Mask=0, GatesPD=0):
        """ Mirror the data and then apply window function and resample in Fourier domain at twice the number. 
            Then take half. Should be better.
            GatesPD is Gates Per Decade 
        """
        
        tt = np.concatenate( [self.T2T[Mask::], self.T2T[Mask::][-1]+self.T2T[Mask::] ] )
        xx = np.concatenate( [self.T2D[Mask::][::-1] , self.T2D[Mask::] ] )
        
        xx,tx = signal.resample( xx, 2*Q, t=tt)#, window='hanning')
        self.T2D = xx[0:Q][::-1]
        #self.T2D = self.T2D[::-1]
        
        xx = np.concatenate( [self.T2N[::-1] , self.T2N] )
        xx,tx = signal.resample( xx, 2*Q, t=tt)#, window='hanning')
        self.T2N = xx[0:Q][::-1]

        # don't reverse times dummy!
        self.T2T = tx[0:Q] #[::-1] ## <-- No! bad Trevor
        self.NE = len(self.T2D)

    def computeSTD(self):
        if self.NS >= 2:
            self.sigma = np.std( np.imag(self.T2N) )
        else:
            #N = np.real(self.T2D)
            s0 = np.std( np.real(self.T2D) )
            #sr = ( np.std( (np.exp(-1j*self.zeta)*N).real ))  
            #si = ( np.std( (np.exp(-1j*self.zeta)*N).imag ))  
            print("s0", s0)
            self.sigma = s0 # np.array((.15)) #s0 #+ (sr+si-s0)/2.   
            #self.phasegain =  (s0 + (sr+si-s0)/2) / s0
            #print("ratio",  s0,  (s0+(sr+si-s0)/2.)/s0  )
            #nr = (np.exp(-1j*self.zeta)*N).real
            #ni = (np.exp(-1j*self.zeta)*N).imag
            #print ("phasegain", self.phasegain)
            #self.T2D.real *= self.phasegain 
            #print("SIGMA CALCULATION HERE", s0, self.sigma)
            
    def computeSTDBurst(self):
        if self.NS >= 2:
            self.sigmaBurst = np.std( np.imag(self.T2Nb) )
        else:
            N = np.real(self.T2Db)
            s0 = np.std( np.real(self.T2Db) )
            sr = ( np.std( (np.exp(-1j*self.zeta)*N).real) )  
            si = ( np.std( (np.exp(-1j*self.zeta)*N).imag) )  
            self.sigmaBurst = s0 #+ (sr+si-s0)/2.   
            #self.T2Db.real *= self.sigmaBurst/s0 #nr #(nr+ni)/2.          
        #self.sigmaBurst = np.std( np.real(self.T2Db) ) 
    
    def gateIntegrateBurst(self, gpd, stackEfficiency):

        self.computeSTDBurst()

        # calculate total number of decades
        nd = np.log10(self.T2Tb[-1]/self.T2Tb[0])   #np.log10(self.T2T[-1]) -  np.log10(self.T2T[-1])
        tdd = np.logspace( np.log10(self.T2Tb[0]), np.log10(self.T2Tb[-1]), (int)(gpd*nd)+1, base=10, endpoint=True) 
        tdl = tdd[0:-1]     # these are left edges
        tdr = tdd[1::]      # these are left edges
        td = (tdl+tdr) / 2. # window centres

        htd = np.zeros( len(td), dtype=complex )
        htn = np.zeros( len(td), dtype=complex )

        isum = np.zeros( len(td), dtype=int )
        
        Vars = np.ones( len(td) ) * self.sigmaBurst**2 #* .15**2

        ii = 0
        for itd in range(len(self.T2Tb)):
            if ( self.T2Tb[itd] > tdr[ii] ):
                if (ii < len(td)-1):
                    ii += 1
                else:
                    break
            isum[ii] += 1
            htd[ii]  += self.T2Db[ itd ]
            #htn[ii] += self.T2N[ itd ]
            Vars[ii] += self.sigmaBurst**2
            #self.emit(SIGNAL("updateProgress(int)"), (int)(1e2*itd/len(self.T2T)))  
        
        # The 1.5 should be 2 in the theory, the 1.5 compensates for correlated noise, etc. Bascically 
        #       we don't see noise reduce like ideal averaging. 
        self.sigmaBurst = np.sqrt( Vars  * (1/isum)**stackEfficiency ) # Var (aX) = a^2 Var(x)

        # Bootstrap
        bootstrap = False
        if bootstrap:
            Means, isumbs = self.bootstrapWindows(np.real(self.T2Db), 300, isum)
            ts,sm = self.smooth( isumbs, np.std(Means, axis=1) )
            for it in range(ii):
                self.sigmaBurst[it] = sm[isum[it]]

        # RESET times where isum == 1
        ii = 0
        while (isum[ii] == 1):
            td[ii] = self.T2Tb[ii]
            ii += 1

        htd /= isum
        htn /= isum

        self.T2Tb = td
        self.T2Db = htd

 
    def gateIntegrate(self, gpd, stackEfficiency):
        """ 
            Gate integrate the signal to gpd, gates per decade
        """
        self.computeSTD()
        self.gpd = gpd

        if self.burst:
            self.gateIntegrateBurst(gpd, stackEfficiency)

        # calculate total number of decades
        nd = np.log10(self.T2T[-1]/self.T2T[0])   #np.log10(self.T2T[-1]) -  np.log10(self.T2T[-1])
        tdd = np.logspace( np.log10(self.T2T[0]), np.log10(self.T2T[-1]), (int)(gpd*nd)+1, base=10, endpoint=True) 
        tdl = tdd[0:-1]     # these are left edges
        tdr = tdd[1::]      # these are left edges
        td = (tdl+tdr) / 2. # window centres

        htd = np.zeros( len(td), dtype=complex )
        htn = np.zeros( len(td), dtype=complex )
        isum = np.zeros( len(td), dtype=int )
        Vars = np.zeros( len(td) ) 

        ii = 0
        for itd in range(len(self.T2T)):
            if ( self.T2T[itd] > tdr[ii] ):
                if (ii < len(td)-1):
                    ii += 1
                else:
                    break
            isum[ii] += 1
            htd[ii]  += self.T2D[ itd ]
            
            #htn[ii] += self.T2N[ itd ]
            Vars[ii] += self.sigma**2
            self.emit(SIGNAL("updateProgress(int)"), (int)(1e2*itd/len(self.T2T)))  
        
        # Theoretical gates 
        # The 1.5 should be 2 in the theory, the 1.5 compensates for correlated noise, etc. Bascically 
        #       we don't see noise reduce like ideal averaging. 
        self.sigma = np.sqrt(Vars  * (1./isum)**stackEfficiency ) # Var (aX) = a^2 Var(x)

        # Bootstrap
        bootstrap = False # False
        bs = open("bootstrap.dat", "a")
        if bootstrap:
            nboot=10000
            #Means, isumbs = self.bootstrapWindows(np.real(self.T2D), 100, isum[isum!=1])
            Means, isumbs = self.bootstrapWindows(np.real(self.T2D), nboot, np.arange(1,isum[-1]+1))
            #ts,sm = self.smooth(isumbs, np.std(Means - np.tile(np.mean(Means,axis=1),(nboot,1)).T  , axis=1, ddof=2)) # unbiased
            # TODO use MAD instead of STD??
            ts,sm = self.smooth(isumbs, np.std(Means, axis=1, ddof=1)) # unbiased
            #scale = (1. + (nboot*isum)/(np.ones(ii+1)*len(self.T2D)))**.125
            #print (scale) 
            for item in (sm[isum[isum!=1]-1]-self.sigma[isum!=1]): #/self.sigma[isum!=1]:
                bs.write( str(item) + " " )
            bs.write("\n")
            #ts,sm = self.smooth( isumbs, np.std(Means, axis=1))
            #bias = np.average(Means,axis=1)
            for it in range(len(self.sigma)):
                self.sigma[it] = sm[isum[it]-1]
            print('sigma boot', self.sigma)
        # RESET times where isum == 1
        ii = 0
        while (isum[ii] == 1):
            td[ii] = self.T2T[ii]
            ii += 1

        htd /= isum
        htn /= isum

        self.T2T = td
        self.T2D = htd

        self.emit(SIGNAL("plotLOG10ENV()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("enableINV()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def bootstrapWindows(self, N, nboot, isum):
        """ Bootstraps noise as a function of gate width
        """
        nc = np.shape(N)[0]

        Means = {}
        Means = np.zeros( (len(isum), nboot)  )
        for ii, nwin in enumerate(isum):  
            for iboot in range(nboot):
                cs = np.random.randint(0,nc-nwin)
                #Means[ii,iboot] = np.mean( N[cs:cs+nwin] )
                Means[ii,iboot] = np.mean( np.random.normal(0, self.sigma[0], nwin))
        return Means, np.array(isum)

    def smooth(self, x, y):
        w = np.ones(len(x))
        w[0] = 100. # force first time gate
        xs = np.arange(1, x[-1]+1, .1) # resample
        #s = UnivariateSpline( np.log(x), np.log(y), s=2.5, w=w)
        s = UnivariateSpline( np.log(x), np.log(y), s=4.5, w=w)
        #s = UnivariateSpline( x, y, s=2.5, w=w)
        ys = s(np.log(xs))
        #ys = s(xs)

        # resample 
        ys = ys[0::10]    
        xs = xs[0::10]    
   
        return xs,np.exp(ys)

    def T2EnvelopeDetect(self, offset=0):

        self.NRS = (int)(2**12)  # Number of zeros to pad with, total record length
        # Zero pad
        DATASTACKZPF = np.append( self.DATASTACK, np.zeros( (self.NE, self.NRS - self.NR) ), axis=1  )
        if (self.NS > 1):
            NOISEZPF = np.append( self.NOISE, np.zeros( (self.NE, self.NRS - self.NR) ), axis=1  )
    
        # TODO check for proper split
        SPLIT = self.NR//2 + offset  # 20 # 15 #2 

        # containers for array
        self.T2D = np.zeros( (self.NE), dtype='complex' ) 
        self.T2N = np.zeros( (self.NE), dtype='complex' ) 
        self.T2T = np.zeros( (self.NE) ) 
        
        # frequency-domain processing
        for ie in range(0,self.NE):

            COLOUR = np.array([ie/self.NE, 0., 1.-ie/self.NE]) # for plots
            
            ###############################################################
            # fold data need to reverse some of these
            #FOLD = ((DATASTACK[ie,0:NR/2] + DATASTACK[ie,:][::-1][0:NR/2])/2.)[::-1] 
            #plt.figure(223)
            #plt.plot( self.TAUE*(float)(ie) + self.TIMES, np.imag(self.DATASTACK[ie]),  color=COLOUR)
            #plt.plot( self.TAUE*(float)(ie) + self.TIMES, np.imag(self.DATA[0,ie]),  color='green')
            #plt.plot( self.TAUE*(float)(ie) + self.TIMES, np.imag(self.DATA[1,ie]),  color='black')
            #plt.plot( self.TAUE*(float)(ie) + self.TIMES, np.imag(self.DATA[2,ie]),  color='cyan')
            #plt.plot( self.TAUE*(float)(ie) + self.TIMES, np.imag(self.DATA[3,ie]),  color='purple')
            #plt.plot( TAUE*(1e-6)*(float)(ie) + TIMES[SPLIT], np.imag(DATASTACK[ie])[SPLIT], 'o', markersize=6, color=COLOUR)
            #plt.plot( TAUE*DT*ie + TIMES, np.real(DATASTACK[ie]),  color=COLOUR[::-1])

            ##################################################################
            # do frequency domain workflow 
            # fold to DC
            ND = np.append( DATASTACKZPF[ie,:][SPLIT::], DATASTACKZPF[ie,:][0:SPLIT] )
            #X = 1e-2*((self.RL/2.)/self.NR) * np.fft.fft(ND, n=self.NRS) # Why 1e-2?
            X = 1e4*(( (self.DT*self.NRS) /2.)/self.NRS) * np.fft.fft(ND, n=self.NRS) # Why 1e-2?
            #X = (self.DT) * np.fft.fft(ND, n=self.NRS) # Why 1e-2?

            if (self.NS > 1):
                NND = np.append( NOISEZPF[ie,:][SPLIT::], NOISEZPF[ie,:][0:SPLIT] )
                #XN = 1e-2*((self.RL/2.)/self.NR) * np.fft.fft(NND, n=self.NRS) # Why 1e-2?
                XN = 1e4*(( (self.DT*self.NRS) /2.)/self.NRS) * np.fft.fft(NND, n=self.NRS) # Why 1e-2?
                self.T2N[ie] = XN[0] #complex( np.real(X[0]),  np.imag(X[0]) )
    
            ##################################################################
            # Save T2 records
            self.T2D[ie] = X[0] #complex( np.real(X[0]),  np.imag(X[0]) )
            self.T2T[ie] = (1.+ie)*(self.TAUE)
            self.emit(SIGNAL("updateProgress(int)"), (int)(1e2*ie/self.NE))   
       
        self.computeSTD() 
        
        self.emit(SIGNAL("plotENV()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("enableINV()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def CorrectedAmplitude(self, joe=0, frank=0):
        """ Phase corrects the record, joe and Frank are empty (ignored) and are workarounds for threading.
        """
        #[E0,df,phi,T2] = quadratureDetect(Q, I, tt)
       
        # mask first few echoes, just for detection 
        NM = 2 
        #[conv,E0,df,phi,T2] = quadratureDetect(np.imag(self.T2D[NM::]), np.real(self.T2D[NM::]), self.T2T[NM::], False, True)
        #[conv,E0,df,phi,T2] = quadratureDetect(np.imag(self.T2D[NM::]), np.real(self.T2D[NM::]), self.T2T[NM::], False, False, True)
        #                                                                              #CorrectFreq=False, BiExp=False, CorrectDC=False)
        [conv,E0,df,phi,T2] = quadratureDetect2(np.imag(self.T2D[NM::]), np.real(self.T2D[NM::]), self.T2T[NM::])
        self.zeta1 = phi

        #if conv: # failed convergence
        #    [conv,E0,df,phi,T2] = quadratureDetect(np.imag(self.T2D[NM::]), np.real(self.T2D[NM::]), self.T2T[NM::], False, False, False)
        self.T2D = self.RotateAmplitude(np.real(self.T2D), np.imag(self.T2D), phi, df, self.T2T) 
        self.computeSTD() 
        #print(str(conv) + "\t" + str(phi)) #, file = "phase.dat")
       
        if self.burst:
            #[conv,E0,df,phi,T2] = quadratureDetect(np.imag(self.T2Db), np.real(self.T2Db), self.T2Tb, False, False, False)
            [conv,E0,df,phi,T2] = quadratureDetect2(np.imag(self.T2Db), np.real(self.T2Db), self.T2Tb) #, False, False, False)
            #if conv: # failed convergence
            #    [conv,E0,df,phi,T2] = quadratureDetect(np.imag(self.T2Db), np.real(self.T2Db), self.T2Tb, False, False, False)
            env = E0*np.exp(-np.array(self.T2Tb)/T2)
            self.T2Db = self.RotateAmplitude(np.real(self.T2Db), np.imag(self.T2Db), phi, df, self.T2Tb)
            self.computeSTDBurst() 
        
        self.emit(SIGNAL("plotENV()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("enableINV()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def RotateAmplitude(self, X, Y, zeta, df, t):
        self.zeta = zeta 
        self.df = df 
        V = X + 1j*Y
        return np.abs(V) * np.exp(1j * (np.angle(V) - zeta - 2.*np.pi*df*t)) 

    def MonoFit(self, mask=2, intercept="False"):
        component='imag'
        #print "MonoRes", component, mask, intercept
        a0 = 0
        if intercept=="True":
            if component == "imag":
                [a0,b0,rt20] =  regressCurve(np.imag(self.T2D[mask::]), self.T2T[mask::], intercept=True)  
            if component == "real":
                [a0,b0,rt20] =  regressCurve(np.real(self.T2D[mask::]), self.T2T[mask::], intercept=True)  
            rfit =  r"$\tilde{T}_2$= %4d [ms]"%(1e3*rt20)

        else:
            if component == "imag":
                [b0,rt20] =  regressCurve(np.imag(self.T2D[mask::]), self.T2T[mask::], intercept=False) #, intercept=True) # 
            if component == "real":
                [b0,rt20] =  regressCurve(np.real(self.T2D[mask::]), self.T2T[mask::], intercept=False) #, intercept=True) # 
            rfit =  r"$\tilde{T}_2$= %4d [ms]"%(1e3*rt20)
        

        extrapolate = False #False #True #False 
        if extrapolate:
            Textr = 3
            nede = np.log10(Textr/self.T2T[-1]) #np.log10(self.T2T[-1]) -  np.log10(self.T2T[-1])
            tex = np.logspace( np.log10(self.T2T[-1]), np.log10(Textr), (int)(self.gpd*nede)+1, base=10, endpoint=True) 
            dex = 1j*  (a0 + b0*np.exp(-np.array(tex)/rt20))
            self.T2T = np.concatenate( (self.T2T, tex) )            
            self.T2D = np.concatenate( (self.T2D, dex) )            
            self.sigma = np.concatenate( (self.sigma, self.sigma[-1]*np.ones(len(dex)) ) )            
            
        #env = b0*np.exp(-np.array(self.T2T)/rt20)
        #self.MonoRes = { 'env':env, 'rfit':rfit, 'b0':b0, 'rt20' :rt20 } #[a0,b0,rt20] }
        env = a0 + b0*np.exp(-np.array(self.T2T)/rt20)
        self.MonoRes = { 'env':env, 'rfit':rfit, 'a0':a0, 'b0':b0, 'rt20' :rt20, 't':self.T2T } #[a0,b0,rt20] }


        self.emit(SIGNAL("plotMONO()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("enableINV()"))   
        self.emit(SIGNAL("doneStatus()"))  

    def BiFit(self, mask=2, intercept="False"):
        component='imag'
        #print "MonoRes", component, mask, intercept
        a0 = 0
        if intercept=="True":
            if component == "imag":
                [a0,b0,rt20,bb0,rrt20] =  regressCurve2(np.imag(self.T2D[mask::]), self.T2T[mask::], sigma2=self.sigma[mask::], intercept=True)  
            if component == "real":
                [a0,b0,rt20,bb0,rrt20] =  regressCurve2(np.real(self.T2D[mask::]), self.T2T[mask::], sigma2=self.sigma[mask::], intercept=True)  
            rfit =  r"$\tilde{T}_2$= %4d [ms]"%(1e3*rt20)

        else:
            if component == "imag":
                [b0,rt20,bb0,rrt20] =  regressCurve2(np.imag(self.T2D[mask::]), self.T2T[mask::], sigma2=self.sigma[mask::], intercept=False) #, intercept=True) # 
            if component == "real":
                [b0,rt20,bb0,rrt20] =  regressCurve2(np.real(self.T2D[mask::]), self.T2T[mask::], sigma2=self.sigma[mask::], intercept=False) #, intercept=True) # 
            rfit =  r"$\tilde{T}_2$= %4d [ms]"%(1e3*rt20)
        

        extrapolate = True # True #True 
        if extrapolate:
            Textr = 3
            nede = np.log10(Textr/self.T2T[-1]) #np.log10(self.T2T[-1]) -  np.log10(self.T2T[-1])
            tex = np.logspace( np.log10(self.T2T[-1]), np.log10(Textr), (int)(self.gpd*nede)+1, base=10, endpoint=True) 
            dex = 1j*  (a0 + b0*np.exp(-np.array(tex)/rt20) + bb0*np.exp(-np.array(tex)/rrt20) )
            self.T2T = np.concatenate( (self.T2T, tex) )            
            self.T2D = np.concatenate( (self.T2D, dex) )            
            self.sigma = np.concatenate( (self.sigma, self.sigma[-1]*np.ones(len(dex)) ) )            
            
        #env = b0*np.exp(-np.array(self.T2T)/rt20)
        #self.MonoRes = { 'env':env, 'rfit':rfit, 'b0':b0, 'rt20' :rt20 } #[a0,b0,rt20] }
        env = a0 + b0*np.exp(-np.array(self.T2T)/rt20) + bb0*np.exp(-np.array(self.T2T)/rrt20)

        #print ("bi fits", a0, b0, rt20, bb0, rrt20)

        self.BiRes = { 'env':env, 'rfit':rfit, 'a0':a0, 'b0':b0, 'rt20' :rt20, 'bb0':bb0, 'rrt20':rrt20, 't':self.T2T } #[a0,b0,rt20] }

        self.emit(SIGNAL("plotBI()"))   
        self.emit(SIGNAL("enableDSP()"))   
        self.emit(SIGNAL("enableINV()"))   
        self.emit(SIGNAL("doneStatus()")) 
        
 
    def DistFit(self, dc=0, mask=0, nT2=30, lowT2=.005, hiT2=3.25, spacing="Log_10", Smooth="Smallest", betaScale=1):
        
        Time = pwcTime()
        Time.setT2(lowT2, hiT2, nT2, spacing) 
        Time.setSampling( self.T2T[mask:] ) 
        Time.generateGenv()

        #self.burst = False
        #print ("MRProc.py ~~ DistFit self.burst override for publication ", self.burst)
        if self.burst:
            # fit burst data first
            Timeb = pwcTime()
            #Timeb.setT2(lowT2, hiT2, nT2, spacing) 
            Timeb.T2Bins = np.copy(Time.T2Bins)
            #if len(Timeb.T2Bins) > 20:
            #    Timeb.T2Bins = Timeb.T2Bins[0:20] # Subset
            Timeb.setSampling( self.T2Tb ) 
            Timeb.generateGenv()

            # invert burst data, always use smooth to encourage fast time inversion
            modelb  = logBarrier(Timeb.Genv, np.imag(self.T2Db)-dc, Timeb.T2Bins, MAXITER=500, sigma=betaScale*self.sigmaBurst, alpha=1e10, smooth=Smooth) 
            modelb2 = np.zeros(nT2)
            modelb2[0:len(modelb)] += modelb 
            #model = modelb2 # TODO remove 
 
            # invert regular data independently, sum inversions 
            #model  = logBarrier(Time.Genv, np.imag(self.T2D[mask:])-dc, MAXITER=500, x_0=modelb, sigma=betaScale*self.sigma[mask:], alpha=1e10, smooth=Smooth)
            #model += modelb 
            
            # Pass burt result as reference model
            model  = logBarrier(Time.Genv, np.imag(self.T2D[mask:])-dc, Time.T2Bins, MAXITER=500, xr=modelb2, sigma=betaScale*self.sigma[mask:], alpha=1e10, smooth=Smooth)

        else:
            model  = logBarrier(Time.Genv, np.imag(self.T2D[mask:])-dc, Time.T2Bins, MAXITER=500, sigma=betaScale*self.sigma[mask:], alpha=1e10, smooth=Smooth) #, smooth=True) #

        # forward model
        env = np.dot(Time.Genv, model) 
        # NOTE this is not thread safe or even thread aware. Need to use MPI style message passing 
        self.DistRes = { 'env':env, 'mod':model, 'Time':Time } #[a0,b0,rt20] } no dice
    
    def DistFitMP(self, q, tag, dc=0, mask=0, nT2=30, lowT2=.005, hiT2=3.25, spacing="Log_10", Smooth="Smallest", betaScale=1):
        """ Multi-processor version of DistFit using Queue. Note that Queue doesn't 
            tolerate large messages. So instead we write them to file, using pickle! 
            It's semi-absurd. But seems to function fast enough. 
        """
        import pickle 
        self.DistFit(dc, mask, nT2, lowT2, hiT2, spacing, Smooth, betaScale)
        self.DistRes['tag'] = tag
        pickle.dump( self.DistRes, open( str(tag)+".p", "wb" ) )
        q.put( tag )
    

if __name__ == "__main__":

    fileName, fileExtension = os.path.splitext(sys.argv[1])
    if fileExtension == ".ors":

        ORS = MRProc()
        ORS.loadORSFile(sys.argv[1])
        ORS.WindowAndStack()
        ORS.T2EnvelopeDetect()
        
        #ORS.ResampleEnvelopeData(200, Mask=2, GatesPD=20)  # resample size 
        
        #print (len(ORS.T2D))
        #exit()
        #mono, rfit, [a0,b0,mt2] = ORS.MonoFit()
        mono, rfit, [b0,mt2] = ORS.MonoFit()
        #ca, caEnv, CAS = ORS.CorrectedAmplitude()
        mymask = 2
        env, mod, Time = ORS.DistFit(0, mymask) #a0)

        # Subtract DC term before 2D fit? Kind of kludgy but I don't see a way around it right now. 
        plt.figure()
        #plt.plot(ORS.T2T, np.imag(ca), label='imag ca')
        #plt.plot(ORS.T2T, np.real(ca), label='real ca')
        #if ORS.NS < 2:
        #    plt.plot(ORS.T2T, np.imag(ORS.T2D), label='imag')
        #    plt.plot(ORS.T2T, np.real(ORS.T2D), label='real')
        #else:
        plt.errorbar(ORS.T2T, np.imag(ORS.T2D), fmt='o', yerr=ORS.sigma, label="data imag")
        plt.errorbar(ORS.T2T, np.real(ORS.T2D), fmt='o', yerr=np.std(ORS.T2D.real), label="data real")
        if ORS.NS > 2:
            plt.errorbar(ORS.T2T, np.imag(ORS.T2N), fmt='o', yerr=ORS.sigma, label="noise imag")
        
        plt.plot(ORS.T2T, mono, color='cyan', linewidth=3, label=rfit)
        #plt.plot(ORS.T2T, caEnv, label="car")
        plt.plot(ORS.T2T[mymask::], env, color='magenta', linewidth=3, label="dist")
        #plt.plot(ORS.T2T, a0+env)
        plt.legend()
        plt.savefig(sys.argv[1].strip(".ors")+"_caenv.pdf", facecolor='white', edgecolor='white', dpi=200)


        # Report RMS error of two results 
        print("RMS error mono-exponential ", np.linalg.norm(np.imag(ORS.T2D) - mono) )
        print("RMS error dist-fit         ", np.linalg.norm(np.imag(ORS.T2D[mymask::]) - env) )
        print("RMS error real  channel    ", np.linalg.norm(np.real(ORS.T2D) ) )
        if ORS.NS > 2:
            print("RMS error noise channel    ", np.linalg.norm(np.imag(ORS.T2N) ) )

        plt.figure()
        plt.plot( ORS.T2T[mymask::],  np.imag(ORS.T2D[mymask::]) - env )
        plt.title("residuals")

        #plt.show()
        #exit()

        plt.figure() 
        plt.plot(Time.T2Bins, mod, color='purple', linewidth=2, label="recovered")
        #plt.plot(Time.T2Bins, guess, color='red', linewidth=2, label="guess")

        axmine = plt.gca()
        axmine.set_xlabel(r"$T_2$ [s]", color='black')
        axmine.set_ylabel(r"$A_0$ [rku]", color='black')
        axmine2 = plt.twinx()
        axmine2.set_ylabel(r"$A_0$ [rku]", color='black')

        theta = np.sum(mod)
        LogMeanT2 = np.exp(np.sum( mod * np.log(Time.T2Bins) ) / theta ) 

        #axmine2.plot(mt2, a0+b0, 'o', markersize=8, label="mono")
        axmine2.plot(mt2, b0, 'o', markersize=8, label="mono")
        #axmine2.plot(Time.T2Bins, guess, '-',color=colour[ic], linewidth=2, label="total")
        axmine2.plot(LogMeanT2, theta, 's', markersize=8, label="$T_{2ML}$")
        axmine2.set_ylim([0, axmine2.get_ylim()[1]])
        leg = plt.legend()

        plt.savefig(sys.argv[1].strip(".ors")+"_2DT2.pdf", facecolor='white', edgecolor='white', dpi=200)

   
        plt.figure(23)
        stats.probplot( np.imag(ORS.T2D) - mono , dist="norm", plot=plt) #, label="mono resid")
        stats.probplot( np.imag(ORS.T2D[mymask::]) - env , dist="norm", plot=plt) #, label="dist resid")
        if ORS.NS > 2:
            stats.probplot( np.imag(ORS.T2N), dist="norm", plot=plt) #, label="noise" )

        plt.savefig(sys.argv[1].strip(".ors")+"_qq.pdf")
 
        plt.show()
        exit()