#=
QMC course 2018

 JULIA
DMC on a helium atom with trial wave function 
 \varphi_T = e^{-alpha*r1}e^{-alpha*r2} ; r1,r2 distance of els from proton
 one el in spin-up and the other in spin-down state (Fermi rule is ok)
  
Measures only total energy
Hamiltonian in a.u. (Z=2) :

H = -1/2* (nabla_1^2 + nabla_2^2) - 2/r1 - 2/r2 + 1/r12  
=#

using Printf
using Distributions
import LinearAlgebra: norm
using Statistics
#using Random

#Random.seed!(1234)


#
# Parameters
#
const order = 2   # 1 or 2
println("He atom DMC using algorithm of order ",order)

const lambda = 0.5       # hbar^2/(2m) in au
const alpha = 27/16      # trial wafe function parameter
const tau = 0.005        # imaginary time step
const blocksize = 1000   # DMC data block size


const NVMC = 1000        # number of VMC steps
const NDMC = 10000       # number of DMC steps
const Nwx = 20000        # max number of walkers
const Nw_target = 1000   # on average this many walkers
const kappa = 0.1        # how zealously DMC tries to keep Nw_target walkers  
const N = 2              # number of electrons 
const D = 3              # dimension  
const Eexact = -2.903724377  # accurate ground state energy, for comparison


const minstep = 1e-5
const maxstep = 10.0



mutable struct Walker
    R     ::Array{Float64,2}
    alive ::Int64
    psi   ::Float64
    E     ::Float64
end

# analytical energy in the trial wave function, for checking
function Eexpec(alpha)
    return alpha^2-27/8*alpha
end

# local energy = H \varphi_T/\varphi_T
function Elocal(R::Array{Float64,2})
    r12 = norm(R[:,1]-R[:,2])
    r1  = norm(R[:,1])
    r2  = norm(R[:,2])
    Elocal = -alpha^2 - (2-alpha)/r1 - (2-alpha)/r2 + 1/r12
    Elocal
end


# initialization 
function init() ::Array{Walker,1}
    global step,Naccept,Ntry
    println("init")
    # 
    # Generate Nwx walker
    #
    walker = [Walker(zeros(Float64,D,N),0,0.0,0.0) for i in 1:Nwx] # R, alive=0, psi=0, E=0
    println("generating $Nw_target walkers")
    # set Nw_target walkers
    for iw in 1:Nw_target
        walker[iw].alive = 1
        R = rand(D,N)  ::Array{Float64,2}           # coordinates
        walker[iw].R = copy(R)
        walker[iw].psi = ln_psi2(R)  # wave function (ln(psi^2))
        walker[iw].E   = Elocal(R)        
    end
    Ntry=0
    Naccept=0
    step = 3.1
    
    println("init done")
    walker
end

# ln(\varphi_T^2) = 2\ln(\varphi_T)
function ln_psi2(R::Array{Float64,2})
    r1 = norm(R[:,1])
    r2 = norm(R[:,2])
    -2*alpha*(r1+r2)
end

function metro(Wold ::Float64, Wnew ::Float64)
    if Wnew>Wold
        return true
    end
    if rand()<exp(Wnew-Wold)
        return true
    else
        return false
    end
        
    
end

# one VMC step
function vmc_step!(R::Array{Float64,2})    
    global Ntry,Naccept
    psi = ln_psi2(R)
    E   = Elocal(R) 
    @inbounds for i in 1:N
        d   = step*(rand(D) .- 0.5 ) ::Array{Float64,1}
        R[:,i] += d
        psi_new = ln_psi2(R)
        diff = psi_new-psi        
        
        accept = false
        Ntry +=1
        if diff>=0.0
            accept = true            
        else
            if exp(diff) > rand() 
                accept = true
            end
        end
        if accept
            Naccept +=1
            psi = psi_new
            E = Elocal(R)
        else
            # revert to old value
            R[:,i] -= d
        end        
    end
    psi,E
end
                
# adjust step to keep VMC acceptance 40-60 %
function adjust_step()
    global step
    acceptance = Naccept/Ntry*100.0    
    if acceptance<40.0
        step *= 0.9
    end
    if acceptance>60.0
        step *= 1.1
    end
    step = max(minstep,step)
    step = min(maxstep,step)
end

function save_R(R::Array{Float64,2},mode="w")    
    open("R",mode) do f
        @inbounds for i in 1: size(R,2)
            for k in 1:size(R,1)
                print(f,R[k,i]," ")
            end
            println(f," ")
        end
    end
end

function save_r(R::Array{Float64,2},mode="w")    
    open("r",mode) do f
        rmin = 1.e5 ::Float64
        rr = 0.0 ::Float64
        @inbounds for i in 1: size(R,2)
            rr = norm(R[:,i]) 
            rmin  = min(rmin,rr)
        end
        println(f,rmin)
    end
