import numpy as np
from numpy import exp, sin
import matplotlib.pyplot as plt
import scipy.integrate as spint
from scipy.optimize import fsolve, fmin, minimize, root, newton
from scipy.linalg import solve_triangular
from scipy.sparse.linalg import eigsh
import scipy
from scipy.integrate import solve_ivp
from param import nm, U0, dtsymp
from kernels import *


def f_kern(x, y, x0, y0, l):
    return kern_num(x,y,x0,y0,l[0], l[1])

def dkdx(x, y, x0, y0, l):
    return dkdx_num(x,y,x0,y0,l[0], l[1])

def dkdy(x, y, x0, y0, l):
    return dkdy_num(x,y,x0,y0,l[0], l[1])

def dkdx0(x, y, x0, y0, l):
    return dkdx0_num(x,y,x0,y0,l[0], l[1])

def dkdy0(x, y, x0, y0, l):
    return dkdy0_num(x,y,x0,y0,l[0], l[1])

def d2kdxdx0(x, y, x0, y0, l):
    return d2kdxdx0_num(x,y,x0,y0,l[0], l[1])

def d2kdydy0(x, y, x0, y0, l):
    return d2kdydy0_num(x,y,x0,y0,l[0], l[1])

def d2kdxdy0(x, y, x0, y0, l):
    return d2kdxdy0_num(x,y,x0,y0,l[0], l[1])

def d2kdydx0(x, y, x0, y0, l):
    return d2kdxdy0(x, y, x0, y0, l)

def d3kdxdx0dy0(x, y, x0, y0, l):
    return d3kdxdx0dy0_num(x,y,x0,y0,l[0], l[1])

def d3kdydy0dy0(x, y, x0, y0, l):
    return d3kdydy0dy0_num(x,y,x0,y0,l[0], l[1])

def d3kdxdy0dy0(x, y, x0, y0, l):
    return d3kdxdy0dy0_num(x,y,x0,y0,l[0], l[1])

def d3kdydx0dy0(x, y, x0, y0, l):
    return d3kdxdy0dy0(x, y, x0, y0, l)

def dkdlx(x, y, x0, y0, l):
    return dkdlx_num(x, y, x0, y0, l[0], l[1])

def dkdly(x, y, x0, y0, l):
    return dkdly_num(x, y, x0, y0, l[0], l[1])

def d3kdxdx0dlx(x, y, x0, y0, l):
    return d3kdxdx0dlx_num(x,y,x0,y0,l[0], l[1])

def d3kdydy0dlx(x, y, x0, y0, l):
    return d3kdydy0dlx_num(x,y,x0,y0,l[0], l[1])

def d3kdxdy0dlx(x, y, x0, y0, l):
    return d3kdxdy0dlx_num(x,y,x0,y0,l[0], l[1])

def d3kdydx0dlx(x, y, x0, y0, l):
    return d3kdxdy0dlx(x, y, x0, y0, l)

def d3kdxdx0dly(x, y, x0, y0, l):
    return d3kdxdx0dly_num(x,y,x0,y0,l[0], l[1])

def d3kdydy0dly(x, y, x0, y0, l):
    return d3kdydy0dly_num(x,y,x0,y0,l[0], l[1])

def d3kdxdy0dly(x, y, x0, y0, l):
    return d3kdxdy0dly_num(x,y,x0,y0,l[0], l[1])

def d3kdydx0dly(x, y, x0, y0, l):
    return d3kdxdy0dly(x, y, x0, y0, l)

def intode(y, t, dtsymp):
    ys = np.zeros([2, len(t)])
    ys[0,0] = y[0]
    ys[1,0] = y[1]

    for kt in range(1,len(t)):
            ys[1, kt]=ys[1, kt-1] - dtsymp*(np.sin(ys[0, kt-1] + np.pi))#-0.1*np.sin(ys[0,kt-1]-0.5*t[kt]))
            ys[0, kt]=ys[0, kt-1] + dtsymp*ys[1, kt]

    # ys[0,:] = np.mod(ys[0,:], 2*np.pi)
    return ys.T

