#=
Timings:
17s
added free-slots
18s
passing η
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
R, alive, psi, E as separate arrays
2.79s

 JULIA
DMC on a helium atom with trial wave function 
 \varphi_T = e^{-α*r1}e^{-α*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
#using LoopVectorization # @turbo

#Random.seed!(1234)


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

# Trial wf parameters
const α = 2.0 # 1.847529

# Constants
const lambda = 0.5       # hbar^2/(2m) in au
const tau = 0.001        # imaginary time step
const blocksize = 10000  # DMC data block size

const NVMC = 1000        # number of VMC steps
const NDMC = 10000000    # 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

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


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

# local energy = H \varphi_T/\varphi_T
@inline function Elocal(R, iw)
    x1, y1, z1 = R[1,1,iw], R[2,1,iw], R[3,1,iw]
    x2, y2, z2 = R[1,2,iw], R[2,2,iw], R[3,2,iw]

    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 -α^2 - (2-α)/r1 - (2-α)/r2 + 1/r12
end

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



# initialization 
function init(R, alive, psi, E) 
    println("init")
    # 
    # Generate Nwx walkers
    #
    alive .= false
    psi .= 0.0
    E .= 0.0
    println("generating $Nw_target walkers")
    # set Nw_target walkers
    for iw in 1:Nw_target
        alive[iw] = true
        R[:,:,iw] = rand(D, N)        # coordinates
        psi[iw] = ln_psi2(R, iw)  # wave function (ln(psi^2))
        E[iw]   = Elocal(R, iw)        
    end
    vmcstat = VMCStats(3.1,0,0)
    println("init done")
    vmcstat
end

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

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

    return -2*α*(r1+r2)
end
@inline function ln_psi2(R)
    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*α*(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, iw, vmcstat::VMCStats)
    

    psi = ln_psi2(R,iw)
    E   = Elocal(R,iw)
    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,iw] += d1
        R[2,i,iw] += d2
        R[3,i,iw] += d3

        psi_new = ln_psi2(R,iw)
        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,iw)
        else
            R[1,i,iw] -= d1
            R[2,i,iw] -= d2
            R[3,i,iw] -= 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,iw,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,iw]," ")
            end
            println(f," ")
        end
    end
end

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


# Branching
function branch!(R::Array{Float64}, copies::Vector{Int}, alive::Vector{Bool}, psi::Vector{Float64}, E::Vector{Float64},
                 free_slots::Vector{Int}   # preallocated buffer
                 )
    nfree = 0

    # --- First pass: collect free slots ---
    @inbounds for i in eachindex(alive)
        if alive[i]
            if copies[i] == 0
                alive[i] = false
                nfree += 1
                free_slots[nfree] = i
            end
        else
            nfree += 1
            free_slots[nfree] = i
        end
    end

    nextfree = 1

    # --- Second pass: branching ---
    @inbounds for iw in eachindex(alive)
        if !alive[iw]
            continue
        end

        c = copies[iw]
        if c <= 1
            continue
        end
       
        for _ in 1:(c - 1)
            if nextfree > nfree
                error("No free walkers available; too large τ ?")
            end

            iw2 = free_slots[nextfree]
            nextfree += 1

            R[:,:,iw2]  .= R[:,:,iw] 
            E[iw2] = E[iw]
            alive[iw2] = true
        end
    end

    return nothing
end


# DMC drift, for ith electron F_i = 2\nabla_i \varphi_T /\varphi_T
@inline function Drift!(F, R, iw)
    x1, y1, z1 = R[1,1,iw], R[2,1,iw], R[3,1,iw]
    x2, y2, z2 = R[1,2,iw], R[2,2,iw], R[3,2,iw]

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

    c = -2*α

    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

@inline function Drift!(F, R)
    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*α

    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, iw, F, R1, R2, η)
    if order == 1
        randn!(η)
        Drift!(F, R, iw)
        @inbounds for j in 1:N, i in 1:D
            R[i,j,iw] += lambda*tau*F[i,j] + sqrt(2*lambda*tau)*η[i,j]
        end

    elseif order == 2
        randn!(η)
        @inbounds for j in 1:N, i in 1:D
            R1[i,j] = R[i,j,iw] + sqrt(lambda*tau)*η[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,iw] = R1[i,j] + lambda*tau*F[i,j]
        end

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

    return nothing
