import numpy as np
import sys
import argparse
import matplotlib.pyplot as plt
#import seaborn as sns
from matplotlib.patches import Ellipse
import matplotlib.transforms as transforms


data_dir = sys.argv[1]
data_1 = data_dir + '/0/samples.npz'
data_2 = data_dir + '/1/samples.npz'


def next_pow_two(n):
    i = 1
    while i < n:
        i = i << 1
    return i

def autocorr_1d(x, norm=True):
    x = np.atleast_1d(x)
    if len(x.shape) != 1:
        raise ValueError("invalid dimensions for 1D autocorrelation function")
    n = next_pow_two(len(x))
    f = np.fft.fft(x - np.mean(x), n=2*n)
    acf = np.fft.ifft(f * np.conjugate(f))[:len(x)].real / (4 * n)
    if norm:
        acf /= acf[0]
    return acf


def window(taus, c):
    m = np.arange(len(taus)) < c * taus
    if np.any(m):
        return np.argmin(m)
    return len(taus) - 1

def acf(x, c=5.0):
    print("len(x) = ", len(x))
    f = np.zeros(x.shape[1])
    for xx in x:
        f += autocorr_1d(xx)
    f /= len(x)
    taus = 2.0 * np.cumsum(f) - 1
    return taus[window(taus, c)]

def acf_ml(x, thin=1, c=5.0):
    import celerite
    from celerite import terms
    from scipy.optimize import minimize
    #initial estimate
    init = acf(x, c=c)
    z = x[:, ::thin]
    N = z.shape[1]
    #build a GP to model the chain
    tau = max(1.0, init / thin)
    kernel = terms.RealTerm(np.log(0.9 * np.var(z)), -np.log(tau),
                            bounds=[(-5.,5.), (-np.log(N), 0.)])
    kernel += terms.RealTerm(np.log(0.1 * np.var(z)), -np.log(.5*tau),
                             bounds=[(-5.,5.), (-np.log(N), 0.)])
    gp = celerite.GP(kernel, mean=np.mean(z))
    gp.compute(np.arange(z.shape[1]))

    def objective(theta):
        gp.set_parameter_vector(theta)

        v, g = zip(*(gp.grad_log_likelihood(z0, quiet=True) for z0 in z))
        return -np.sum(v), -np.sum(g, axis=0)

    theta0 = gp.get_parameter_vector()
    bounds = gp.get_parameter_bounds()
    sol = minimize(objective, theta0, jac=True, bounds=bounds)
    gp.set_parameter_vector(sol.x)

    a, c = kernel.coefficients[:2]
    return thin * 2 * np.sum(a / c) / np.sum(a)


def get_whole_chain(data):
    a = data['a']
    p = data['p']
    n = len(p)
    chain = np.zeros((n,2))
    chain[:,0] = a
    chain[:,1] = p
    print("chain is", chain)
    return chain


def ess(data):
    chain = get_whole_chain(data).T
    N = chain.shape[1]
    return N / acf(chain)

def datafile_name(data_dir, i):
    return data_dir + f"/{i}/samples.npz"



def print_mu_sd(Ne, x, name):
    print(f"Mean of {name}: {np.mean(x)}. Standard Error: {np.std(x) / np.sqrt(Ne)}. N = {x.size}. Ne = {Ne}")

def analyze_data(dr, **kwargs):
    data = np.load(dr)
    variables = data.files
    #interval, variables = variables[0], variables[1:]
    Ne = ess(data)
    for v in variables:
        print_mu_sd(Ne, data[v],v)
        plt.figure()
        plt.title(v)
        plt.hist(data[v])

    plt.show()



def plot_p_stokes_trace(dr, args):
    ptrue = args.ptrue
    atrue = args.atrue
    
    data = np.load(dr)
    p = data['p']
    a = data['a']
    #from https://matplotlib.org/3.1.1/gallery/statistics/confidence_ellipse.html#sphx-glr-gallery-statistics-confidence-ellipse-py
    cov = np.cov(p,a)
    pearson = cov[0,1]/np.sqrt(cov[0,0] * cov[1,1])
    rx, ry = np.sqrt(1 + pearson), np.sqrt(1 - pearson)
    conf_ellipse = Ellipse((0,0),width=2*rx,height=2*ry,
                           facecolor='none',edgecolor='green')
    sx = np.sqrt(cov[0,0]) * args.conf_std
    sy = np.sqrt(cov[1,1]) * args.conf_std
    trans = transforms.Affine2D().rotate_deg(45).scale(sx,sy).translate(np.mean(p),np.mean(a))
    fig, ax = plt.subplots()
    ax.set_title(f'p vs a (true value red, sample mean +/- {args.conf_std} std blue/green)')
    plt.xlabel('p')
    plt.ylabel('a')
    ax.scatter([ptrue],[atrue],c='r')
    ax.scatter(np.mean(p),np.mean(a),c='b')
    ax.scatter(p,a,c='k')
    conf_ellipse.set_transform(trans + ax.transData)
    ax.add_patch(conf_ellipse)
    plt.savefig(f"{dr}_trace.png")

        
def analyze_chains(chains, args, **kwargs):
    for i, ch in enumerate(chains):
        print(f"""--------------------------------------------
        Chain {i} variable-wise statistics:
        """)
        analyze_data(ch, **kwargs)
        plot_p_stokes_trace(ch, args)



if __name__ == '__main__':
    parser = argparse.ArgumentParser('Analyze the saved output of a Markov chain generated by PyMC3')
    parser.add_argument('data_dir',help='Directory in which the chains are stored')
    parser.add_argument('-n','--num_chains',type=int,help='How many independent chains are stored in the directory?',default=2)

    parser.add_argument('-ptrue','--ptrue',help='True value of p',type=float,default=1.2)
    parser.add_argument('-atrue','--atrue',help='True value of a',type=float,default=1.0)
    
    parser.add_argument('-std','--conf_std',help='How many standard deviations for the confidence interval? Default 2.0',type=float,default=2.0)
    args = parser.parse_args()

    data_dir = args.data_dir[:-1] if args.data_dir[-1] == '/' else args.data_dir
    n = args.num_chains
    
    chains = [datafile_name(data_dir,i) for i in range(n)]

    analyze_chains(chains, args)