def build_K(xin, x0in, l, K):
    # set up covariance matrix
    N = K.shape[1]//2
    N0 = K.shape[0]//2
    x0 = x0in[0:N]
    x = xin[0:N]
    y0 = x0in[N:2*N]
    y = xin[N:2*N]
    for k in range(N0):
        for lk in range(N):
            K[k,lk] = d2kdxdx0(
                x0[k], y0[k], x[lk], y[lk], l)
            K[N0+k,lk] = d2kdxdy0(
                 x0[k], y0[k], x[lk], y[lk], l)
            K[k,N+lk] = d2kdydx0(
                 x0[k], y0[k], x[lk], y[lk], l)
            K[N0+k,N+lk] = d2kdydy0(
                x0[k], y0[k], x[lk], y[lk], l)

def buildK(x, y, x0, y0, l, K):
    # set up covariance matrix
    N = K.shape[1]//2
    N0 = K.shape[0]//2

    for k in range(N0):
        for lk in range(N):
            K[k,lk] = d2kdxdx0(
                x0[k], y0[k], x[lk], y[lk], l)
            K[N0+k,lk] = d2kdxdy0(
                 x0[k], y0[k], x[lk], y[lk], l)
            K[k,N+lk] = d2kdydx0(
                 x0[k], y0[k], x[lk], y[lk], l)
            K[N0+k,N+lk] = d2kdydy0(
                x0[k], y0[k], x[lk], y[lk], l)


def build_dK(x, y, x0, y0, l):
    # set up covariance matrix
    N = len(x)
    N0 = len(x0)
    k11 = np.empty((N0, N))
    k12 = np.empty((N0, N))
    k21 = np.empty((N0, N))
    k22 = np.empty((N0, N))

    dK = []

    for k in range(N0):
        for lk in range(N):
              k11[k,lk] = d3kdxdx0dlx(
                  x0[k], y0[k], x[lk], y[lk], l)
              k21[k,lk] = d3kdxdy0dlx(
                  x0[k], y0[k], x[lk], y[lk], l)
              k12[k,lk] = d3kdydx0dlx(
                  x0[k], y0[k], x[lk], y[lk], l)
              k22[k,lk] = d3kdydy0dlx(
                  x0[k], y0[k], x[lk], y[lk], l)

    dK.append(np.vstack([
        np.hstack([k11, k12]),
        np.hstack([k21, k22])
    ]))

    for k in range(N0):
        for lk in range(N):
             k11[k,lk] = d3kdxdx0dly(
                 x0[k], y0[k], x[lk], y[lk], l)
             k21[k,lk] = d3kdxdy0dly(
                 x0[k], y0[k], x[lk], y[lk], l)
             k12[k,lk] = d3kdydx0dly(
                  x0[k], y0[k], x[lk], y[lk], l)
             k22[k,lk] = d3kdydy0dly(
                 x0[k], y0[k], x[lk], y[lk], l)

    dK.append(np.vstack([
        np.hstack([k11, k12]),
        np.hstack([k21, k22])
    ]))

    return dK


def build_dKdy0(x, y, x0, y0,l):
    # set up covariance matrix
    N = len(x)
    N0 = len(x0)
    k11 = np.empty((N0, N))
    k12 = np.empty((N0, N))
    k21 = np.empty((N0, N))
    k22 = np.empty((N0, N))

    for k in range(N0):
        for lk in range(N):
             k11[k,lk] = d3kdxdx0dy0(
                 x0[k], y0[k], x[lk], y[lk], l)
             k21[k,lk] = d3kdxdy0dy0(
                  x0[k], y0[k], x[lk], y[lk], l )
             k12[k,lk] = d3kdydx0dy0(
                  x0[k], y0[k], x[lk], y[lk], l)
             k22[k,lk] = d3kdydy0dy0(
                 x0[k], y0[k], x[lk], y[lk], l)

    K = np.vstack([
        np.hstack([k11, k12]),
        np.hstack([k21, k22])
    ])
    return K


