import time
import numpy as np
import scipy.linalg as lin
import scipy.stats as stat
import matplotlib.pyplot as plt
import kernel_qp as qp

def log_likelihood(theta, t, Y):
    tLen = len(t.reshape(-1))
    Nw   = theta.shape[0]
    Y    = np.ascontiguousarray(Y.reshape(-1))

    K       = qp.gamma_qp(t, theta)
    Kinv_mu = np.empty((Nw, tLen))

    for i in range(Nw):
        K[i], lower = lin.cho_factor(K[i])
        Kinv_mu[i]  = lin.cho_solve((K[i], lower), Y)
    muT_Kinv_mu = (Y*Kinv_mu).sum(axis=1)

    HalfLogDetK = np.log(K.diagonal(axis1=1, axis2=2)).sum(axis=1)

    return -0.5*tLen*np.log(2*np.pi) - HalfLogDetK - 0.5*muT_Kinv_mu

def log_prior(theta):
    out  = stat.norm.logpdf(theta[:,0],    loc=0.0005, scale=0.00025) # Sc
    out += stat.norm.logpdf(theta[:,1],    loc=25.3,   scale=2.0)     # Pqp
    out += stat.uniform.logpdf(theta[:,2], loc=0.1,    scale=1.5)     # lambda_P
    out += stat.norm.logpdf(theta[:,3],    loc=50.0,   scale=25.0)    # lambda_e
    out += stat.norm.logpdf(theta[:,4],    loc=0.0005, scale=0.00025) # sig
    return out.squeeze()

def log_posterior(theta, t, Y):
    return log_likelihood(theta, t, Y) + log_prior(theta)

def proposal(size=1, a=2.0):
    U = stat.uniform.rvs(size=size)
    return (U*(a - 1) + 1)**2/a

def main():
    Nparam = 5;
    Ni  = int(1e4)
    Nw  = int(1e2)
    Nw2 = int(Nw/2)
    Ns  = int(Ni*Nw)

    Z       = proposal(size=(Ni, Nw2, 2))
    r       = stat.uniform.rvs(size=(Ni, Nw2, 2))
    logrZNp = np.log(r*(Z**(1 - Nparam)))
    rw      = stat.randint.rvs(low=0, high=Nw2, size=(Ni, Nw2, 2))
    
    x01 = stat.norm.rvs(loc=0.0005, scale=0.00025, size=(Nw2, 2)) # Sc
    x02 = stat.norm.rvs(loc=25.3, scale=2.0,  size=(Nw2, 2))      # Pqp
    x03 = stat.uniform.rvs(loc=0.1, scale=1.5,  size=(Nw2, 2))    # lambda_P
    x04 = stat.norm.rvs(loc=50.0, scale=25.0, size=(Nw2, 2))      # lambda_e
    x05 = stat.norm.rvs(loc=0.0005, scale=0.00025, size=(Nw2, 2)) # sig2
    
    theta0 = np.array([x01, x02, x03, x04, x05])
    theta  = np.zeros((Nparam, Ni, Nw2, 2))
    acpt   = np.zeros(Z.shape)
    
    fname = r'C:\Users\Nicholas\Documents\Walsworth_group\solar_RVs\solarSindexCutDaily.txt'
    data  = np.loadtxt(fname, skiprows=1)
    t     = data[:,0]
    Y     = data[:,2]
    
    for l in range(0, Ni):
        if l % 1 == 0: print(100*l/Ni, '% complete')
        
        j                 = rw[l,:,0]
        thetaP            = theta0[:,j,1] + Z[l,:,0]*(theta0[:,:,0] - theta0[:,j,1])
        logPthetaP        = log_posterior(thetaP.T, t, Y)
        logPtheta0        = log_posterior(theta0[:,:,0].T, t, Y)
        idx               = (logPthetaP - logPtheta0) >= logrZNp[l,:,0]
        acpt[l,idx,0]    += 1
        theta[:,l, idx,0] = thetaP[:,idx]
        theta[:,l,~idx,0] = theta0[:,~idx,0]
        theta0[:,:,0]     = theta[:,l,:,0]
        
        j                 = rw[l,:,1]
        thetaP            = theta0[:,j,0] + Z[l,:,1]*(theta0[:,:,1] - theta0[:,j,0])
        logPthetaP        = log_posterior(thetaP.T, t, Y)
        logPtheta0        = log_posterior(theta0[:,:,1].T, t, Y)
        idx               = (logPthetaP - logPtheta0) >= logrZNp[l,:,1]
        acpt[l,idx,1]    += 1
        theta[:,l, idx,1] = thetaP[:,idx]
        theta[:,l,~idx,1] = theta0[:,~idx,1]
        theta0[:,:,1]     = theta[:,l,:,1]

    fname = r'C:\Users\Nicholas\Documents\Walsworth_group\solar_RVs\S_samples20'
    np.save(fname, theta)
    fname = r'C:\Users\Nicholas\Documents\Walsworth_group\solar_RVs\accept20'
    np.save(fname, acpt)

if __name__ == '__main__': main()