# propagators.jl
# Let's get those ψ moving.

"""
    SiteEnergyFromDipoles(dipoles,S,E)

Calculate site-energies for sites, from potential generated by dipoles at
other sites.

dipoles - dipole matrix (e*a_0)
S - previous site energies (Hartree)

returns:
S - Updated site energies for TightBinding Hamiltonian (Hartree)
"""
function SiteEnergyFromDipoles(dipoles,S,E)
    S = zeros(N)
    for i in 1:N
        S[i]+=-field_ext*(r*i)
        for j in 1:N
            if (j==i)
                # avoid infinity in self energies
                continue
            end
            # Contribution to site energy (1 e- at site) from dipoles
            S[i]+=sign(j-i)*dipoles[j]/(r*(i-j))^2
        end
    end
    S = S
    H = diagm(E,-1)+diagm(S)+diagm(E,1)
    return S,H
end

"""
    FieldFromDensity(density)

Calculate field at each dipole site from the charge density at every other site

density - electron density found from the wavefunction |Psi|^2

returns:
Field - electrostatic field
"""
function FieldFromDensity(density)
    field=zeros(N)
    for i in 1:N
        Fsum=zero(eltype(density)) # Calculate total field at each
        for j in 1:N
            if (j==i)
                continue # avoid infinite self energies
            end
            Fsum+=-sign(i-j)*density[j]/(i-j)^2 # How much do the dipoles respond to the electron density?
        end                         # check power of r for field from point charge.
        field[i]=Fsum/r^2
    end
    return field
end


"""
    FieldFromDipole(dipoles)

Calculate field at each dipole site from the induced dipole-induced dipole
interations at every other site

dipoles - existsing induced dipole

returns:
FieldDI - electrostatic field
"""
function FieldFromDipole(dipole)
    FieldDI=zeros(N)
    for i in 1:N
        FDIsum=zero(eltype(dipole)) # Calculate total field at each
        for j in 1:N
            if (j==i)
                continue # avoid infinite self energies
            end
            FDIsum+=-9*sign(j-i)*dipole[j]/((r*(j-i))^7) # How much do the dipoles respond to the electron density?
        end
        FieldDI[i]=FDIsum
    end
    return FieldDI
end

"""
    UpdateDipole(field, dipole, dampening, alpha)

Update dipole vector by solving the equation E = \alpha*p

field  - electrostatic field from electron density at every other site
dipole - current dipole vector to be updated
dampening - Contains the exponential decay of the dipole response (exp(-TimeStep*DecayRate))
alpha - polarisability of sites (assuming uniformity)

returns:
dipole - updates dipole moments induced by electrostatic fields
"""
function UpdateDipole(field, dipole, dampening, alpha=1.0)
    alphainv = 1*ones(N)/alpha#*dipolestrength
    # diagonal elements from polarisability of sites
    M = diagm(alphainv)
    n=0
    # Off-diagonal elements from separation between sites
    for i in 1:N
        for j in 1:N
            if (j==i)
                continue
            end
            M[n+j] = 1/(r*(j-i))^3
        end
        n+=N
    end
    new_dipole = \(M,field)
    # Update dipoles with a fraction from the new dipole matrix
    dipole += dampening*(new_dipole-dipole)
    return dipole
end

"Original JMF function; dipoles by exponential relaxation to t=+∞ solution."
function dipoles_by_relaxation(dipoles,density,dampening)
    for i in 1:N
        relaxeddipole=zero(eltype(density)) # fix type instability; initialise to zero in the supplied type of density
        for j in 1:N
            if (j==i)
                continue # avoid infinite self energies
            end
            relaxeddipole+=density[j]/(i-j)^1 # How much do the dipoles respond to the electron density?
        end

        dipoles[i]+=dampening*(relaxeddipole-dipoles[i]) # Approaches the infinite time limit via Zeno's dichotomy
    end
    return dipoles
end