def nlp_with_grad(hyp, xtrain, ytrain, ztrain, sig_n):
    K = np.empty((2*len(xtrain), 2*len(xtrain)))
    buildK(xtrain, ytrain, xtrain, ytrain, hyp, K)
    Ky = K + sig_n**2*np.diag(np.ones(np.shape(K)[0]))
    Kyinv = np.linalg.inv(Ky)                # invert GP matrix
    alpha = Kyinv.dot(ztrain)
    nlp_val = 0.5*ztrain.T.dot(alpha) + 0.5*np.linalg.slogdet(Ky)[1]

    dK = build_dK(xtrain, ytrain, xtrain, ytrain, hyp)

    nlp_grad = np.array([
        -0.5*alpha.T.dot(dK[0].dot(alpha)) + 0.5*np.trace(Kyinv.dot(dK[0])),
        -0.5*alpha.T.dot(dK[1].dot(alpha)) + 0.5*np.trace(Kyinv.dot(dK[1]))
    ])

    return nlp_val, nlp_grad

def nlp(hyp, xtrain, ytrain, ztrain, sig_n):
    #try:
    #    return nlpgen(hyp, xtrain, ytrain, ztrain, sig_n)
    #except:
    #    pass
    K = np.empty((2*len(xtrain), 2*len(xtrain)))
    buildK(xtrain, ytrain, xtrain, ytrain, hyp, K)
    Ky = K + sig_n**2*np.diag(np.ones(np.shape(K)[0]))
    Kyinv = np.linalg.pinv(Ky)                # invert GP matrix
    return 0.5*ztrain.T.dot(Kyinv.dot(ztrain)) + 0.5*np.linalg.slogdet(Ky)[1]

def gpsolve(Ky, ft):
    L = np.linalg.cholesky(Ky)
    alpha = solve_triangular(
        L.T, solve_triangular(L, ft, lower=True, check_finite=False),
        lower=False, check_finite=False)

    return L, alpha

def nlpgen(hyp, xtrain, ytrain, ztrain, sig_n):
    K = np.empty((2*len(xtrain), 2*len(xtrain)))
    buildK(xtrain, ytrain, xtrain, ytrain, hyp, K)
    # print(np.linalg.eigvals(K))
    Ky = K + sig_n**2*np.diag(np.ones(2*len(xtrain)))  # add noise variance
    L, alpha = gpsolve(Ky, ztrain)
    return 0.5*ztrain.T.dot(alpha) + np.sum(np.log(L.diagonal()))

def energy(x):
    return x[1]**2/2 + U0*(1 + np.cos(x[0]))

def calc_tau_ex(q,p):
    kappa = energy([q, p])/(2*U0)
    kappinv = 1/kappa
    if kappa < 1:
        tau = 4*scipy.special.ellipk(kappa)
    else:
        tau = 2.0*scipy.special.ellipk(kappinv)/np.sqrt(kappa)
    return tau

def buildKreg(xtrainreg, ytrainreg, lp):

    N = len(xtrainreg)
        # ztrain = np.concatenate((ztrain1, ztrain2))
    Kp = np.empty((N, N))

    for k in range(N):
        for lk in range(N):
            Kp[k,lk] = f_kern(
                     xtrainreg[k], ytrainreg[k], xtrainreg[lk], ytrainreg[lk], lp)
    return Kp

def build_dKreg(xtrainreg, ytrainreg, lp):

    N = len(xtrainreg)
        # ztrain = np.concatenate((ztrain1, ztrain2))
    Kp = np.empty((N, N))

    dK = []
    for k in range(N):
        for lk in range(N):
            Kp[k,lk] = dkdlx(
                     xtrainreg[k], ytrainreg[k], xtrainreg[lk], ytrainreg[lk], lp)

    dK.append(Kp.copy())

    for k in range(N):
        for lk in range(N):
            Kp[k,lk] = dkdly(
                     xtrainreg[k], ytrainreg[k], xtrainreg[lk], ytrainreg[lk], lp)

    dK.append(Kp.copy())

    return dK

