# -*- coding: utf-8 -*-
import os
import csv
import numpy as np
import matplotlib.pyplot as plt
import random as rd
import math

class Splinefit:
    """Spline fit tool

    Reference of the algorithm:
    Christian H. Reinsch, Numerische Mathematik 10, 177-183 (1967)

    """

    _x = None
    _y = None
    _dy = None
    _xx = None
    _yy = None

    _t = None
    _mag = None
    _dmag = None
    _tt = None
    _magmag = None

    _tt_default = np.logspace(0,6,61)

    _f = None
    _df = None
    _ff = None

# =====================================================================
# =====================================================================
# Private methods
# =====================================================================
# =====================================================================

    def _clear_xydyxxyy(self):
        self._x = None
        self._y = None
        self._dy = None
        self._xx = None
        self._yy = None

    def _clear_tmagdmagttmagmag(self):
        self._t = None
        self._mag = None
        self._dmag = None
        self._tt = None
        self._magmag = None
        self._f = None
        self._df = None
        self._ff = None

    def _load_data(self,full_filename:str, col_t1:int, col_t2:int, col_mag:int, col_dmag:int):
        """
        """
        # =================
        path, filename = os.path.split(full_filename)
        rootname, extension = os.path.splitext(filename)
        # -------------------------------
        method = ""
        if (extension=='.txt'):
            method = "dsv" ; # Delimiter Separator Values
            delimiter = ' '
        if (extension=='.csv'):
            method = "dsv" ; # Delimiter Separator Values
            delimiter = ','
        # -------------------------------
        if (method == "dsv"):
            try:
                with open(full_filename, 'r', encoding='utf-8', newline='') as fid:
                    reader = csv.reader(fid, delimiter=delimiter, skipinitialspace=True)
                    try:
                        self._clear_tmagdmagttmagmag()
                        klig = 0
                        self._t=[]
                        self._mag=[]
                        self._dmag=[]
                        for ligne in reader:
                            klig+=1
                            ncol = len(ligne)
                            if (ligne[0])[0]=="#":
                                continue
                            if ncol<3:
                                continue
                            t1 = float(ligne[col_t1])
                            t2 = float(ligne[col_t2])
                            mag = float(ligne[col_mag])
                            dmag = float(ligne[col_dmag])
                            self._t.append((t1+t2)/2.)
                            self._mag.append(mag)
                            self._dmag.append(dmag)
                            # print(f"ligne={ligne} lig={klig} ncol={ncol}")
                    except csv.Error as e:
                        print(f'Problem file {full_filename}, line {klig}: {e}')
                    self._t = np.array(self._t)
                    self._mag = np.array(self._mag)
                    self._dmag = np.array(self._dmag)
                    self._clear_xydyxxyy()
                return ""
            except OSError:
                print("Problem with file "+full_filename)

# =====================================================================
# =====================================================================
# Private methods getter/setter
# =====================================================================
# =====================================================================

    def _set_x(self, x:np.ndarray):
        self._x = np.array(x)
        self._clear_tmagdmagttmagmag()

    def _get_x(self):
        return self._x

    def _set_y(self, y:np.ndarray):
        self._y = np.array(y)
        self._clear_tmagdmagttmagmag()

    def _get_y(self):
        return self._y

    def _set_dy(self, dy:np.ndarray):
        self._dy = np.array(dy)
        self._clear_tmagdmagttmagmag()

    def _get_dy(self):
        return self._dy

    def _set_xx(self, xx:np.ndarray):
        self._xx = np.array(xx)
        self._clear_tmagdmagttmagmag()

    def _get_xx(self):
        return self._xx

    def _set_yy(self, yy:np.ndarray):
        pass

    def _get_yy(self):
        return self._yy

    def _set_t(self, t:np.ndarray):
        self._t = np.array(t)
        self._clear_xydyxxyy()

    def _get_t(self):
        return self._t

    def _set_mag(self, mag:np.ndarray):
        self._mag = np.array(mag)
        self._clear_xydyxxyy()

    def _get_mag(self):
        return self._mag

    def _set_dmag(self, dmag:np.ndarray):
        self._dmag = np.array(dmag)
        self._clear_xydyxxyy()

    def _get_dmag(self):
        return self._dmag

    def _set_tt(self, tt:np.ndarray):
        self._tt = np.array(tt)
        self._clear_xydyxxyy()

    def _get_tt(self):
        return self._tt

    def _set_magmag(self, magmag:np.ndarray):
        pass

    def _get_magmag(self):
        return self._magmag

    def _set_tt_default(self, tt_default:np.ndarray):
        self._tt_default = np.array(tt_default)

    def _get_tt_default(self):
        return self._tt_default

    def _set_f(self, f:np.ndarray):
        pass

    def _get_f(self):
        return self._f

    def _set_df(self, df:np.ndarray):
        pass

    def _get_df(self):
        return self._df

    def _set_ff(self, ff:np.ndarray):
        pass

    def _get_ff(self):
        return self._ff

