#=
Timings:
17s
added free-slots
18s
passing eta
17.6s
F=Drift(r) -> Drift!(F,R)
11.6s
Rold to MMatrix
12s
major StaticArray use
6.8s;
internal timing:
@time main 4.0s
replacing global Ntry ... with ::VMCStats
3.5s
empty!(free_slots)
3.3s

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  
=#
# activate dmc_project
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))

using Printf
using Distributions
import LinearAlgebra: norm
using Statistics
using Random
using StaticArrays
using BenchmarkTools

#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.01        # 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 DEAD  = 0
const ALIVE = 1
const CHILD = 2

const minstep = 1e-5
const maxstep = 10.0

const Coord = MMatrix{D,N,Float64,D*N}

mutable struct Walker
    R     ::Coord
    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
@inline function Elocal(R::Coord)
    x1, y1, z1 = R[1,1], R[2,1], R[3,1]
    x2, y2, z2 = R[1,2], R[2,2], R[3,2]

    r1  = sqrt(x1*x1 + y1*y1 + z1*z1)
    r2  = sqrt(x2*x2 + y2*y2 + z2*z2)

    dx  = x1 - x2
    dy  = y1 - y2
    dz  = z1 - z2
    r12 = sqrt(dx*dx + dy*dy + dz*dz)

    return -alpha^2 - (2-alpha)/r1 - (2-alpha)/r2 + 1/r12
end

mutable struct VMCStats
    step::Float64
    Naccept::Int
    Ntry::Int
end



