import autograd.numpy as np

from pyCHAMP.wavefunction.wf_base import WF
from pyCHAMP.optimizer.minimize import Minimize
from pyCHAMP.sampler.metropolis import Metropolis
from pyCHAMP.sampler.hamiltonian import Hamiltonian
from pyCHAMP.solver.vmc import VMC


class Hydrogen(WF):

    def __init__(self, nelec, ndim):
        WF.__init__(self, nelec, ndim)

    def values(self, parameters, pos):
        """ Compute the value of the wave function.

        Args:
                parameters : parameters of th wf
                x: position of the electron

        Returns: values of psi
        """

        beta = parameters[0]
        if pos.ndim == 1:
            pos = pos.reshape(1, -1)

        r = np.sqrt(np.sum(pos**2, 1))
        return 2*np.exp(-beta*r).reshape(-1, 1)

    def nuclear_potential(self, pos):
        r = np.sqrt(np.sum(pos**2, 1))
        rm1 = - 1. / r
        return rm1.reshape(-1, 1)

    def electronic_potential(self, pos):
        return 0


if __name__ == "__main__":

    wf = Hydrogen(nelec=1, ndim=3)
    sampler = Metropolis(nwalkers=1000, nstep=1000, step_size=3,
                         nelec=1, ndim=3, domain={'min': -5, 'max': 5})
    sampler = Hamiltonian(nwalkers=1000, nstep=1000,
                          step_size=3, nelec=1, ndim=3)
    optimizer = Minimize(method='bfgs', maxiter=25, tol=1E-4)

    # VMS solver
    vmc = VMC(wf=wf, sampler=sampler, optimizer=optimizer)

    # single point
    opt_param = [1.]
    pos, e, s = vmc.single_point(opt_param)
    print('Energy   : ', e)
    print('Variance : ', s)
    vmc.plot_density(pos)

    # optimization
    init_param = [0.5]
    vmc.optimize(init_param)
    vmc.plot_history()

[((44, 14, 45, 71), 'pyCHAMP.sampler.metropolis.Metropolis', 'Metropolis', (), '', False, 'from pyCHAMP.sampler.metropolis import Metropolis\n'), ((46, 14, 47, 55), 'pyCHAMP.sampler.hamiltonian.Hamiltonian', 'Hamiltonian', (), '', False, 'from pyCHAMP.sampler.hamiltonian import Hamiltonian\n'), ((48, 16, 48, 61), 'pyCHAMP.optimizer.minimize.Minimize', 'Minimize', (), '', False, 'from pyCHAMP.optimizer.minimize import Minimize\n'), ((51, 10, 51, 58), 'pyCHAMP.solver.vmc.VMC', 'VMC', (), '', False, 'from pyCHAMP.solver.vmc import VMC\n'), ((13, 8, 13, 38), 'pyCHAMP.wavefunction.wf_base.WF.__init__', 'WF.__init__', ({(13, 20, 13, 24): 'self', (13, 26, 13, 31): 'nelec', (13, 33, 13, 37): 'ndim'}, {}), '(self, nelec, ndim)', False, 'from pyCHAMP.wavefunction.wf_base import WF\n'), ((29, 20, 29, 37), 'autograd.numpy.sum', 'np.sum', ({(29, 27, 29, 33): 'pos ** 2', (29, 35, 29, 36): '1'}, {}), '(pos ** 2, 1)', True, 'import autograd.numpy as np\n'), ((33, 20, 33, 37), 'autograd.numpy.sum', 'np.sum', ({(33, 27, 33, 33): 'pos ** 2', (33, 35, 33, 36): '1'}, {}), '(pos ** 2, 1)', True, 'import autograd.numpy as np\n'), ((30, 17, 30, 32), 'autograd.numpy.exp', 'np.exp', ({(30, 24, 30, 31): '(-beta * r)'}, {}), '(-beta * r)', True, 'import autograd.numpy as np\n')]