def nlpreg(hyp, xtrain, ytrain, ztrain, sig_n):
    K = buildKreg(xtrain, ytrain, hyp)
    Ky = K + sig_n**2*np.diag(np.ones(np.shape(K)[0]))
    Kyinv = np.linalg.pinv(Ky)                # invert GP matrix
    return 0.5*ztrain.T.dot(Kyinv.dot(ztrain)) + 0.5*np.linalg.slogdet(Ky)[1]

def nlpreg_with_grad(hyp, xtrain, ytrain, ztrain, sig_n):
    K = buildKreg(xtrain, ytrain, hyp)
    dK = build_dKreg(xtrain, ytrain, hyp)
    Ky = K + sig_n**2*np.diag(np.ones(np.shape(K)[0]))
    Kyinv = np.linalg.inv(Ky)                # invert GP matrix
    alpha = Kyinv.dot(ztrain)

    nlp_val = 0.5*ztrain.T.dot(alpha) + 0.5*np.linalg.slogdet(Ky)[1]

    nlp_grad = np.array([
        -0.5*alpha.T.dot(dK[0].dot(alpha)) + 0.5*np.trace(Kyinv.dot(dK[0])),
        -0.5*alpha.T.dot(dK[1].dot(alpha)) + 0.5*np.trace(Kyinv.dot(dK[1]))
    ])

    return nlp_val, nlp_grad

def guessP(x, y, l, xtrainp, ytrainp, ztrainp, Kyinvp):
    Ntest = 1
    N = len(xtrainp)
    Kstar = np.empty((Ntest, N))
    for k in range(Ntest):
        for lk in range(N):
            Kstar[k,lk] = f_kern(
                 x[k], y[k], xtrainp[lk], ytrainp[lk], l)
    Ef = Kstar.dot(Kyinvp.dot(ztrainp))
    return Ef

def calcQ(x,y, xtrain, ytrain, l, Kyinv, ztrain):
    # get \Delta q from GP on mixed grid.
    # temporarily, also \delta p is passed from the function
    Kstar = np.empty((2, 2*len(xtrain)))
    buildK(xtrain, ytrain, [x], [y], l, Kstar)
    qGP = Kstar.dot(Kyinv.dot(ztrain))
    f = qGP[1]
    return f, qGP[0]

def Pnewton(P, x, y, l, xtrain, ytrain, Kyinv, ztrain):
    Kstar = np.empty((2, 2*len(xtrain)))
    buildK(xtrain, ytrain, x, P, l, Kstar)
    pGP = Kstar.dot(Kyinv.dot(ztrain))
    f = pGP[0] - y + P
    return f

def calcP(x,y, l, lp, xtrainp, ytrainp, ztrainp, Kyinvp, xtrain, ytrain, ztrain, Kyinv):
    # as P is given in an implicit relation, use newton to solve for P
    # use the second GP on regular grid (q,p) for a first guess for P
    pgss = guessP([x], [y], lp, xtrainp, ytrainp, ztrainp, Kyinvp)
    res, r = newton(Pnewton, pgss, full_output=True, maxiter=5, disp=False,
       args = (np.array([x]), np.array([y]), l, xtrain, ytrain, Kyinv, ztrain))
    return res