# initialization 
function init() 
    println("init")
    # 
    # Generate Nwx walker
    #
    walker = [Walker(zero(Coord),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 = Coord(rand(D, N))        # coordinates
        walker[iw].R = copy(R)
        walker[iw].psi = ln_psi2(R)  # wave function (ln(psi^2))
        walker[iw].E   = Elocal(R)        
    end
    vmcstat = VMCStats(3.1,0,0)
    println("init done")
    walker, vmcstat
end

# ln(\varphi_T^2) = 2\ln(\varphi_T)
@inline function ln_psi2(R::Coord)
    x1, y1, z1 = R[1,1], R[2,1], R[3,1]
    x2, y2, z2 = R[1,2], R[2,2], R[3,2]

    r1 = sqrt(x1*x1 + y1*y1 + z1*z1)
    r2 = sqrt(x2*x2 + y2*y2 + z2*z2)

    return -2*alpha*(r1+r2)
end

@inline 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::Coord, vmcstat::VMCStats)
    

    psi = ln_psi2(R)
    E   = Elocal(R)
    step = vmcstat.step

    @inbounds for i in 1:N
        d1 = step*(rand() - 0.5)
        d2 = step*(rand() - 0.5)
        d3 = step*(rand() - 0.5)

        R[1,i] += d1
        R[2,i] += d2
        R[3,i] += d3

        psi_new = ln_psi2(R)
        diff = psi_new - psi

        vmcstat.Ntry += 1
        accept = diff >= 0.0 || exp(diff) > rand()

        if accept
            vmcstat.Naccept += 1
            psi = psi_new
            E = Elocal(R)
        else
            R[1,i] -= d1
            R[2,i] -= d2
            R[3,i] -= d3
        end
    end

    return psi, E
end

                
# adjust step to keep VMC acceptance 40-60 %
function adjust_step(vmcstat::VMCStats)
    acceptance = vmcstat.Naccept/vmcstat.Ntry*100.0    
    if acceptance<40.0
        vmcstat.step *= 0.9
    end
    if acceptance>60.0
        vmcstat.step *= 1.1
    end
    vmcstat.step = max(minstep,vmcstat.step)
    vmcstat.step = min(maxstep,vmcstat.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!(F::Coord, R::Coord)
    x1, y1, z1 = R[1,1], R[2,1], R[3,1]
    x2, y2, z2 = R[1,2], R[2,2], R[3,2]

    r1 = sqrt(x1*x1 + y1*y1 + z1*z1)
    r2 = sqrt(x2*x2 + y2*y2 + z2*z2)

    c = -2*alpha

    F[1,1] = c*x1/r1
    F[2,1] = c*y1/r1
    F[3,1] = c*z1/r1

    F[1,2] = c*x2/r2
    F[2,2] = c*y2/r2
    F[3,2] = c*z2/r2

    return nothing
end



# one diffusion+drift step in DMC
@inline function diffusion_drift_step!(R::Coord, F::Coord, R1::Coord, R2::Coord, eta::Coord)
    if order == 1
        randn!(eta)
        Drift!(F, R)
        @inbounds for j in 1:N, i in 1:D
            R[i,j] += lambda*tau*F[i,j] + sqrt(2*lambda*tau)*eta[i,j]
        end

    elseif order == 2
        randn!(eta)
        @inbounds for j in 1:N, i in 1:D
            R1[i,j] = R[i,j] + sqrt(lambda*tau)*eta[i,j]
        end
        Drift!(F, R1)
        @inbounds for j in 1:N, i in 1:D
            R2[i,j] = R1[i,j] + 0.5*lambda*tau*F[i,j]
        end

        Drift!(F, R2)
        @inbounds for j in 1:N, i in 1:D
            R[i,j] = R1[i,j] + lambda*tau*F[i,j]
        end

        randn!(eta)
        @inbounds for j in 1:N, i in 1:D
            R[i,j] += sqrt(lambda*tau)*eta[i,j]
        end
    end

    return nothing
end


function lnG(Rnew::Coord, Rold::Coord, Fold::Coord)
    if order != 1
        println("lnG only for 1st order code")
        exit(1)
    end

    Drift!(Fold, Rold)

    s = 0.0
    @inbounds for j in 1:N, i in 1:D
        x = Rnew[i,j] - Rold[i,j] - lambda*tau*Fold[i,j]
        s += x*x
    end

    return -s / (4*lambda*tau)
end


        
function main()
    #
    # Main program 
    # 
    # initialize
    walker, vmcstat = 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, vmcstat)
            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, vmcstat)
            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(vmcstat)
    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
    copies = zeros(Int64,Nwx)
    Ntry = 0 ::Int64
    Nacc = 0 ::Int64

    Rold = zero(Coord)
    F    = zero(Coord)
    R1   = zero(Coord)
    R2   = zero(Coord)
    eta  = zero(Coord)
    
    free_slots = Int[]
    sizehint!(free_slots, Nwx)

    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
            @inbounds Rold  .= walker[iw].R
            #
            diffusion_drift_step!(walker[iw].R, F, R1, R2, eta)
            # 
            EL::Float64 = Elocal(walker[iw].R)    

            if order==1
                # metropolis
                # careful with the order of R's 
                Wnew::Float64 = lnG(Rold, walker[iw].R, F)+ln_psi2(walker[iw].R)  # lnG(new->old) + lnpsi2(new)
                Wold::Float64 = lnG(walker[iw].R, Rold, F)+ln_psi2(Rold)          # 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

        empty!(free_slots)
        @inbounds for i in 1:Nwx
            if walker[i].alive == 0
                push!(free_slots, i)
            end
        end
        
        next_free = 1
        # Branching
        @inbounds for iw in 1:Nwx
            if walker[iw].alive != 1; continue; end            
            if copies[iw] == 1
                continue
            elseif copies[iw] == 0
                walker[iw].alive = 0
                continue
            end
            srcR = walker[iw].R
            srcE = walker[iw].E
            srcpsi = walker[iw].psi

            for _ in 1:copies[iw]-1
                if next_free > length(free_slots)
                    break
                end
               
                
                inew = free_slots[next_free]
                next_free += 1

                walker[inew].R = copy(srcR)
                walker[inew].E = srcE
                walker[inew].psi = srcpsi
                walker[inew].alive = 2
            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

@time main()