end


function lnG(Rnew, Rold, Fold)
    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 
    #
    R = Array{Float64}(undef, D, N, Nwx)
    alive = Array{Bool}(undef, Nwx)
    psi = Array{Float64}(undef, Nwx)
    E = Array{Float64}(undef, Nwx)
    
    
    # initialize
    vmcstat = init(R, alive, psi, E)   
    # thermalization
    println("thermalizing")
    Nw = Nw_target # for VMC
    @inbounds for i in 1:10
       @inbounds for iw in 1:Nw         
           _psi, _E = vmc_step!(R, iw, vmcstat)
           psi[iw] = _psi
           E[iw]   = _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_sum = 0.0 
        @inbounds for iw in 1:Nw         
            _psi, _E = vmc_step!(R, iw, vmcstat)
            psi[iw] = _psi
            E[iw]   = _E
            E_sum += _E 
        end
        E_ave += E_sum/Nw        
        nE +=1
        
        if ivmc%10 == 0            
            #println("VMC Elocal=$(E/Nw)  <E>=$(E_ave/nE)  <E>_analytical=$(Eexpec(α))")
            @printf("VMC  E = %.10f <E> = %.10f  <E>_analytical = %.10f\n",E_sum/Nw,E_ave/nE, Eexpec(α))            
        end
        adjust_step(vmcstat)
    end
    
    #
    # DMC
    #
    ET::Float64 = E_ave/nE  # trial energy, to be updated
    println("ET = $ET")
    
    Ntherm = floor(Int,1.0/tau) # start DMC measurement after excitations have died out; system specific !
    idmc::Int = -Ntherm

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

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

    
    idat::Int = 0 
    copies = zeros(Int,Nwx)
    Ntry::Int = 0 
    Nacc::Int = 0 

    
    Rold = zero(Coord)
    Rnew = zero(Coord)
    F    = zero(Coord)
    R1   = zero(Coord)
    R2   = zero(Coord)
    η  = zero(Coord)
    
    free_slots = Int[]
    sizehint!(free_slots, Nwx)
    
    E_sum::Float64 = 0.0  
    n::Int = 0

    while idmc<NDMC
        #
        # take a DMC step in each walker
        #
        
        E_sum = 0.0  
        n = 0
        
        @inbounds for iw in 1:Nwx
            copies[iw] = 0
            if alive[iw]==false; continue; end  
            ELold = E[iw]
            Rnew = R[:,:,iw] # alias
            @inbounds Rold .= Rnew
            #
            diffusion_drift_step!(R, iw, F, R1, R2, η)
            # 
            EL = Elocal(R, iw)
            if order==1
                # Metropolis accept/reject                
                Wnew = lnG(Rold, Rnew, F)+ln_psi2(Rnew)   # lnG(new->old) + lnpsi2(new)
                Wold = lnG(Rnew, Rold, F)+ln_psi2(Rold)   # lnG(old->new) + lnpsi2(old)
                accept = metro(Wold,Wnew)
            else
                accept = true
            end
            Ntry +=1
            if accept
                Nacc +=1
                E[iw] = EL
            else
                EL = ELold
                E[iw] = ELold
                Rnew .= Rold                
            end

            E_sum  += EL
            n  += 1
            if idat>0
                Edat  += EL
                E2dat += EL^2
                Ndat  +=1  
            end
            
            weight  = exp(-tau*(0.5*(ELold+EL)-ET))  # symmetric under R<->R'
            copies[iw] = floor(Int,weight + rand()) 
        end
        # Branching
        branch!(R, copies, alive, psi, E, free_slots)

        Nw = sum(alive)
        E_ave += E_sum/n # Don't use Nw as energy counter, it's not correct
        nE += 1

        idmc +=1
        
       
        
        if Nw==Nwx
            error("Hit max number of walkers")
        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_sum/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
            push!(Eb,Edat/Ndat)
            push!(E2b,E2dat/Ndat)
            Eave = mean(Eb)
            EStdError  =  0.0
            
            if size(Eb,1)>1
                EStdError = std(Eb)/sqrt(size(Eb,1))  # sigma/sqrt(N) for N data points
            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()