from math import sqrt, pi

import numpy as np
from ase.units import Bohr

from gpaw.utilities import pack
from gpaw.analyse.hirshfeld import HirshfeldDensity
from gpaw.utilities.tools import coordinates


def wignerseitz(gd, atoms, scale=None):
    """Determine which atom is closest to each grid point.

    The atom distances might be scaled by the scale list."""
    if scale is None:
        scale = [1.] * len(atoms)
    else:
        assert len(scale) == len(atoms)
    r_vG, R2min_G = coordinates(gd, atoms[0].position / Bohr)
    R2min_G *= scale[0] ** 2
    index_G = gd.zeros(dtype=int)
    for a, atom in enumerate(atoms[1:]):
        r_vG, r2_G = coordinates(gd, atom.position / Bohr)
        r2_G *= scale[a + 1] ** 2
        index_G = np.where(R2min_G > r2_G, a + 1, index_G)
        R2min_G = np.where(R2min_G > r2_G, r2_G, R2min_G)
    return index_G


class WignerSeitz:

    def __init__(self, gd, atoms,
                 calculator=None, scale=None):
        """Find the grid points nearest to the atoms"""

        self.atoms = atoms
        self.gd = gd
        self.calculator = calculator

        self.atom_index = wignerseitz(gd, atoms, scale)

    def expand(self, density):
        """Expand a smooth density in Wigner-Seitz cells around the atoms"""
        n = len(self.atoms)
        weights = np.empty((n,))
        for a in range(n):
            mask = np.where(self.atom_index == a, density, 0.0)
            # XXX Optimize! No need to integrate in zero-region
            weights[a] = self.gd.integrate(mask)

        return weights

    def expand_density(self, nt_G, s, nspins):
        """Get the weights of spin-density in Wigner-Seitz cells
        around the atoms. The spin index and number of spins are
        needed for the augmentation sphere corrections."""
        weights_a = self.expand(nt_G)
        for a, nucleus in enumerate(self.atoms):
            weights_a[a] += nucleus.get_density_correction(s, nspins)
        return weights_a

    def expand_wave_function(self, psit_G, u, n):
        """Get the weights of wave function in Wigner-Seitz cells
        around the atoms. The spin-k-point index u and band number n
        are needed for the augmentation sphere corrections."""

        assert psit_G.dtype == float
        # smooth part
        weigths = self.expand(psit_G ** 2)

        # add augmentation sphere corrections
        for a, nucleus in enumerate(self.atoms):
            P_i = nucleus.P_uni[u, n]
            P_p = pack(np.outer(P_i, P_i))
            Delta_p = sqrt(4 * pi) * nucleus.setup.Delta_pL[:, 0]
            weigths[a] += np.dot(Delta_p, P_p)

        return weigths

    def get_charges(self, den_g):
        """Charge on the atom according to the Wigner-Seitz partitioning

        Can be applied to any density den_g.
        """
        assert den_g.shape == tuple(self.gd.n_c)
        charges = []
        for atom, q in zip(self.atoms, self.expand(den_g)):
            charges.append(atom.number - q)
        return charges

    def get_effective_volume_ratio(self, atom_index):
        """Effective volume to free volume ratio.

        After: Tkatchenko and Scheffler PRL 102 (2009) 073005
        """
        atoms = self.atoms
        finegd = self.gd

        den_g, gd = self.calculator.density.get_all_electron_density(atoms)
        assert gd == finegd
        denfree_g, gd = self.hdensity.get_density([atom_index])
        assert gd == finegd

        # the atoms r^3 grid
        position = self.atoms[atom_index].position / Bohr
        r_vg, r2_g = coordinates(finegd, origin=position)
        r3_g = r2_g * np.sqrt(r2_g)

        weight_g = np.where(self.atom_index == atom_index, 1.0, 0.0)

        nom = finegd.integrate(r3_g * den_g[0] * weight_g)
        denom = finegd.integrate(r3_g * denfree_g)

        return nom / denom

    def get_effective_volume_ratios(self):
        """Return the list of effective volume to free volume ratios."""
        ratios = []
        self.hdensity = HirshfeldDensity(self.calculator)
        for a, atom in enumerate(self.atoms):
            ratios.append(self.get_effective_volume_ratio(a))
        return np.array(ratios)


class LDOSbyBand:

    """Base class for a band by band LDOS"""

    def by_element(self):
        # get element indicees
        elemi = {}
        for i, nucleus in enumerate(self.paw.atoms):
            symbol = nucleus.setup.symbol
            if symbol in elemi:
                elemi[symbol].append(i)
            else:
                elemi[symbol] = [i]
        for key in elemi.keys():
            elemi[key] = self.get(elemi[key])
        return elemi


'''
class WignerSeitzLDOS(LDOSbyBand):

    """Class to get the unfolded LDOS defined by Wigner-Seitz cells"""

    def __init__(self, paw):
        self.paw = paw
        self.ws = WignerSeitz(paw.gd, paw.atoms)

        nu = paw.nkpts * paw.nspins
        ldos = np.empty((nu, paw.nbands, len(paw.atoms)))
        for u, kpt in enumerate(paw.kpt_u):
            for n, psit_G in enumerate(kpt.psit_nG):
                # XXX variable ws undefined.  Probably self.ws
                ldos[u, n, :] = ws.expand_wave_function(psit_G, u, n)

    def write(self, filename, slavewrite=False):
        if self.world.rank == MASTER or slavewrite:
            paw = self.paw
            f = open(filename, 'w')

            nn = len(paw.atoms)

            for k in range(paw.nkpts):
                for s in range(paw.nspins):
                    u = s * paw.nkpts + k
                    for n in range(paw.nbands):
                        # avery: Added dummy loop body to make compiling work.
                        1
'''
