from math import *
from scipy.special import erf, sph_harm, eval_genlaguerre

import numpy as np

import element


# Atomic Coulomb potential, modelled using HGH pseudopotentials.
# See https://arxiv.org/abs/cond-mat/9803286 for details.
class Atom:
    def __init__(self, s, symbol, pos):
        self.symbol = symbol
        self.params = element.Element(symbol) # get element data
        self.position = pos

        # Set up local grid in spherical coordinates
        self.r     = np.zeros((s.Nres, s.Nres, s.Nres))
        self.theta = np.zeros((s.Nres, s.Nres, s.Nres))
        self.phi   = 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):
                    x = s.rx[i, j, k]
                    y = s.ry[i, j, k]
                    z = s.rz[i, j, k]

                    r = sqrt((x - pos[0])**2 + (y - pos[1])**2 + (z - pos[2])**2)
                    if r == 0:
                        theta = 0
                        phi   = 0
                    else:
                        theta = acos((z - pos[2]) / r)
                        phi   = atan2((y - pos[1]), (x - pos[0]))

                    self.r[i, j, k]     = r
                    self.theta[i, j, k] = theta
                    self.phi[i, j, k]   = phi

        # Initialize and cache orbitals for nonlocal potential contribution
        self.orbitals = {}
        for l in range(0, 3):
            self.orbitals[l] = {}
            for i in range(1, 4):
                p = self.get_projector(l, i)
                self.orbitals[l][i] = {}
                for m in range(-l, l + 1):
                    Y = self.cubic_harmonic(m, l)
                    self.orbitals[l][i][m] = p * Y


    # Local potential representing short-range interactions.
    # This is Eq.(1), but without the long-range erf-term:
    def get_local_potential(self):
        V = np.zeros_like(self.r)

        rr = self.r / self.params.rloc # relative r
        for i in range(4):
            V += self.params.Cloc[i] * rr**(2 * i) * np.exp(-0.5 * rr**2)

        return V


    # The local erf-term decays as 1/r, which is problematic for periodicity.
    # Instead, we replace it with a Hartree potential from a Gaussian charge.
    # This is exact: multiply first term of Eq.(5) by g^2 (from Poisson).
    def get_local_density(self):
        alpha = 0.5 / self.params.rloc**2
        rho = 4 * alpha**(3/2) * np.exp(-alpha * self.r**2) / sqrt(pi)

        # Normalize in spherical coordinates: 4pi int rho(r) r^2 dr = Zion
        rho *= self.params.Zion / (4 * pi)
        return rho


    # Given a psi in real space, apply nonlocal potential by integration
    # Vpsi = sum_{lijm} |p_i> h_{ij} <p_j | psi>
    def apply_nonlocal_potential(self, psi):
        Vpsi = np.zeros_like(psi)

        # Note that these if-statements provide a significant speedup
        for l in range(0, 3):
            if self.params.rl[l] != 0:
                for i in range(1, 4):
                    for j in range(1, 4):
                        if self.params.h[l][i][j] != 0:
                            for m in range(-l, l + 1):
                                # Inner product normalization is handled by caller
                                ip = np.vdot(self.orbitals[l][j][m], psi)
                                Vpsi += self.orbitals[l][i][m] * self.params.h[l][i][j] * ip

        return Vpsi


    # Calculate projectors p_i^l(r) for nonlocal potential, see Eq.(3)
    def get_projector(self, l, i):
        p = np.zeros_like(self.r)

        if self.params.rl[l] != 0:
            p =  sqrt(2) * self.r**(l + 2 * (i - 1))
            p *= np.exp(-0.5 * (self.r / self.params.rl[l])**2)
            p /= self.params.rl[l]**(l + 0.5 * (4 * i - 1)) 
            p /= sqrt(gamma(l + 0.5 * (4 * i - 1)))

        return p


    # Guess this atom's electron density (to initialize main DFT loop)
    def guess_density(self):
        # We choose a spherical gaussian with standard deviation sigma
        sigma = 1
        rho = (sqrt(2 / pi) / sigma**3) * np.exp(- 0.5 * self.r**2 / sigma**2)

        # Normalize in spherical coordinates: 4pi int rho(r) r^2 dr = Zion
        rho /= self.params.Zion / (4 * pi)
        return rho


    # Guess this atom's eigenstates (to initialize eigensolver).
    # We simply choose hydrogen orbitals, without any modifications:
    def guess_orbitals(self, Norb):
        psis = []
        for n in range(1, 4):
            rho = 2 * self.r / n
            for l in range(0, n):
                for m in range(-l, l + 1):
                    psi = self.cubic_harmonic(m, l)
                    psi *= np.exp(-0.5 * rho) * rho**l
                    psi *= eval_genlaguerre(n - l - 1, 2 * l + 1, rho)
                    psis.append(psi)

        return psis[0 : Norb]


    # Use cubic harmonics instead of spherical harmonics. This makes no
    # physical difference, but it avoids confusion when visualizing.
    def cubic_harmonic(self, m, l):
        # Note that phi and theta are swapped according to SciPy
        Yp = sph_harm( abs(m), l, self.phi, self.theta)
        Yn = sph_harm(-abs(m), l, self.phi, self.theta)

        f = np.zeros_like(self.r)
        if m == -3:
            f = (1  / sqrt(2)) * (Yn - Yp)
        if m == -2:
            f = (1j / sqrt(2)) * (Yn - Yp)
        if m == -1:
            f = (1  / sqrt(2)) * (Yn - Yp)
        if m == 0:
            f = Yp
        if m == 1:
            f = (1j / sqrt(2)) * (Yn + Yp)
        if m == 2:
            f = (1  / sqrt(2)) * (Yn + Yp)
        if m == 3:
            f = (1j / sqrt(2)) * (Yn + Yp)

        return np.real(f)