end

# DMC drift, for ith electron F_i = 2\nabla_i \varphi_T /\varphi_T
@inline function Drift(R::Array{Float64,2})
    # 2*nabla_i \varphi_T / \varphi_T
    r1 = norm(R[:,1])
    r2 = norm(R[:,2])
    -2*alpha* hcat(R[:,1]/r1,R[:,2]/r2)
end

# one diffusion+drift step in DMC
@inline function diffusion_drift_step(R ::Array{Float64,2}, F ::Array{Float64,2}, R1 ::Array{Float64,2}, R2 ::Array{Float64,2})

    
    if order==1
        eta = reshape(randn(N*D),(D,N))
        F = Drift(R)
        @. R += lambda*tau*F + sqrt(2*lambda*tau)*eta # elementwise
        
    elseif order==2
        
        #
        # diffusion(tau/2)-drift(tau)-diffusion(tau/2)
        # --------------------------------------------
        # S. Chin Phys. Rev. A42, 6991 (1990)
        # Eq. (25)
        # y = x + sqrt(dt/2)*eta
        # x = y + dt*G(y+dt/2*G(y))+ sqrt(dt/2)*eta
        #
        # Note: our diffusion tau is 2*chin's dt, so here
        #
        # R1 = R + sqrt(lambda*tau)*eta
        # R2 = R1+0.5*lambda*tau*G(R1)
        # R = R1 + lambda*tau*G(R2)         
        # R = R + sqrt(lambda*tau)*eta
        
        
        
        # tau/2 diffusion           
        #eta = reshape(randn(N*D),(D,N)) # don't know if this is good enough
        eta = reshape(rand(Normal(0,1),D*N),(D,N))
        @. R1 = R + sqrt(lambda*tau)*eta    
        # tau drift            
        F = Drift(R1)                    
        @. R2 = R1 + 0.5*lambda*tau*F  
        F = Drift(R2)                    
        @. R = R1 + lambda*tau*F         
        # tau/2 diffusion
        #eta = reshape(randn(N*D),(D,N))
        eta = reshape(rand(Normal(0,1),D*N),(D,N))
        @. R += sqrt(lambda*tau)*eta
        

    end
        
end



function lnG(Rnew ::Array{Float64,2}, Rold ::Array{Float64,2}, Fold ::Array{Float64,2})     
    # no need to compute EL(R)-ET, it's used symmetrically
    # notice order:
    # Rnew = Rold + lambda*tau*F(Rold) + sqrt(2*lambda*tau)*eta
    # so
    # lnG(Rnew,Rold,tau) has exponent  Rnew - Rold - lambda*tau*Fold
    if order != 1
        println("lnG only for 1st order code")
        exit(1)
    end
    dRF = 0.0 ::Float64
    
    Fold = Drift(Rold) 
    dRF -= norm(Rnew - Rold - lambda*tau*Fold)^2
    lnG = dRF/(4*lambda*tau) 
    lnG
