# -*- coding: utf-8 -*-
"""
Created on Fri Sep 28 17:16:12 2012

@author: Edmondo
"""
import numpy
import matplotlib.pyplot as plt
from scipy.signal import hilbert, resample, medfilt, lfilter, butter, buttord, filtfilt
from scipy.integrate import cumtrapz
from scipy.interpolate import RectBivariateSpline, InterpolatedUnivariateSpline
#from ftubase import smoothd

import inspect

class _AxInfo(object):
    """Keeps the axis label and unit in a single class.
    """
    def __init__(self, name = '', label=None, unit=None):
        self.name = name
        self.label = label
        self.unit = unit
        
    def addlu(self, label, unit):
        """Change the label and the unit to the current object
        """
        self.label = label
        self.unit = unit
        
    def parse1d(self, exprl):
        """The expression coming from a 1d FTU channel is decomposed
        in label and unit
        """
        xxl = str(exprl).split('(')
        self.label = xxl[0]
        self.unit = xxl[1].split(')')[0]
        
    def _exposednames(self):
        """ expose the name of the label and unit.
        It basically correspond to self.name+'label and self.name+'unit'
        A dictionary is returned
        """
        return {'labelname':self.name+'label', 
                'unitname':self.name+'unit'}
    def getnames(self):
        """Return a dictionary whith the label and utni name associated
        with self. It should be used in the methid __getattr__ of Ftudata
        in order to understand which _AxInfo should be called to get the
        proper label
        """
        names = self._exposednames()
        return {names['labelname']:self, names['unitname']:self}
        
    def torstr(self):
        """A list of string representations are returned. To be used in the 
        __repr__() method of Ftudata
        """
        rstr = []
        if self.label is not None:
            names = self._exposednames()
            rstr += ["%8s: %s" % (names['labelname'], self.label)]
            rstr += ["%8s: %s" % (names['unitname'], self.unit)]
        return rstr
    
    def __getitem__(self, name):
        names = self._exposednames()
        if name == names['labelname']:
            return self.label
        elif name == names['unitname']:
            return self.unit
        raise AttributeError('No {0} attribute is present'.format(name))

