from math import *

import numpy as np
import numpy.fft

import visualize


# Normalize wavefunction with respect to volume element dV
def normalize(state, dV):
    norm2 = np.vdot(state, state) * dV
    return state / sqrt(abs(norm2))


# Gram-Schmidt orthonormalization with respect to dV
def orthonormalize(basis, dV):
    for n in range(len(basis)):
        psi = basis[n]

        for i in range(n):
            psi -= basis[i] * np.vdot(basis[i], psi) * dV

        basis[n] = normalize(psi, dV)

    return basis


# Zero all components with kinetic energy less than Ecut
def enforce_cutoff(s, psis_q, kq2):
    for i in range(s.Nres):
        for j in range(s.Nres):
            for k in range(s.Nres):
                if kq2[i, j, k] > 2 * s.Ecut:
                    for psi_q in psis_q:
                        psi_q[i, j, k] = 0
    return orthonormalize(psis_q, s.dq)


# Apply the Kohn-Sham Hamiltonian to a normalized state psi
def apply_hamiltonian(s, psi_q, Veff_r, kq2):
    # Apply kinetic energy in q-domain
    Tpsi_q = 0.5 * kq2 * psi_q

    # Apply local+nonlocal potential in r-domain
    psi_r = np.fft.ifftn(psi_q)
    # We don't normalize: the iFFT ensures a correctly normalized result
    Vpsi_r = Veff_r * psi_r # local part
    for a in s.atoms: # nonlocal part
        Vnon_r = a.apply_nonlocal_potential(psi_r) * s.dr
        Vpsi_r += Vnon_r
    Vpsi_q = np.fft.fftn(Vpsi_r)

    return Tpsi_q + Vpsi_q


# Resisual vector R, which the eigensolver tries to minimize
def get_residual(s, psi_q, Hpsi_q, kq2):
    ene = np.vdot(psi_q, Hpsi_q) * s.dq # assuming normalization
    R_q = Hpsi_q - ene * psi_q

    # Calculate preconditioned residual PR for faster convergence:
    # The ideal PR satisfies PR = (H - ene)^-1 R, such that (H - ene) PR = R.
    # This operator (H - ene) is hard to invert, so we approximate it using
    # the kinetic energy only: -0.5 \nabla^2 PR = R (the Poisson equation).
    PR_q = 2 * R_q / (kq2 + 1)
    # This is an even rougher approximation, but it works well enough.

    return PR_q


# One iteration of the block Davidson eigensolver algorithm
def diagonalize_subspace(s, psis, V, kq2):
    # Apply Hamiltonian to psis and cache results
    Hbasis = []
    for n in range(len(psis)):
        Hbasis.append(apply_hamiltonian(s, psis[n], V, kq2))

    # Calculate corresponding residuals
    resids = []
    for n in range(len(psis)):
        R = get_residual(s, psis[n], Hbasis[n], kq2)
        resids.append(R)

    # Augment the input basis (psis) with the residuals
    basis = psis + resids
    basis = enforce_cutoff(s, basis, kq2)

    Ndim  = len(basis)

    # Apply Hamiltonian to orthonormalized residuals and cache results
    for n in range(len(psis), Ndim):
        Hbasis.append(apply_hamiltonian(s, basis[n], V, kq2))

    # Subspace Hamiltonian matrix with elements <basis[i] | H | basis[j]>
    Hsub = np.zeros((Ndim, Ndim), dtype="complex128")
    # Exploit Hermiticity: only calculate the diagonal and upper triangle,
    # then copy the elements from the upper to the lower triangle.
    for r in range(Ndim):
        for c in range(r, Ndim):
            Hsub[r, c] = np.vdot(basis[r], Hbasis[c]) * s.dq
    for r in range(Ndim):
        for c in range(0, r):
            Hsub[r, c] = np.conj(Hsub[c, r])

    # Find the eigenvectors and eigenvalues of Hsub
    evals, evecs = np.linalg.eigh(Hsub)

    # Express the eigenvectors in the original basis (input psis)
    result = []
    for n in range(Ndim // 2):
        psi = np.zeros((s.Nres, s.Nres, s.Nres), dtype="complex128")
        for b in range(Ndim):
            psi += basis[b] * evecs[b, n]
        result.append(psi)

    return evals, result


# Small optimization: cache |k+q|^2 ahead of time
def cache_kinetic_factor(s, kpoint):
    kq2 = np.zeros((s.Nres, s.Nres, s.Nres))
    for i in range(s.Nres):
        for j in range(s.Nres):
            for k in range(s.Nres):
                q = np.array((s.qx[i, j, k], s.qy[i, j, k], s.qz[i, j, k]))
                kq = kpoint + q
                kq2[i, j, k] = np.dot(kq, kq)
    return kq2


# Initial guess from atoms' individual (approximate) orbitals
def guess_orbitals(s):
    psis_q = []
    for a in s.atoms:
        psis_r = a.guess_orbitals(s.Norb)

        # Convert guesses to Fourier space
        for psi_r in psis_r:
            psi_q = np.fft.fftn(psi_r)
            psis_q.append(psi_q)

    return orthonormalize(psis_q, s.dq)


# Get lowest eigenvalues and eigenstates of Kohn-Sham equation
def solve_kohnsham_equation(s, V, k, prefix="", thresh=1e-4):
    kq2 = cache_kinetic_factor(s, k)

    # Initialize psis, which we will refine shortly
    psis_q = guess_orbitals(s)

    print(prefix + "Eigensolver cycle 1 running...")
    cycles = 0
    improv = 1
    energy_old = 0
    # Iterative eigensolver loop
    while improv > thresh:
        cycles += 1

        # Eliminate high-frequency noise (little effect in practice)
        psis_q = enforce_cutoff(s, psis_q, kq2)

        # Refine eigenstates using one step of the block Davidson method
        energies, psis_q = diagonalize_subspace(s, psis_q, V, kq2)

        # To measure progress, we look at the change of (N+1)th eigenvalue,
        # where N is the number of bands we are interested in converging.
        energy_new = energies[s.Neig]
        improv = abs(energy_new - energy_old) / abs(energy_new)
        energy_old = energy_new

        # Clear line and (re)print progress information
        print("\033[A\033[2K" + prefix, end="")
        print("Eigensolver cycle {} done".format(cycles), end="")
        if cycles > 1:
            print(", improvement {:.2e}".format(improv), end="")
        print("")

    # Convert psis to real space and normalize them there
    psis_r = []
    for n in range(len(psis_q)):
        psi_r = np.fft.ifftn(psis_q[n])
        psis_r.append(normalize(psi_r, s.dr))

    # energies is a NumPy array right now
    return energies.tolist(), psis_r, psis_q