from math import pi, sin, cos
import numpy as np

from ase.units import Bohr
import gpaw.cgpaw as cgpaw
from gpaw.poisson import PoissonSolver
from gpaw.pes.state import State
from gpaw.analyse.expandyl import AngularIntegral


class PlaneWave(State):

    """Plane wave state."""

    def get_grid(self, k_c, r0=None):
        if r0 is None:
            r0_c = np.array([0., 0., 0.])
        else:
            r0_c = np.array(r0)
        pw_G = self.gd.zeros(dtype=complex)
        python = True
        python = False
        if python:
            r_c = np.array([0., 0., 0.])
            gd = self.gd
            for i in range(gd.n_c[0]):
                r_c[0] = (gd.beg_c[0] + i) * gd.h_c[0]
                for j in range(gd.n_c[1]):
                    r_c[1] = (gd.beg_c[1] + j) * gd.h_c[1]
                    for k in range(gd.n_c[2]):
                        r_c[2] = (gd.beg_c[2] + k) * gd.h_c[2]
                        phase = np.dot(k_c, r_c - r0_c)
                        pw_G[i, j, k] = cos(phase) + 1j * sin(phase)
        else:
            gd = self.gd
            cgpaw.plane_wave_grid(gd.beg_c, gd.end_c,
                                  gd.h_cv.diagonal().copy(),
                                  k_c, r0_c, pw_G)
        pw_G /= (2 * pi) ** (3. / 2.)
        return pw_G


class ZerothOrder1(State):

    def __init__(self, calculator):
        self.calculator = calculator
        State.__init__(self, calculator.gd)
        self.pw = PlaneWave(self.gd)

        # get effective potential
        hamiltonian = self.calculator.hamiltonian
        if 1:
            self.vt_G = hamiltonian.vt_sG[0]  # XXX treat spin
        else:
            self.vt_G = np.where(hamiltonian.vt_sG[0] > 0,
                                 0.0, hamiltonian.vt_sG[0])
        self.intvt = self.gd.integrate(self.vt_G)
        print('# int(vt_G)=', self.intvt, (self.vt_G > 0).any())

        self.solve()

    def get_grid(self, k_c, r0):
        print('Correction:', self.gd.integrate(self.corrt_G))
        if not hasattr(self, 'written'):
            print('r0', r0, r0 * Bohr)
            dR = 0.3
            AI = AngularIntegral(r0 * Bohr, self.gd, dR=dR)
            r_R = AI.radii()
            psi_R = AI.average(self.corrt_G)
            v_R = AI.average(self.vt_G)
            with open('radial_dR' + str(dR) + '.dat', 'w') as fd:
                print('# R  v(R)    psi(R)', file=fd)
                for r, psi, v in zip(r_R, psi_R, v_R):
                    print(r, psi, v, file=fd)

            self.written = True

        return self.pw.get_grid(k_c, r0) - 1e6 * self.corrt_G

    def solve(self):
        hamiltonian = self.calculator.hamiltonian
        self.poisson = PoissonSolver('fd', nn=hamiltonian.poisson.nn)
        self.poisson.set_grid_descriptor(self.gd)
        self.poisson.initialize()

        corrt_G = self.gd.empty()
        self.poisson.solve(corrt_G, self.vt_G, charge=None)
        corrt_G /= (2 * pi) ** (5. / 2)

        self.corrt_G = corrt_G