def applymap(l, lp, Q0map, P0map, xtrainp, ytrainp, ztrainp, Kyinvp, xtrain, ytrain, ztrain, Kyinv):
    Nq = Q0map.shape[0]
    Np = P0map.shape[1]
    #init
    pmap = np.zeros([nm, Nq, Np])
    pmaptest = np.zeros([nm, Nq, Np])
    qmap = np.zeros([nm, Nq, Np])
    H = np.zeros([nm, Nq, Np])
    #set initial conditions
    pmap[0,:,:] = P0map
    pmaptest[0, :, :] = P0map
    qmap[0,:,:] = Q0map
    H[0,:,:] = energy([Q0map, P0map])

    # loop through all points on grid (q, p) and all time steps
    for i in range(0,nm-1):
        for k in range(0, Nq):
            for lk in range(0, Np):
                # first: set new P
                pmap[i+1, k, lk] = calcP(qmap[i,k,lk], pmap[i, k, lk], l, lp, xtrainp, ytrainp, ztrainp, Kyinvp, xtrain, ytrain, ztrain, Kyinv)
                # pmap[1, k, lk] = pe0[k, lk]
                # pmap[2, k, lk] = p1[k, lk]
        for k in range(0, Nq):
            for lk in range(0, Np):
                if np.isnan(pmap[i+1, k, lk]):
                    qmap[i+1,k,lk] = np.nan
                else:
                    # then: set new Q via calculating \Delta q and adding q
                    qmap[i+1, k, lk], dpmap = calcQ(qmap[i,k,lk], pmap[i+1,k,lk],xtrain, ytrain, l, Kyinv, ztrain)# + qmap[i, k, lk]
                    qmap[i+1, k, lk] = np.mod(qmap[i+1,k,lk]+ qmap[i, k, lk], 2.0*np.pi)
                    # testing GP on mixed grid for p -> THIS does NOT work
                    # P = p + \delta p -> as the GP is fitted for -\delta p -> additional minus
                    pmaptest[i+1, k, lk] = pmap[i,k,lk] - dpmap
                H[i+1, k, lk] = energy([qmap[i+1, k, lk], pmap[i+1, k, lk]])
    return qmap, pmap, H, pmaptest

def dydt(y_, t):
    ydot = np.zeros([2])
    ydot[0] = y_[1] # dx/dt
    ydot[1] = -np.sin(y_[0] + np.pi)# - eps*np.sin(y_[0] - om*(t))#- 0.5*np.pi) # dpx/dt
    return ydot

def dydt_ivp(t, y):
    ydot = np.zeros([2])
    ydot[0] = y[1] # dx/dt
    ydot[1] = -np.sin(y[0]+np.pi)# - eps*np.sin(y_[0] - om*(t))#- 0.5*np.pi) # dpx/dt
    return ydot

def integrate_pendulum(q0, p0, t):
    # yint = np.zeros([len(t), 2, len(q0)]) # initial values for y
    ysint = np.zeros([len(t), 2, len(q0)]) # initial values for y

    # for k in range(len(q0)):
    #     yint[:, :, k] = spint.odeint(dydt, [q0[k], p0[k]], t)
        # ysint[:,:, k] = intode([q0[k], p0[k]], t, dtsymp)
        # ysint[:,0, k] = np.mod(ysint[:,0, k], 2*np.pi)
    ysint = []
    for ik in range(len(q0)):
        res_int = solve_ivp(dydt_ivp, [t[0], t[-1]], np.array((q0[ik], p0[ik])), max_step=0.001, method='RK45')
        temp = res_int.y
        # temp[0] = np.mod(temp[0], 2*np.pi)
        ysint.append(temp)

    return ysint
# compute log-likelihood according to RW, p.19
def solve_cholesky(L, b):
    return solve_triangular(
        L.T, solve_triangular(L, b, lower=True, check_finite=False),
        lower=False, check_finite=False)

# negative log-posterior
def nll_chol(hyp, x, y, build_K=build_K):
    K = np.empty((len(x), len(x)))
    build_K(x, x, hyp[:-1], K)
    Ky = K + np.abs(hyp[-1])*np.diag(np.ones(len(x)))
    L = np.linalg.cholesky(Ky)
    alpha = solve_cholesky(L, y)
    ret = 0.5*y.T.dot(alpha) + np.sum(np.log(L.diagonal()))
    #print(hyp, ret)
    return ret