# =====================================================================
# =====================================================================
# Public Property methods
# =====================================================================
# =====================================================================

    x = property(_get_x, _set_x)
    y = property(_get_y, _set_y)
    dy = property(_get_dy, _set_dy)
    xx = property(_get_xx, _set_xx)
    yy = property(_get_yy, _set_yy)

    t = property(_get_t, _set_t)
    mag = property(_get_mag, _set_mag)
    dmag = property(_get_dmag, _set_dmag)
    tt = property(_get_tt, _set_tt)
    magmag = property(_get_magmag, _set_magmag)

    f = property(_get_f, _set_f)
    df = property(_get_df, _set_df)
    ff = property(_get_ff, _set_ff)

# =====================================================================
# =====================================================================
# Methods for users
# =====================================================================
# =====================================================================

    def fit(self, x:np.ndarray, y:np.ndarray, dy:(float, np.ndarray), s:float, xx:np.ndarray) -> np.ndarray:
        """Compute an array which fit points defined by arrays x and y.

        Args:
            x: Input array of x coordinates
            y: Input array of y coordinates
            dy: Input uncertainties for y coordinates. If only one value is given, it is applied to all elements of y. Else use an array of the same number of elements to define different uncertainties for each element of y.
            s: The smooth parameter. Should be equal to the number of y elements if dy is well defined.
            xx: Output array of x coordinates. Be careful, the xx values must be inside the mini,maxi of the x array. No extrapolation is possible.

        Returns:
            The array of y coordinates fit along the xx coordinates.

        """
        # x[1..n1..n2]
        # y[1..n1..n2]
        # dy[1..n1..n2] 1 sigma errors (do not put zeros !!!)
        # s = smooth parameter >=0 (=nb x points if dy are 1 sigma errors)
        # xx[1..nn] vector of x values to compute
        # xx[1]>=x[n1+2] and xx[nn]<=x[n2-1]
        # Internal indexes start at 1
        # x and xx are sorted before the calculation
        #
        # verify the dimensions
        nx=len(x)
        ny=len(y)
        if isinstance(dy,(int,float))==True:
            dy=np.array([dy],dtype=float)
        elif isinstance(dy,np.ndarray)==True:
            if self._dmag !=  None:
                if isinstance(self._dmag.tolist(),(int,float)):
                    dy=np.array([self._dmag],dtype=float)
        ndy=len(dy)
        if (nx < 3):
            raise Exception(f"Length of x ({nx} counts) must be >= 3.")
        if (ny != nx):
            raise Exception(f"Length of y ({ny} counts) not equal to length of x ({nx} counts).")
        if ((ndy!=ny) and (ndy>1)):
            raise Exception(f"Length of dy ({ndy} counts) not equal to length of y ({ny} counts).")

        # sort the vectors with x increasing
        inds = np.argsort(x)
        x = np.array([x[i] for i in inds])
        y = np.array([y[i] for i in inds])
        if (ndy==ny):
            dy = np.array([dy[i] for i in inds])
        inds = np.argsort(xx)
        xx = np.array([xx[i] for i in inds])
        self._xx = xx

        # estimation of std of the y vector if needed
        if (ndy==1) and (dy[0]>0):
            std = dy[0]
        else:
            d = np.zeros(len(y)-1, dtype=float)
            for i in range(len(y)-1):
                d[i] = y[i+1]-y[i]
            # empirical estimation of sigma
            std = np.std(d)/math.sqrt(2.0)
        if (std==0):
            std = 1

        # verify values of the dy vector if needed
        if (ndy==1):
            dy=np.full(ny,std)
        elif (ndy==ny):
            for i in range(ny):
                if (dy[i]<=0):
                    dy[i]=std
        else:
            dy = np.full(dy,std)

        # x vector must not have same absissa
        sames=np.full(nx,0)
        for i in range(nx-1):
            if (x[i]==x[i+1]):
                if (sames[i]==0):
                    sames[i]=i
                sames[i+1]=sames[i]
        i=0
        ii=0
        xx0=[]
        yy0=[]
        dyy0=[]
        while (i<nx):
            i1=-1
            i2=i1
            if (sames[i]>0):
                i1=i
                iii=i+1
                while (iii<nx):
                    if (sames[iii]==0):
                        i2=iii-1
                        break
                    iii+=1
                mean = 0
                dy2 = 0
                #print(f"(1.7) i1={i1} i2={i2}")
                for iii in range (i1,i2+1):
                    mean += y[iii]
                    dy2 += dy[iii]*dy[iii]
                    #print(f"(1.8) iii={iii}")
                mean = mean/(i2-i1+1)
                dy2 = math.sqrt(dy2)/(i2-i1+1)
            xx0.append(x[i])
            #print(f"(1.80) i={i}")
            if (i1>=0):
                yy0.append(mean)
                dyy0.append(dy2)
                i+=(i2-i1)
            else:
                yy0.append(y[i])
                dyy0.append(dy[i])
            i+=1
            ii+=1
        self._x=np.array(xx0)
        self._y=np.array(yy0)
        self._dy=np.array(dyy0)
        self._xx=np.array(xx)

        # shift all indexes by an offset +1
        x = np.array(0.0)
        y = np.array(0.0)
        dy = np.array(0.0)
        xx = np.array(0.0)
        x = np.append(x,np.array(self._x))
        y = np.append(y,np.array(self._y))
        dy = np.append(dy,np.array(self._dy))
        xx = np.append(xx,np.array(self._xx))

        # create the output vector (filled by 0)
        nn=len(xx)-1
        ff=np.zeros(nn+1, dtype=float)

        # start the algorithm
        n1=1
        n2=len(x)-1; # -1 !
        n=(n2+1)-(n1-1)+2
        r=np.zeros(n, dtype=float)
        r1=np.zeros(n, dtype=float)
        r2=np.zeros(n, dtype=float)
        t=np.zeros(n, dtype=float)
        t1=np.zeros(n, dtype=float)
        u=np.zeros(n, dtype=float)
        v=np.zeros(n, dtype=float)
        a=np.zeros(n, dtype=float)
        b=np.zeros(n, dtype=float)
        c=np.zeros(n, dtype=float)
        d=np.zeros(n, dtype=float)

        m1=n1-1
        m2=n2+1
        r[m1]=0
        r[n1]=0
        r1[n2]=0
        r2[n2]=0
        r2[m2]=0
        u[m1]=0
        u[n1]=0
        u[n2]=0
        u[m2]=0
        p=0
        m1=n1+1
        m2=n2-1
        h=x[m1]-x[n1]
        f=(y[m1]-y[n1])/h
        for i in range(m1,m2+1):
            g=h
            h=x[i+1]-x[i]
            e=f
            f=(y[i+1]-y[i])/h
            a[i]=f-e
            t[i]=2*(g+h)/3
            t1[i]=h/3
            r2[i]=dy[i-1]/g
            r[i]=dy[i+1]/h
            r1[i]=-dy[i]/g-dy[i]/h
        for i in range(m1,m2+1):
            b[i]=r[i]*r[i]+r1[i]*r1[i]+r2[i]*r2[i]
            c[i]=r[i]*r1[i+1]+r1[i]*r2[i+1]
            d[i]=r[i]*r2[i+2]
        f2=-s
        while True:
            #:next_interation
            for i in range(m1,m2+1):
                r1[i-1]=f*r[i-1]
                r2[i-2]=g*r[i-2]
                r[i]=1/(p*b[i]+t[i]-f*r1[i-1]-g*r2[i-2])
                u[i]=a[i]-r1[i-1]*u[i-1]-r2[i-2]*u[i-2]
                f=p*c[i]+t1[i]-h*r1[i-1]
                g=h
                h=d[i]*p
            for i in range(m2,m1-1,-1):
                u[i]=r[i]*u[i]-r1[i]*u[i+1]-r2[i]*u[i+2]
            e=0
            h=0
            for i in range(n1,m2+1):
                g=h
                h=(u[i+1]-u[i])/(x[i+1]-x[i])
                v[i]=(h-g)*dy[i]*dy[i]
                e=e+v[i]*(h-g)
            g=-h*dy[n2]*dy[n2]
            v[n2]=g
            e=e-g*h
            g=f2
            f2=e*p*p
            if ((f2>=s) or (f2<=g)):
                break
            f=0
            h=(v[m1]-v[n1])/(x[m1]-x[n1])
            for i in range(m1,m2+1):
                g=h
                h=(v[i+1]-v[i])/(x[i+1]-x[i])
                g=h-g-r1[i-1]*r[i-1]-r2[i-2]*r[i-2]
                f=f+g*r[i]*g
                r[i]=g
            h=e-p*f
            if (h<=0):
                break
            p=p+(s-f2)/((math.sqrt(s/e)+p)*h)
            # goto next_iteration;
        # use negative branch of square root, if the sequence of absissae x[i] is strictly decreasing
        for i in range(n1,n2+1):
            a[i]=y[i]-p*v[i]
            c[i]=u[i]
        for i in range(n1,m2+1):
            h=x[i+1]-x[i]
            d[i]=(c[i+1]-c[i])/(3*h)
            b[i]=(a[i+1]-a[i])/h-(h*d[i]+c[i])*h
        # --- compute the final vector
        for ii in range(1,nn+1):
            ff[ii]=0
            for i in range(n1,n2):
                if ((xx[ii]>=x[i]) and (xx[ii]<=x[i+1])):
                    h=xx[ii]-x[i]
                    ff[ii]=((d[i]*h+c[i])*h+b[i])*h+a[i]
                    break

        # --- shift indexes of ff
        yy = ff[1:]
        self._yy = yy

        return yy

    def fitmag(self, s:float, tt:np.ndarray=None):
        if (self._t is None) or (self._mag is None) or (self._dmag is None):
            raise Exception("No input data. Use methods import_data or t, mag, dmag")
        if (self._tt is None):
            t=self._tt_default
        else:
            t=self._tt
        tt = []
        # --- cancel absissa outside x range
        for i in range(len(t)):
            if t[i]<self._t[0]:
                continue
            elif t[i]>self._t[-1]:
                continue
            tt.append(t[i])
        # --- nomalization of s
        s *= len(self._t)
        self.fit(self._t,self._mag,self._dmag,s,tt)
        self._t = self._x
        self._mag = self._y
        self._dmag = self._dy
        self._tt = self._xx
        self._magmag = self._yy
        self._clear_xydyxxyy()
        return (self._tt, self._magmag)

    def fitmag2mjy(self, s:float, mag2mjy:float, tt:np.ndarray=None):
        self.fitmag(s,tt)
        tt = self.tt
        fmax = mag2mjy*1e3*np.power(10,-0.4*(self.mag-self.dmag));  # mJy
        fmin = mag2mjy*1e3*np.power(10,-0.4*(self.mag+self.dmag));  # mJy
        self._f = (fmin+fmax)/2
        self._df = fmax - self._f
        self._ff = mag2mjy*1e3*np.power(10,-0.4*self.magmag);  # mJy
        return (self._tt, self._ff)

    def plot(self):
        mode = 0
        if self._f is not None:
            mode = 3
        elif self._mag is not None:
            mode = 2
        elif self._y is not None:
            mode = 1
        # ---
        if mode==1:
            x = self._x
            y = self._y
            dy = self._dy
            xx = self._xx
            yy = self._yy
            plt.figure(1)
            fig, ax=plt.subplots(1,1,figsize=(5.0,5.0))
            plt.plot(x,y,'ob');
            if isinstance(dy,np.ndarray)==True:
                if isinstance(dy.tolist(),(int,float)):
                    dy=np.array([dy],dtype=float)
                ndy=len(dy)
                for i in range(ndy):
                    vx = np.array([x[i], x[i]])
                    vy = np.array([y[i]+dy[i], y[i]-dy[i]])
                    plt.plot(vx,vy,'-b')
            if yy is not None:
                plt.plot(xx,yy,'r-');
            plt.xlabel('x')
            plt.ylabel('y')
            plt.grid(True)
            plt.show()
            # plt.savefig('/Users/rosa/Desktop/Utiles/Prueba.png')
            plt.close()
        elif mode==2:
            t = self._t
            mag = self._mag
            dmag = self._dmag
            tt = self._tt
            magmag = self._magmag
            plt.figure(1)
            fig, ax=plt.subplots(1,1,figsize=(5.0,5.0))
            plt.semilogx(t,mag,'ob');
            plt.gca().invert_yaxis()
            if isinstance(dmag,np.ndarray)==True:
                if isinstance(dmag.tolist(),(int,float)):
                    dmag=np.array([dmag],dtype=float)
                ndmag=len(dmag)
                for i in range(ndmag):
                    vx = np.array([t[i], t[i]])
                    vy = np.array([mag[i]+dmag[i], mag[i]-dmag[i]])
                    plt.semilogx(vx,vy,'-b')
            if tt is not None:
                plt.semilogx(tt,magmag,'r-');
            plt.xlabel('Time since trigger')
            plt.ylabel('Magnitude')
            plt.grid(True)
            plt.show()
            plt.close()
        elif mode==3:
            t = self._t
            f = self._f
            df = self._df
            tt = self._tt
            ff = self._ff
            plt.figure(1)
            fig, ax=plt.subplots(1,1,figsize=(5.0,5.0))
            plt.loglog(t,f,'ob');
            if isinstance(df,np.ndarray)==True:
                if isinstance(df.tolist(),(int,float)):
                    df=np.array([df],dtype=float)
                ndf=len(df)
                for i in range(ndf):
                    vx = np.array([t[i], t[i]])
                    vy = np.array([f[i]+df[i], f[i]-df[i]])
                    plt.semilogx(vx,vy,'-b')
            if tt is not None:
                plt.semilogx(tt,ff,'r-');
            plt.xlabel('Time since trigger')
            plt.ylabel('Flux density')
            plt.grid(True)
            plt.show()
            plt.close()

    def import_data(self,full_filename:str, col_t1:int, col_t2:int, col_mag:int, col_dmag:int):
        return self._load_data(full_filename,col_t1,col_t2,col_mag,col_dmag)

    def clear(self):
        self._clear_xydyxxyy()
        self._clear_tmagdmagttmagmag()

