import numpy as np
from ase.data.molecules import molecule
from ase.units import Bohr
from gpaw import GPAW
from gpaw.poisson import PoissonSolver
from gpaw.mpi import rank
from gpaw.spline import Spline
from gpaw.atom.basis import QuasiGaussian
from gpaw.lfc import NewLocalizedFunctionsCollection as LFC

def testcalculation():
    system = molecule('H2O')
    system.center(vacuum=5.0)
    calc = GPAW(mode='lcao', basis='dzp',
                poissonsolver=PoissonSolver(relax='GS', eps=1e-8))

    system.set_calculator(calc)
    system.get_potential_energy()
    return calc
    
def get_all_electron_charge_distributions(calc, ref=4):
    atoms = calc.get_atoms()
    density = calc.density

    n, gd = density.new_get_all_electron_density(atoms, gridrefinement=ref)
    assert n.shape[0] == 1, 'No spin polarization implemented'
    n = n[0]

    print 'Electron charge', gd.integrate(n)
    
    rchar = 0.25
    rcut = 1.7 * rchar
    r = np.linspace(0.0, rcut, 256)
    
    gaussian = QuasiGaussian(1.0 / rchar**2, rcut)
    deltafunc = gaussian(r)
    
    delta = Spline(0, r[-1], deltafunc)
    spos_ac = atoms.get_scaled_positions() % 1.0
    setups = calc.wfs.setups

    lfc = LFC(gd,
              [[delta] for _ in setups],
              integral=[float(setup.Z) for setup in setups])
    lfc.set_positions(spos_ac)
    Z = gd.zeros()
    lfc.add(Z)
    print 'Nuclear charge', gd.integrate(Z)
    
    rho = Z - n
    print 'Total charge', gd.integrate(rho)
    return n, Z, rho, gd

def get_all_electron_electrostatic_potential(rho, gd):
    poisson = PoissonSolver(relax='GS')
    poisson.set_grid_descriptor(gd)
    poisson.initialize()

    phi = gd.zeros()
    poisson.solve(phi, rho)
    return phi

def plot(arrays, titles, layout, gd):
    import pylab as pl
    for i, (arr, title) in enumerate(zip(arrays, titles)):
        arr = gd.collect(arr)

        if rank == 0:
            xsum = arr.sum(axis=0)
            pl.subplot(layout[0],layout[1],i + 1)
            pl.ylabel(title)
            pl.imshow(xsum)
    if rank == 0:
        pl.show()

def main():
    calc = testcalculation()
    n, Z, rho, gd = get_all_electron_charge_distributions(calc)
    phi = get_all_electron_electrostatic_potential(rho, gd)
    plot([n, Z, rho, phi], ['n', 'Z', 'rho', 'phi'], (2, 2), gd)

if __name__ == '__main__':
    main()