"""
    AdiabaticPropagation(S,dipoles,E)

Build Hamiltonian from Dipoles (via SiteEnergyFromDipoles), diagonalise and
update dipoles responding to ground state electron density. A wrapper
function.

Time is ill-defined in this simulation - it takes you towards equilibrium.

S       - (Hartree)
dipoles -
E       - (Hartree)

returns
S       -
psi     -
Density -
dipoles -
"""
function AdiabaticPropagation(S,dipoles,density,E,dampening,verbose::Bool=true,dipole_int::Bool=false)
    S,H = SiteEnergyFromDipoles(dipoles,S,E)
    psi=eigvecs(H)[:,1] # gnd state

    if verbose
        println("Adiabtic State energy: <psi|H|psi> = ",psi'*H*psi)
        KE=diagm(E,-1)+diagm(E,+1)
        PE=diagm(S)
        println("KE: <psi|KE|psi> = ",psi'*KE*psi)
        println("PE: <psi|PE|psi> = ",psi'*PE*psi)
    end

    density=psi.^2
    Field=FieldFromDensity(density)+field_ext
    dipoles = UpdateDipole(Field, dipoles, dampening)

    if dipole_int
        FieldDI = FieldFromDipole(dipoles)
        dipoles = UpdateDipole(Field+FieldDI, dipoles, dampening)
    end

    return S,H,psi,density,dipoles
end

"""
    UnitaryPropagation(dipoles,S,E,psi,dt,dampening; slices::Int=1)

Self-consistent response of the lattice with unitary (time dependent)
evolution of the wavefunction.
"""
function UnitaryPropagation(dipoles,density,S,E,psi,dt,dampening; slices::Int=1)
    S,H=SiteEnergyFromDipoles(dipoles,S,E)

    H=diagm(E,-1)+diagm(S)+diagm(E,+1)

    En=eigvals(H)[1]
    println("First Eigval (adiabatic): ",En)

#    psi=eigvecs(H)[:,1] # gnd state; reproduces adiabtic state (eigval) above
    println("State energy: <psi|H|psi> = ",psi'*H*psi)
    KE=diagm(E,-1)+diagm(E,+1)
    PE=diagm(S)
    println("KE: <psi|KE|psi> = ",psi'*KE*psi)
    println("PE: <psi|PE|psi> = ",psi'*PE*psi)

    psi=TimeDependentPropagation(psi,H,dt,slices=slices,decompose=false,verbose=false)

    density=abs(psi.^2) # can be Complex!

    Field=FieldFromDensity(density)+field_ext
    dipoles = UpdateDipole(Field, dipoles, dipolestrength)

    return S,H,psi,density,dipoles
end

"""
    TimeDependentPropagation(psi,H,dt;slices::Int=1,decompose::Bool=false,verbose::Bool=false)

Propagate Wavefunction directly with Hamiltonian and time dependent Schrodinger equation.

Psi (N) is the wavefunction; H the Hamiltonian (NxN); dt the length of time to
propgate along; slices is how many slices to decompose the Unitary operator
into; verbose sets the display of info on the unitary nature of U, whether U*U'
~= I.

psi         - wavefunction
H           - TightBinding Hamiltonian (Hartree)
dt          - Time Step (0.02418884326505 fs)
slices      - sub-time step (0.02418884326505 fs)
decompose   - Perform Trotter Expansion (Boolean)
verbose     - Print unitary transformation (Boolean)

returns:
psi         - wavefunction
"""
function TimeDependentPropagation(psi,H,dt;slices::Int=1,decompose::Bool=false,verbose::Bool=false,test::Bool=false) # propagate directly using full Hamiltonian=T+V
    dt=dt/slices
    U=eye(H) # identiy matrix same size + type as H

    for i=1:slices
        if decompose
            S,J=Decompose_H(H) # split into diagonal and off-diag terms
            # Trotter decomposition
            U*=expm(-im*J*dt/2)*expm(-im*S*dt)*expm(-im*J*dt/2)
        else
            U*=expm(-im*H*dt)
        end
    end

    psi=U*psi # OK; we've built out unitary time evolution operator U, now apply it

    if verbose # for debugging / introspection
        @printf(" (matrix-squaring slices: %d) U: \n",slices)
        display(U)
        println("\n\tUU': (in any sane world this should be = Identity )\n")
        display(U*U') # Why you no unitary?

        println("\nPre normalised Norm of psi: ",norm(psi))
    end

    normalisation=sqrt(sum(abs.(psi.^2))) # Nb: for normalisation of WAVEFUNCTION; must SQUAREROOT the sum of density
    #println("Normalising Psi by dividing by: ",normalisation)
    psi=psi/normalisation # Normalise propagated wavefunction

    if test
        psi = [psi, U]
    end

    return psi
end

"""
    TimeDependentPropagation(psi,H,dt,E)

Propagate (in time) using eigenvalue

psi     - wavefunction
H       - TightBinding Hamiltonian (Hartree)
dt      - Time step (0.02418884326505 fs)
E       - Energy Eigenvalues (Hartree)

returns:
psi     - wavefunction
"""
function TimeDependentPropagation(psi,H,dt,E)
    psi=exp(-im*E*dt)*psi
    println("Pre normalised Norm of psi: ",norm(psi))
#    psi/=norm(psi*psi') # Normalise propagated wavefunction
    return psi
end
