import numpy as np
import argparse

# this function calculates the S-matrix for the given set of t- and v-values
def CalculateSMatrix(t, v):
    tv = np.array([[vk*ti*1j for vk in v] for ti in t])
    return np.exp(-tv)

# calculate the weight matrix from the errors of the sM(s)
def CalculateWMatrix(ey):
    try:
        wvec = 1.0/ey**2
        return np.diag(wvec)
    except:
        raise RuntimeError("Could not calculate the weights, check if the errors are > 0")

# this calculates the StWS matrix
def CalulateSWSMatrix(S,W):
    return np.matmul(np.conjugate(S.T),np.matmul(W,S))

# this calculates the effective Sigma matrix from SWS and lambda value (LVal)
def CalcSigmaLambda(SWS,reg):
    M = SWS
    M += reg * np.eye(SWS.shape[0])
    try:
        res = np.linalg.inv(M)
        return res
    except:
        raise RuntimeError("Failed to get effective covariance matrix. Play with rmin/rmax/NumPts/lambda values")

# save Pearson's coefficient correlation matrix
def SaveCorrMatrix(Sigma):
    N = len(Sigma)
    Res = np.zeros(Sigma.shape)
    for i in range(0,N):
        for j in range(0,N):
            Res[i][j] = Sigma[i][j]/np.sqrt(Sigma[i][i] * Sigma[j][j])
    np.savetxt("corr.mt", Res)            
    
    MaxOffDiag = np.max(Res-np.diag(np.diag(Res)))
    MinOffDiag = np.min(Res-np.diag(np.diag(Res)))
    print("Min/max off-diagonal correlations are %f and %f" % (MinOffDiag,MaxOffDiag))

# this function returns rwsLSSA-d values
def GetXValues(Sigma,S,W,y):
    x = np.matmul(np.matmul(Sigma.T,np.matmul(np.conjugate(S.T),W)),y)
    return x
    
# this function returns uncertanites rwsLSSA-d values
def GetXErrors(Sigma,SWS,reg):
    res = np.sqrt(np.diag(Sigma)) # this is pure error from Sigma
    # and, now we extrapolate the experimental errors
    res *= np.sqrt(np.diag(SWS) + reg)
    res /= np.sqrt(np.diag(SWS))
    return res

# this function unextrapolated uncertanites rwsLSSA-d values
def GetXErrorsUncorr(Sigma,SWS,reg):
    res = np.sqrt(np.diag(Sigma)) # this is pure error from Sigma
    return res

# this function prints values of different functionals, for the L-curve method
def PrintPureFunctional(S,W,x,y,reg):
    dQ = np.matmul(S,x) - y
    Fpure = 0.5 * np.matmul(dQ,np.matmul(W,dQ))
    Freg  = 0.5 * np.dot(x,x)
    
    print("%.7e %.7e %.7e (L-curve: reg.param., pure LSQ, reg. LSQ)\n" % (reg,Fpure,Freg))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='This script performs the regularized weighted least-squares spectral analysis (rwsLSSA) for arbitrary time-dependent data')

    parser.add_argument('-f', '--file', type=str, default=None,
                    help="File name of the file containing experimental observables. Default is None (no file).") 

    parser.add_argument('--vmin', type=float, default=None,
                    help="Minimal frequency value of the spectrum, in THz. Default is None (automatic determination).") 

    parser.add_argument('--vmax', type=float, default=None,
                    help="Maximal frequency value of the spectrum, in THz. Default is None (automatic determination).") 

    parser.add_argument('-n', '--NumPts', type=int, default=None,
                    help="Number of points for spectrum. Default is 10.") 

    parser.add_argument('--damp', type=float, default=None,
                    help="Damping parameter value (in THz**2). Default is None (not applied).") 
                    
    parser.add_argument('--reg', type=float, default=0.0,
                    help="L2-regularization parameter. Default is 0.0 (no regularization applied).")                     

    parser.add_argument('--AutoReg', action='store_true', 
                    help="Flag to apply the automatic a priori regularization parameter choice.")

    parser.add_argument('--PrintOmega', action='store_true', 
                    help="Flag to plot frequencies as angular frequencies w=2*pi*v, but not as normal frequencies.")

    args,unknown = parser.parse_known_args() # parse the arguments
    
    if args.file is None:
        print("No file, no rwLSSA analysis")
        exit()

    # load data
    # first row is the t (in ps),
    # second row is the signal (in arb. units)
    # third row is the error of signal (the same units, as signal).
    data = np.loadtxt(args.file)
    
    if args.damp is None:
        t = data[:,0]
        y = data[:,1]
        ey = data[:,2]
    else:
        t = data[:,0]
        gamma = args.damp
        y = data[:,1] * np.exp(-gamma*t**2)
        ey = data[:,2] * np.exp(-gamma*t**2)
    
    inds = np.argsort(t)
    t = t[inds]
    y = y[inds]
    ey = ey[inds]
    
    
    print("Loaded data from file "+args.file)
    
    dt=np.average(t[1:]-t[:-1]) # mean step
    vmax = 2.*np.pi/dt    
    ttotal=(max(t)-min(t)) # time span
    dvmin = 2.*np.pi/ttotal


    print("Mean time step is %.5f ps, time span is %.5f ps" % (dt,ttotal) )
    print("Minimal dv from estimation 2*pi/(max(t)-min(t)) is %.5f THz" % (dvmin) )
    print("Maximal value of v from estimation 2*pi/dt is %.5f THz" % (vmax) )
    
    if args.vmin is None:
        args.vmin = 0.0
        
    if args.vmax is None:
        args.vmax = vmax
    
    if args.NumPts is None:
        args.NumPts = int( (args.vmax-args.vmin)/dvmin ) - 1
    
    # set the v-values according to pre-sets
    v = np.linspace(args.vmin,args.vmax,args.NumPts)
    
    print("Current setting of the dv is %.5f THz" % (np.mean(v[1:]-v[:-1])))
    
    S = CalculateSMatrix(t, v)
    W = CalculateWMatrix(ey)
    SWS = CalulateSWSMatrix(S,W)
    #print(S.shape,W.shape,SWS.shape)

    # choose automatic regularization criterion
    if args.AutoReg:
        mdSWS = np.mean(np.diag(SWS))
        mY2 = np.dot(y,np.dot(W,y))/(len(v)*np.sum(np.diag(W)))
        args.reg = np.real(mdSWS/(1. + mY2 * mdSWS))
        print("Automatically chosen regularization parameter is %f" % (args.reg))
    
    Sigma = CalcSigmaLambda(SWS,args.reg)
    
    # calculate the spectrum
    x  = GetXValues(Sigma,S,W,y)
    # calculate the RDF uncertainties
    ex = GetXErrors(Sigma,SWS,args.reg)
    uex = GetXErrorsUncorr(Sigma,SWS,args.reg)
    # now, save the resuling RDF function
    MaxVal = max(x)
    
    if not args.PrintOmega:
        v = v/(2.*np.pi)
        
    np.savetxt("spectrum.res", np.stack([v,np.absolute(x),np.real(ex),np.real(uex),np.angle(x)],axis=-1),header=' max.value = %15.10f' % (MaxVal))

    PrintPureFunctional(S,W,x,y,args.reg)
    SaveCorrMatrix(Sigma)    
        
    exit()