def nll(hyp, x, y, neig=8, build_K=build_K):
    K = np.empty((len(x), len(x)))
    build_K(x, x, np.abs(hyp[:-1]), K)
    Ky = K + np.abs(hyp[-1])*np.diag(np.ones(len(x)))
    w, Q = eigsh(Ky, neig, tol=max(1e-6*np.abs(hyp[-1]), 1e-15))
    while np.abs(w[0]-hyp[-1])/hyp[-1] > 1e-6 and neig < len(x):
        if neig > 0.05*len(x):  # TODO: get more stringent criterion
            try:
                return nll_chol(hyp, x, y, build_K)

            except:
                print('Warning! Fallback to eig solver!')
        neig =  2*neig
        w, Q = eigsh(Ky, neig, tol=max(1e-6*hyp[-1], 1e-15))

    alpha = Q.dot(np.diag(1.0/w).dot(Q.T.dot(y)))

    ret = 0.5*y.T.dot(alpha) + 0.5*(np.sum(np.log(w)) + (len(x)-neig)*np.log(np.abs(hyp[-1])))
    #print(hyp, ret)
    return ret

def plot_pendulum(ysint):
    Np = ysint.shape[3]
    Nq = ysint.shape[2]
    plt.figure()
    for kp in range(Np):
        for kq in range(Nq):
            plt.plot(ysint[:,0,kq,kp], ysint[:,1,kq,kp],'.', color = 'k', markersize = 1.0)
            plt.plot(ysint[0,0,kq,kp],ysint[0,1,kq,kp],'o', color = 'b', markersize = 4.9)
            plt.plot(ysint[-1,0,kq,kp],ysint[-1,1,kq,kp],'o', color = 'r', markersize = 4.9)
    plt.xlabel(r'$q(n{\tau})$')
    plt.ylabel(r'$p(n{\tau})$')

def plot_map(qmap, pmap, H, Q0map, P0map):
    plt.figure()
    for i in range(0, qmap.shape[2]):
        plt.plot(qmap[:,:,i], pmap[:,:,i], ',', markersize = 1.5)
        # plt.plot(np.mod(qmap[:,:,i]-np.pi, 2*np.pi)-np.pi, pmap[:,:,i], ',', color = 'k', markersize = 1.5)
        # plt.plot(qmap_RBF[:,:,i], pmap_RBF[:,:,i], '.', color = 'k', markersize = 1.5)
    # plt.ylim([-4, 4])
    plt.xlabel('q')
    plt.ylabel('p')

    plt.figure()
    for i in range(0, qmap.shape[2]):
        # i = 0
        plt.plot(H[:,:,i]/energy([Q0map[:,i], P0map[:,i]]))#/H_RBF[0,:,i])
    plt.xlabel(r'$\tau$')
    plt.ylabel('E/E_0')

def gamma(d):
    x=1.0000
    for i in range(20):
        x = x-(pow(x,d+1)-x-1)/((d+1)*pow(x,d)-1)
    return x

def bluenoise(d,n, method, qmin, qmax, pmin, pmax):

    g = gamma(d)
    alpha = np.zeros(d)
    for j in range(d):
        alpha[j] = pow(1/g,j+1) %1
    z = np.zeros((n, d))
    # z = []
    i = 0
    iz = 0
    while iz < n:
        random_point = (0.5 + alpha*(i+1)) %1
        random_point[0] = random_point[0]*(np.abs(qmin)+qmax)-np.abs(qmin)
        random_point[1] = random_point[1]*(np.abs(pmin)+pmax)+pmin
        if (method == 'trapped') & (energy(random_point)/2 < 1.0):
            z[iz] = random_point
            iz = iz + 1
        elif (method == 'passing') & (energy(random_point)/2 > 1.0):
            z[iz] = random_point
            iz = iz + 1
        elif (method == 'all'):
            z[iz] = random_point
            iz = iz + 1
        i = i + 1
        #     while len(points) < num_points:
        # random_point = Point(temp[i,:])
        # if (random_point.within(poly)):
        #     points.append(random_point)
        # i = i + 1
    return z
