#
# 1-dimensional DMC code
#
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
#
# global constants
#        h2m   - hbar^2/mass
h2m = 1.0

#=============================================================================
def V_HO(x):
    return 0.5*x**2
#=============================================================================
def V_Morse(x):
    return 0.5*(np.exp(-2*x)-2*np.exp(-x))
#=============================================================================
def trial_function(x):
    return 1.0/(1+x**2)
#=============================================================================
def Psi0_HO(x):
    return (np.exp(-0.5*x**2)/np.pi**0.25)
#=============================================================================
def Psi0_Morse(x):
    return np.sqrt(2)*np.exp(-np.exp(-x)-0.5*x)
#=============================================================================
def h_psi_trial(poten,psi_trial,xwalkers):
    nw = xwalkers.shape[-1]
    dx = 1.e-5
    psic = psi_trial(xwalkers)
    xm = xwalkers-dx*np.ones(nw)
    xp = xwalkers+dx*np.ones(nw)
    t_psi = -0.5*h2m*(psi_trial(xp)+psi_trial(xm)-2*psic)/dx**2
    v_psi = poten(xwalkers)*psic
    h_psi = t_psi+v_psi
    return h_psi
#=============================================================================
#
# DMC step subroutine:
# input: 
#        poten - the interaction, e.q V_HO or V_Morse
#        xold  - old (current) walkers positions
#        eold  - old (current) energy estimate 
#        ntarget - desired number of walkers 
#        dtau  - time step
# output: 
#        xnew  - new walkers positions
#        enew  - new energy estimate 
#
def dmc_step(poten,xold,eold,ntarget,dtau):

    ##########################################
    #                                        #
    #                COMPLETE                #
    #                                        #
    #   with numpy this can take 14 lines    #
    #                                        #
    ##########################################
    # once start working remove the following line
    xnew = xold ; enew=eold
    ##########################################
    
    return(xnew,enew)
#=============================================================================
#
# DMC manage subroutine: iterates the DMC step and plots
# the wave function every "mprint" steps.
#
def manage_dmc(poten,xwalkers0,enr,nwalkers0,tau,nsteps,mprint=50):
    global ee_grow_list,ee_hpsi_list,xwalkers

    nw0=nwalkers0
    dtau=tau/nsteps
    xwalkers = xwalkers0

    # plot the walkers distribution
    hist0, bins = np.histogram(xwalkers,bins=60,density=True,range=(-4.0,4.0))
    xbins = [0.5*(bins[i]+bins[i+1]) for i in range(0,len(bins)-1)]
    plt.plot(xbins,hist0)

    ee_grow_list = np.zeros(nsteps)
    ee_hpsi_list = np.zeros(nsteps)
    ee_grow = enr
    for istep in range(nsteps):

        xwalkers,ee_grow = dmc_step(poten,xwalkers,ee_grow,nw0,dtau)
        ee_grow_list[istep]= ee_grow

        h_psi = h_psi_trial(poten,trial_function,xwalkers)
        ee_hpsi = np.sum(h_psi)/np.sum(trial_function(xwalkers))
        ee_hpsi_list[istep] = ee_hpsi
        
        nw = xwalkers.shape[-1]
        print(f"\tDMC step {istep:6d}   nwalkers={nw:9d}    "+
              f"E={ee_grow:11.6f}  {ee_hpsi:11.6f}")

        if (np.remainder(istep,mprint)==0):
            hist, bins = np.histogram(xwalkers,bins=bins,density=True)
            plt.plot(xbins,hist)

    plt.grid(True,linestyle=':')
    plt.xlim((-4, 4))
    pltfile.savefig()
    plt.close()

#=============================================================================
def PlotEnergyTau(tau,enr0=0.5,denr=0.1):
#   plot energy convergence
    ns = ee_grow_list.shape[-1]
    tau_list = np.linspace(0,tau,ns)
    plt.plot(tau_list,ee_grow_list,label='MC - growth')
    plt.plot(tau_list,ee_hpsi_list,label='MC - H$\Psi$')
    plt.plot(tau_list,enr0*np.ones(nsteps),'--')
    plt.grid(True,linestyle=':')
    plt.legend(loc='best',frameon=False)
    plt.ylim((enr0-denr,enr0+denr))
    pltfile.savefig()
    plt.close()
#=============================================================================
def PlotWF(xwalkers0,xwalkersc,Psi0):

    hist0, bins = np.histogram(xwalkers0,bins=60,density=True,range=(-4.0,4.0))
    histc, bins = np.histogram(xwalkersc,bins=bins,density=True)
    xbins = [0.5*(bins[i]+bins[i+1]) for i in range(0,len(bins)-1)]
    plt.plot(xbins,hist0)
    plt.plot(xbins,histc)

    xprt = np.asarray(xbins)
    dx = xprt[1]-xprt[0]
    plt.plot(xprt,Psi0(xprt)/(np.sum(Psi0(xprt))*dx),color='black',lw=2.)

    plt.grid(True,linestyle=':')
    plt.xlim((-4, 4))
    pltfile.savefig()
    plt.close()

#=============================================================================
#
# The driver part of the code
#
# Fixing random state for reproducibility
np.random.seed(19670315)

# define the potential
poten = V_HO
psi0  = Psi0_HO

# opens a pdf file for all the plots
pltfile = PdfPages(poten.__name__+'.pdf')

# initial guess for energy and walkers
nwalkers=10000
enr=2.0
xwalkers0 = 1.5+2*np.random.ranf((nwalkers,))

# set tau and the number of walkers
# dtau = tau/nsteps
tau=40.0
nsteps=400
manage_dmc(poten,xwalkers0,enr,nwalkers,tau,nsteps)

PlotWF(xwalkers0,xwalkers,psi0)
PlotEnergyTau(tau,enr0=0.5,denr=0.04)

pltfile.close()
#=============================================================================