end

        
function main()
    #
    # Main program 
    # 
    # initialize
    walker = init() 
    # thermalization
    println("thermalizing")
    Nw = Nw_target # for VMC
    @inbounds for i in 1:10
       @inbounds for iw in 1:Nw         
            psi,E = vmc_step!(walker[iw].R)
            walker[iw].psi = psi
            walker[iw].E   = E
        end
    end

    println("thermalization done")

    filename = string("E_heatom_",order,"_tau=",tau)
    println("output: ",filename)
    # init file output
    open(filename,"w") do f
        println(f," ")
    end
    #
    # VMC
    #
    E_ave = 0
    nE  = 0
    @inbounds for ivmc in 1:NVMC
        E = 0.0 
        @inbounds for iw in 1:Nw         
            psiw ,Ew = vmc_step!(walker[iw].R)
            walker[iw].psi = psiw
            walker[iw].E   = Ew
            E += Ew 
        end
        E_ave += E/Nw        
        nE +=1
        
        if ivmc%10 == 0            
            #println("VMC Elocal=$(E/Nw)  <E>=$(E_ave/nE)  <E>_analytical=$(Eexpec(alpha))")
            @printf("VMC  E = %.10f <E> = %.10f  <E>_analytical = %.10f\n",E/Nw,E_ave/nE, Eexpec(alpha))            
        end
        adjust_step()
    end
    
    #
    # DMC
    #
    ET = E_ave/nE  # trial energy, to be updated
    println("ET = $ET")
    
    Ntherm = floor(Int64,1.0/tau) # start DMC measurement after excitations have died out; system specific !
    idmc = -Ntherm

    E_ave = 0
    nE = 0
    
    Edat  = 0.0 ::Float64  # single block E 
    E2dat = 0.0 ::Float64  # single block E^2
    Ndat = 0  

    Eb = Array{Float64,1}() # block data E
    E2b = Array{Float64,1}() # block data E^2


    
    idat = 0 ::Int64
    Rold = zeros(Float64,(D,N))
    F    = zeros(Float64,D,N)

    copies = zeros(Int64,Nwx)
    Ntry = 0 ::Int64
    Nacc = 0 ::Int64

    R1 = zeros(Float64,(D,N))
    R2 = zeros(Float64,(D,N))
   
    
    while idmc<NDMC
        #
        # take a DMC step in each walker
        #
        
        E = 0.0  ::Float64
        n = 0
        @inbounds for iw in 1:Nwx
            copies[iw] = 0
            if walker[iw].alive==0; continue; end  # not alive, skip
            ELold = walker[iw].E  ::Float64
            Rold  = copy(walker[iw].R)
            #
            diffusion_drift_step(walker[iw].R, F, R1, R2)
            # 
            
            EL = Elocal(walker[iw].R)    ::Float64

            if order==1
                # metropolis
                # careful with the order of R's 
                Wnew = lnG(Rold, walker[iw].R, F)+ln_psi2(walker[iw].R) ::Float64 # lnG(new->old) + lnpsi2(new)
                Wold = lnG(walker[iw].R, Rold, F)+ln_psi2(Rold)         ::Float64 # lnG(old->new) + lnpsi2(old)
                                
                accept = metro(Wold,Wnew)
            else
                accept = true
            end
            Ntry +=1
            if accept
                Nacc +=1
                walker[iw].E = EL
            else
                EL = ELold
                walker[iw].E = ELold
                walker[iw].R = Rold                
            end

            E  += EL
            n  += 1
            if idat>0
                Edat  += EL
                E2dat += EL^2
                Ndat  +=1  
            end
            
            weight  = exp(-tau*(0.5*(ELold+EL)-ET)) ::Float64 # symmetric under R<->R'
            copies[iw]  = floor(Int64,weight + rand()) ::Int64
        end
        # Branching
        @inbounds for iw in 1:Nwx
            if walker[iw].alive !=1; continue; end  # not alive, or child, skip
            if copies[iw]==1; continue; end        # one copy is already there
            if copies[iw]==0
                # walker dies
                walker[iw].alive = 0
                continue
            end
            # copies[iw]>1
            # copy the walker to empty slots 
            @inbounds for inew in 1:copies[iw]-1   # with copies=3 this is 1,2
                @inbounds for iw2 in 1:Nwx
                    if walker[iw2].alive>0; continue; end 
                    # free slot
                    walker[iw2].R = copy(walker[iw].R)
                    walker[iw2].E =  walker[iw].E 
                    walker[iw2].alive = 2 # NOTE mark as child
                    break
                end
            end
        end
       @inbounds for iw in 1:Nwx
            if walker[iw].alive == 2
                walker[iw].alive =1
            end
        end

        Nw = sum([walker[i].alive for i in 1:Nwx])
        E_ave += E/n # don't use Nw as energy counter, it's not correct
        nE += 1

        idmc +=1
        
       
        
        if Nw==Nwx
            println("hit max number of walkers")
            exit(1)
        end

        # update local energy; on average Nw_target walkers
        # A too large factor kappa will cause a bad feedback, a too small will let number of walkers get too large
        ET = E_ave/nE + kappa* log(Nw_target/Nw)

        if idmc<=0
            @printf("DMC %10d  E = %.10f <E> = %.10f ET = %.10f <E>_exact = %.10f  %6d Walkers \n",
                    idmc, E/Nw, E_ave/nE, ET, Eexact, Nw)
            if order==1;@printf("DMC acceptance = %.5f %%\n",Nacc*100.0/Ntry);end
        end

        # block data
        # ==========
        # screen and file output
        if idat== blocksize            
            append!(Eb,Edat/Ndat)
            append!(E2b,E2dat/Ndat)
            Eave = mean(Eb) :: Float64            
            if size(Eb,1)>1
                EStdError  = std(Eb)/sqrt(size(Eb,1)) ::Float64 # sigma/sqrt(N) for N data points
            else
                EStdError  =  0.0 ::Float64
            end
            @printf("DMC %10d E %.10f <E> = %.10f +/- %.10f ET = %.10f <E>_exact = %.10f  %6d Walkers \n",
                    idmc, Edat/Ndat ,Eave, EStdError, ET, Eexact, Nw)
            open(filename,"a") do f
                println(f,tau," ",Eave," ",EStdError," ",Eexact)
            end
            
            Edat  = 0
            E2dat = 0
            Ndat = 0
            idat = 0
            
        end
        
        # 
        if idmc==0
            println("THERMALIZATION ENDS")
            E_ave = 0
            nE = 0
            idat = 1
        end
        if idmc>0
            idat += 1
        end
        
    end
end

main()