# =====================================================================
# =====================================================================
# Special methods
# =====================================================================
# =====================================================================

    def __init__(self):
        self._clear_xydyxxyy()
        self._clear_tmagdmagttmagmag()
        self._tt_default = np.logspace(0,6,61)

# =====================================================================
# =====================================================================
# Test if main
# =====================================================================
# =====================================================================

if __name__ == "__main__":

    example = 3
    print("Example       = {}".format(example))

    if example == 1:
        """
        case where we define (time,mag) inputs by hands
        """
        fiter = Splinefit();
        fiter.clear()
        fiter.t   = [1e2, 3e2, 1e3, 5e3, 2e4]
        fiter.mag = [14.5, 15.0, 15.8, 16.7, 18.3]
        fiter.dmag = 0 # we put 0 when we have no error available for data
        s = 0.5 # smoothing factor
        # --- fit mag -> mag
        fiter.fitmag(s)
        fiter.plot()
        # --- fit mag -> mJy
        mag2mjy = 3500 ; # F(mJy) for a zero magnitude
        fiter.fitmag2mjy(s,mag2mjy)
        fiter.plot()

    if example == 2:
        """
        Case here we read a text file of (time,mag)
        """
        fiter = Splinefit();
        fiter.clear()
        grb_text_file = os.getcwd() + "/grb180418a_r_ratir.txt"
        # col 0 = t1
        # col 1 = t2
        # col 3 = mag
        # col 4 = dmag
        fiter.import_data(grb_text_file,0,1,3,4)
        s = 1 # smoothing factor
        # --- fit mag -> mag
        fiter.fitmag(s)
        fiter.plot()
        # --- fit mag -> mJy
        mag2mjy = 3500 ; # F(mJy) for a zero magnitude
        fiter.fitmag2mjy(s,mag2mjy)
        fiter.plot()

    if example == 3:
        """
        Case of an example of (x,y) inputs as sinus
        """
        fiter = Splinefit();
        fiter.clear()
        # --- Define an example x,y
        x=np.linspace(0,20,100)
        y1=np.sin(x/1+1)
        y2=[]
        for i in range(len(x)):
            y2.append(0.2*(rd.random()-0.5))
        y=y1+y2
        # --- Define dy as a same value for all data
        # =0 means the uncertainties are estimated by the fiter itself
        dy = 0
        # --- Define a smooth factor
        s=0.01*len(x)
        # --- Define a resampling
        xx = np.linspace(x[0],x[-1],100)
        # --- Use the splinefit
        fiter.fit(x,y,dy,s,xx)
        fiter.plot()
        # --- print output data
        xx = fiter.xx
        yy = fiter.yy
        print(f"xx={xx}")
        print(f"yy={yy}")

    if example == 4:
        """
        """
        fiter = Splinefit();
        t  = [1e2, 3e2, 1e3, 5e3, 2e4]
        mag = [14.5, 15.0, 15.8, 16.7, 18.3]
        dmag = 0
        s = 1*len(t)
        tfit = np.linspace(t[0],t[-1],100)
        magfit = fiter.fit(t, mag, dmag, s, tfit)
        # you can now plot(tfit,magfit)