class V2d(object):
    """It defines object with an R a Z and a V fields, various operation are 
    defined on them.
    """
    def __init__(self, r=None, t=None, v=None, ier=None):
        self.r = r
        self.t = t
        self.v = v
        if ier is None:
            ier = 0
        self.ier = ier

    def copy(self):
        """return a copy of the given object
        """
        res = self._getacopy()
        return res
        
    def _getacopy(self, copydata=True, inplace=False):
        """Return a copy of the given object.
        copydata: if True also the field v is copied
        inplace:  if True the actual object is reeturned isntead of a copy
        """
        if inplace:
            return self
            
        other = self.__class__()
        if self.r is not None:
            other.r = self.r.copy()
        if self.t is not None:
            other.t = self.t.copy()
        other.ier = self.ier
        if copydata and self.v is not None:
            other.v = self.v.copy()
        return other
                                      
    def __repr__(self):
        try:
            nr = len(self.r)
        except TypeError:
            nr = 0
            
        try:
            nt = len(self.t)
        except TypeError:
            nt = 0
            
        rstr = []
        rstr += ["    v: Float %dx%d" % (nt, nr)]
        rstr += ["    t: Float  %d" % nt]
        rstr += ["    r: Float  %d" % nr]
        rstr += ["  ier:        %d" % self.ier]
        return "\n".join(rstr)

    def ok(self):
        """Return True if the object contains correct data
        """
        return self.ier == 0 and self.v.size > 0
        
    def vtline(self, *arg, **args):
        """Plot a vertical line for each time which is present 
        """
        for t in self.t:
            plt.axvline(t, *arg, **args)
            
    def plot(self, t = numpy.nan, r = numpy.nan, func=None, ax=None, **args):
        """Plot the data, various keyword are defined:
            t:  if not is NAN the time slice closest to the given time is plotted 
        """
        if ax is None:
            ax = plt
            
        if func is None:
            func = lambda x:x
            
        if (numpy.isnan(t) and numpy.isnan(r)):
            temp = ax.plot(self.t, func(self.v), **args)
        elif (numpy.isnan(t)):
            ir = numpy.argmin(abs(self.r - r))
            temp = ax.plot(self.t, func(self.v[:,ir]), **args)    
        elif (numpy.isnan(r)):
            it = numpy.argmin(abs(self.t - t))
            temp = ax.plot(self.r, func(self.v[it,:]),  **args)
        return temp
        
    def get(self, r=numpy.nan, t=numpy.nan):
        if (numpy.isnan(t) and numpy.isnan(r)):
            return (self.r, self.t, self.v)
        elif (numpy.isnan(t)):
            ir = numpy.argmin(abs(self.r - r))
            if len(self.v.shape) == 1:
                return (self.t,self.v)
            else:
                return (self.t, self.v[:,ir])
        elif (numpy.isnan(r)):
            it = numpy.argmin(abs(self.t - t)) 
            if len(self.v.shape) == 1:
                return (self.r, self.v[it])
            else:
                return (self.r, self.v[it,:])
        else:
            ir = numpy.argmin(abs(self.r - r))
            it = numpy.argmin(abs(self.t - t))
            if len(self.v.shape) == 1:
                return self.v[it]
            else:
                return self.v[it,ir]
            
    def slice(self,  r=numpy.nan, t=numpy.nan):
        
        result = self._getacopy(copydata=False)

        if not self.ok():
            return result
        if (numpy.isnan(t) and numpy.isnan(r)):
            return result
        elif (numpy.isnan(t)):
            ir = numpy.argmin(abs(self.r - r))
            result.r = self.r[ir][numpy.newaxis]
            result.v = self.v[:,ir][:,numpy.newaxis]
        elif (numpy.isnan(r)):
            it = numpy.argmin(abs(self.t - t)) 
            result.t = self.t[it][numpy.newaxis]
            result.v = self.v[it,:][numpy.newaxis,:]
        else:
            ir = numpy.argmin(abs(self.r - r))
            it = numpy.argmin(abs(self.t - t))
            result.r = self.r[ir][numpy.newaxis]
            result.t = self.t[it][numpy.newaxis]
            result.v = self.v[it,ir][numpy.newaxis,numpy.newaxis]
        return result
        
    def get_r_profile(self, t = numpy.nan):
        it = numpy.argmin(abs(self.t - t))
        return (self.r, self.v[it,:])

    def get_t_profile(self, r = numpy.nan):
        ir = numpy.argmin(abs(self.r - r))
        return (self.t, self.v[:,ir])

    def trange(self, tlim, inplace=False):
        """Restrict the time axis in the specified range
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        if not self.ok():
            return result
        
        tlimm = numpy.asarray(tlim, dtype=float)
        if tlimm.size == 2:
            idt = (self.t >= numpy.min(tlimm)) & (self.t <= numpy.max(tlimm))
            result.t = self.t[idt]
            if len(self.v.shape)==1:
                result.v = self.v[idt]
            else:
                result.v = self.v[idt,:]
        return result
         
    def tresample(self, n, inplace=False):
        """Resample the data along the time axis.
        Calling:
            obj.tresample(n,inplace=[True,False])
            n specify the number of time points it will get 
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        v, t = resample(self.v, n, t=self.t, axis=0)
        result.v = v
        result.t = t
        return result
        
    def mean(self, trange):
        if not self.ok():
            return None
        trange = numpy.asarray(trange, dtype=float)
        idt = (self.t >= numpy.min(trange)) & (self.t <= numpy.max(trange))
        if len(self.v.shape) == 2:
            return numpy.mean(self.v[idt,:])
        else:
            return numpy.mean(self.v[idt])

    def smooth(self, inplace=False):
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
        
        if len(self.v.shape) == 1:
            #result.v, rough = smoothd(self.v)
            pass
        elif len(self.v.shape) == 2:
            vv = numpy.zeros(self.v.T.shape)
            for v,sv in zip(vv,self.v.T):
             #   v[:], rough = smoothd(sv)
               pass
            result.v = vv.T
        return result 
            
    def medfilt(self, nm=5, inplace=True):
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
        if len(self.v.shape) == 1:
            result.v = medfilt(self.v, kernel_size=nm)
        elif len(self.v.shape) == 2:
            vv = numpy.zeros(self.v.T.shape)
            for v,sv in zip(vv,self.v.T):
                v[:] = medfilt(sv, kernel_size=nm)
            result.v = vv.T
        return result 
    
    def specgram(self, ax=None, *args, **kwargs):
        if self.v is None:
            return
        if ax is None:
            ax = plt
        dt = numpy.mean(numpy.diff(self.t))
        Fs = 1.0 / dt
        
        xextent = [self.t.min(), self.t.max()] # not correct TODO 
        ax.specgram(self.v, *args, Fs=Fs, xextent = xextent, **kwargs)
        
    def filter(self, Wn, ftype='low', inplace=False, symmetric=True):
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
        
        dt = numpy.mean(numpy.diff(self.t))
        
        FNy2 = 1.0 / 2.0 / dt
        Wno = Wn / FNy2
        if ftype == 'low':
            bord =  buttord(Wno,Wno*1.7,1,10)
        elif ftype == 'high':
            bord =  buttord(Wno*0.6,Wno,1,10)
        else:
            raise  Exception('Ma che fai')
            
        bf,af = butter(*bord,btype=ftype)
        if len(self.v.shape) == 1:
            if symmetric:
                result.v = filtfilt(bf,af, self.v)
            else:
                result.v = lfilter(bf, af, self.v)
        elif len(self.v.shape) == 2:
            result.v = lfilter(bf, af, self.v, axis=0)
        return result

    def integrate(self, Wn=None, inplace=False, symmetric=True):
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
        
        if len(self.v.shape) == 1:
            ii = cumtrapz(self.v, x=self.t)
            result.v = numpy.hstack((numpy.zeros(1),ii))
        else:
            assert(False)
        if Wn is not None:
            result.filter(Wn,ftype='high',inplace=True, symmetric=symmetric)
        return result
            
    def hilbert(self, inplace=False):
        """return the hilbert transform of the signal.
        The original signal is left unchanged.
        """
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
            
        if 'axis' in inspect.getargspec(hilbert)[0]:
            result.v = hilbert(self.v, axis=0)
        else:
            if len(self.v.shape) == 1:
                result.v = hilbert(self.v)
            elif len(self.v.shape) == 2:
                vv = numpy.zeros(self.v.T.shape,dtype=numpy.complex)
                for v,sv in zip(vv,self.v.T):
                    v[:] = hilbert(sv)
                result.v = vv.T
        return result
        
    def packhilbert(self, packlen=500, inplace=False):
        """Return the Hilbert transform, a packet algorithm is used
        """
        if self.v is None:
            return
        result = self._getacopy(copydata=False, inplace=inplace)
        nt = self.v.shape[0]
        npacket = ((nt - 1) // packlen) + 1
        for i in range(npacket):
            ibase = numpy.clip(packlen*(i - 1), 0, nt-1)
            iplus = numpy.clip(packlen*(i + 2), 0, nt)
            istart = numpy.clip(packlen*i, 0, nt-1)
            iend = numpy.clip(packlen*(i + 1), 0, nt)
    
    def clip(self, vmin, vmax, inplace=False): 
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = numpy.clip(self.v,vmin, vmax)
        return result
    
    def abs(self, inplace=False): 
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = abs(self.v)
        return result
        
    def real(self, inplace=False):
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = numpy.real(self.v)
        return result
        
    def imag(self, inplace=False):
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = numpy.imag(self.v)
        return result
        
    def angle(self, inplace=False):
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = numpy.angle(self.v)
        return result
        
    def unwrap(self, axis=0, inplace=False):
        """Convenience method for numpy.unwrap.
        Contrary to the unwrap convention, it is applied on the first
        dimension (that generally correspond to the time coordinate)
        Calling:
            obj.unwrap(axis=<axis>)
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = numpy.unwrap(self.v, axis=axis)
        return result
        
    def unifunct(self, funct=None, inplace=False, **args):
        """Apply a function to the actual data.
        Calling:
            obj.unifunct(funct=<function>, inplace=[True,False])
            funct could be a lambda expression and if not specified
            it is funct=lambda x:x
        """
        if funct is None:
            funct = lambda x:x
        result = self._getacopy(copydata=False, inplace=inplace)
        result.v = funct(self.v, **args)
        return result
        
    def difft(self, inplace=False):
        """Derivative along the t direction.
        The original signal is left unchanged.
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        vdiff = self.v[1:, :] - self.v[:-1, :]
        tdiff = self.t[1:] - self.t[:-1]
        tmea = 0.5*(self.t[1:] + self.t[:-1])
        result.t = tmea
        result.v = vdiff/tdiff[:, numpy.newaxis]
        return result
        
    def dt(self, inplace=False):
        """Derivative along the t direction.
        The original signal is left unchanged.
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        if self.v.shape == 2:
            vdt, vdr = numpy.gradient(self.v)
        else:
            vdt = numpy.gradient(self.v)
        dtt = numpy.gradient(self.t)
        result.v = vdt/dtt
        return result
        
    def dr(self, inplace=False):
        """Gradient along the radial direction
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        vdt, vdr = numpy.gradient(self.v)
        dr = numpy.gradient(self.r)
        result.v = vdr/dr
        return result
        
    def integ_r(self, inplace=False):
        """Integrate along the radial direction
        """
        result = self._getacopy(copydata=False, inplace=inplace)
        v = numpy.hstack((numpy.zeros((self.v.shape[0],1)),cumtrapz(self.v,x=self.r)))
        result.v = v
        return result
        
    @staticmethod
    def show():
        """It is equivalent to plt.show(). Just for convenience
        """
        plt.show()
         
    def movezeroto(self, tzero, inplace=False):
        """Move the zero of the t axis to tzero.
        The t coordinate is changed so that t(new) = t(old) - tzero.
        This is useful when something interesting is happening
        at different times on the same or a different channel
        """
        result = self._getacopy(inplace=inplace)
        result.t = self.t - tzero
        return result
        
    def __add__(self, other):
        result = self._getacopy(copydata=False)
        if isinstance(other, V2d):
            result.v = self.v + other.v
        else:
            result.v = self.v + other
        return result  

    def __radd__(self, other):
        result = self._getacopy(copydata=False)
        result.v = other + self.v
        return result

    def __iadd__(self, other):
        self.v = self.v + other
        return self
        
    def __sub__(self, other):
        result = self._getacopy(copydata=False)
        if isinstance(other, V2d):
            result.v = self.v - other.v
        else:
            result.v = self.v - other
        return result

    def __rsub__(self, other):
        result = self._getacopy(copydata=False)
        result.v = other - self.v
        return result

    def __isub__(self, other):
        self.v = self.v - other
        return self
        
    def __mul__(self, other):
        result = self._getacopy(copydata=False)
        if isinstance(other, V2d):
            result.v = self.v * other.v
        else:
            result.v = self.v * other
        return result
        
    def __rmul__(self, other):
        result = self._getacopy(copydata=False)
        result.v = other * self.v
        return result
        
    def __imul__(self, other):
        self.v = self.v * other
        return self

    def __div__(self, other):
        result = self._getacopy(copydata=False)
        if isinstance(other, V2d):
            result.v = self.v / other.v
        else:
            result.v = self.v / other
        return result
        
    def __rdiv__(self, other):
        result = self._getacopy(copydata=False)
        result.v = other / self.v
        return result
        
    def __idiv__(self, other):
        self.v = self.v / other
        return self

    def __pow__(self, other):
        result = self._getacopy(copydata=False)
        result.v = self.v**other
        return result
        
    def __ipow__(self, other):
        self.v = self.v ** other
        return self
        
    def __neg__(self):
        result = self._getacopy(copydata=False)
        result.v = -self.v
        return result
        
    def __pos__(self):
        return self
    
    def function(self, obj=None, r=None, t=None, **args):
        if obj is not None:
            r = obj.r
            t = obj.t
        rr = r
        tt = t
        if rr is None:
            rr = 0
        if tt is None:
            tt = 0

        if self.r is None:
            func_base = InterpolatedUnivariateSpline(self.t,self.v)
            func = lambda t, r: func_base(t)[:,numpy.newaxis] + numpy.zeros((1,numpy.asarray(r).size))
        elif self.t is None:
            func_base = InterpolatedUnivariateSpline(self.r,self.v)
            func = lambda t, r: func_base(r) + numpy.zeros((numpy.asarray(t).size,1))
        else:
            func = RectBivariateSpline(self.t, self.r, self.v, **args)
        return V2d(r,t,func(tt, rr))

class NamedV2d(V2d):
    """A V2d class with added label and unit to each axis
    """

    def __init__(self, **args):
        super(NamedV2d, self).__init__(**args)
        self.tax = _AxInfo('t')
        self.rax = _AxInfo('r')
        self.vax = _AxInfo('')

    def _getacopy(self, *arg, **args):
        other = super(NamedV2d, self)._getacopy(*arg, **args)
        if other is not self:
            other.tax = self.tax
            other.rax = self.rax
            other.vax = self.vax
        return other
        
    def __repr__(self):
        basestr = super(NamedV2d, self).__repr__()
        rstr = []
        rstr += self.tax.torstr()
        rstr += self.rax.torstr()
        rstr += self.vax.torstr()
        if len(rstr) == 0:
            return basestr
        else:
            return "\n".join(rstr)+"\n\n"+basestr

    def __getattr__(self, name):
        namedict = self.tax.getnames()
        namedict.update(self.rax.getnames())
        namedict.update(self.vax.getnames())
        if name in namedict:
            return namedict[name][name]
        else:
            raise AttributeError('No {0} attribute is present'.format(name))
