#!/usr/bin/env python3

import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt
import scipy.special as sp

def Main():

    N = 40
    nbar = 1

    x = np.arange(0.16, 0.18, 1e-4)
    y = []
    z = []
    for zt in x:

        f, err = find_f(N, nbar, zt)
        energy = Energy(f, N, nbar, zt)
        nvar   = nVariance(f, N, nbar)
        print(f'{zt:.1e}\t{energy:+.2f}\t{nvar:+.2f}\t{err:+.2e}')
        y.append(energy)
        z.append(nvar)

    y = [k-y[0] for k in y]
    w = [abs(z[i+1]-z[i-1]) for i in range(1, len(z)-1)]
    print(f'zt_c ~ {x[np.argmax(w)+1]}')

    fig, ax = plt.subplots(1, 1)
    ax.plot(x, y, label='$\\varepsilon_t-\\varepsilon_0$', lw=0.5)
    ax.plot(x, z, label='$v$', lw=0.5)
    ax.plot(x[1:-1], w, label='$\\delta v$', lw=0.5)
    ax.set_xlabel('$zt$')
    ax.legend()
    fig.savefig('E_zt.pdf')
    plt.close(fig)

    return

def find_f(N, nbar, dt):
    k = np.arange(N)
    x0 = np.zeros(N+2)
    x0[:N] = np.sqrt(np.exp(-nbar)*nbar**k/sp.factorial(k))
    x0[-2], x0[-1] = -1, -1

    x = opt.root(Lprime, x0=x0, method='lm', args=(N, nbar, dt)).x

    err = np.sqrt(np.sum(Lprime(x, N, nbar, dt)**2))
    return x[:N], err

def Lprime(x, N, nbar, zt):
    f, λ = x[:N], x[N:]

    k = np.arange(N)
    sqrt_k = np.sqrt(k)

    u = np.sum(f[:-1]*f[1:]*sqrt_k[1:])

    L_f = np.zeros(N)
    L_f[1:]  -= 2*zt*u*f[:-1]*sqrt_k[1:]
    L_f[:-1] -= 2*zt*u*f[1:]*sqrt_k[1:]

    L_f += k**2*f - 2*λ[0]*f - 2*λ[1]*k*f

    c0 = np.sum(f**2) - 1
    c1 = np.sum(k*f**2) - nbar

    return np.append(L_f, [c0, c1])

def Energy(f, N, nbar, zt):
    k = np.arange(N)
    sqrt_k = np.sqrt(k)

    u = np.sum(f[:-1]*f[1:]*sqrt_k[1:])
    v = np.sum((f*k)**2)
    return v/2 - nbar**2 - zt*u**2

def nVariance(f, N, nbar):
    k = np.arange(N)
    v = np.sum((f*k)**2)
    return v - nbar**2

Main()