import numpy as np
import pylab as pl
from ase.data.molecules import molecule
from ase.parallel import world
from gpaw import GPAW
from gpaw.poisson import PoissonSolver
from gpaw.dipole_correction import DipoleCorrectionPoissonSolver

p1 = PoissonSolver()
p2 = DipoleCorrectionPoissonSolver(PoissonSolver(), 2)

calc1 = GPAW(mode='lcao', basis='sz', poissonsolver=p1)
calc2 = GPAW(mode='lcao', basis='sz', poissonsolver=p2)


system = molecule('H2O')
system.center(vacuum=4.0)
system.set_pbc([1, 1, 0])

for i, c in enumerate([calc1, calc2]):
    system.set_calculator(c)
    system.get_potential_energy()
    vt_g = c.hamiltonian.vt_sg.sum(0)#c.get_effective_potential()
    #pl.plot(vt_g.sum(0).sum(0), label='%d' % i)
#pl.legend()

pl.figure()
rhot_g = calc1.density.rhot_g
phit1_g = np.zeros_like(rhot_g)
phit2_g = np.zeros_like(rhot_g)

p1.solve(phit1_g, rhot_g)
p2.solve(phit2_g, rhot_g)

dphit_g = phit2_g - phit1_g

pl.plot(phit1_g.sum(0).sum(0), label='phit1')
pl.plot(phit2_g.sum(0).sum(0), label='phit2')
pl.plot(dphit_g.sum(0).sum(0), label='corr')
pl.legend()

pl.show()