import warnings
from math import sqrt, pi
import numpy as np

from ase.units import Ha
from gpaw.mpi import world
from gpaw.density import redistribute_array, redistribute_atomic_matrices
from gpaw.sphere.lebedev import weight_n
from gpaw.utilities import pack, pack_atomic_matrices, unpack_atomic_matrices
from gpaw.xc.gllb import safe_sqr
from gpaw.xc.gllb.contribution import Contribution

# XXX Work in process
debug = False


def d(*args):
    if debug:
        print(args)


class ResponsePotential:
    """Container for response potential"""
    def __init__(self, response, vt_sG, vt_sg, D_asp, Dresp_asp):
        self.response = response
        self.vt_sG = vt_sG
        self.vt_sg = vt_sg
        self.D_asp = D_asp
        self.Dresp_asp = Dresp_asp

    def redistribute(self, new_response):
        old_response = self.response
        new_wfs = new_response.wfs
        new_density = new_response.density
        self.vt_sG = redistribute_array(self.vt_sG,
                                        old_response.density.gd,
                                        new_density.gd,
                                        new_wfs.nspins,
                                        new_wfs.kptband_comm)
        if self.vt_sg is not None:
            self.vt_sg = redistribute_array(self.vt_sg,
                                            old_response.density.finegd,
                                            new_density.finegd,
                                            new_wfs.nspins,
                                            new_wfs.kptband_comm)

        def redist_asp(D_asp):
            return redistribute_atomic_matrices(D_asp,
                                                new_density.gd,
                                                new_wfs.nspins,
                                                new_density.setups,
                                                new_density.atom_partition,
                                                new_wfs.kptband_comm)

        self.D_asp = redist_asp(self.D_asp)
        self.Dresp_asp = redist_asp(self.Dresp_asp)
        self.response = new_response


class C_Response(Contribution):
    def __init__(self, weight, coefficients, damp=1e-10):
        Contribution.__init__(self, weight)
        d('In c_Response __init__', self)
        self.coefficients = coefficients
        self.vt_sg = None
        self.vt_sG = None
        self.D_asp = None
        self.Dresp_asp = None
        self.Ddist_asp = None
        self.Drespdist_asp = None
        self.damp = damp
        self.fix_potential = False

    def set_damp(self, damp):
        self.damp = damp

    def get_name(self):
        return 'RESPONSE'

    def get_desc(self):
        return self.coefficients.get_description()

    def initialize(self, density, hamiltonian, wfs):
        Contribution.initialize(self, density, hamiltonian, wfs)
        self.coefficients.initialize(wfs)
        if self.Dresp_asp is None:
            assert self.density.D_asp is None
        # XXX With the deprecated `fixdensity=True` option this
        # initialize() function is called both before AND after read()!
        # Thus, the second call of initialize() would discard the read
        # potential unless we check if it was already allocated.
        # Remove this if statement once `fixdensity=True` option has
        # been removed from the calculator.
        if self.vt_sg is None:
            self.vt_sG = self.gd.empty(self.nspins)
            self.vt_sg = self.finegd.empty(self.nspins)

    def initialize_1d(self, ae):
        Contribution.initialize_1d(self, ae)
        self.coefficients.initialize_1d(ae)

    def initialize_from_other_response(self, response):
        pot = ResponsePotential(response,
                                response.vt_sG,
                                response.vt_sg,
                                response.D_asp,
                                response.Dresp_asp)
        pot.redistribute(self)
        self.vt_sG = pot.vt_sG
        self.vt_sg = pot.vt_sg
        self.D_asp = pot.D_asp
        self.Dresp_asp = pot.Dresp_asp
        self.Ddist_asp = self.distribute_D_asp(pot.D_asp)
        self.Drespdist_asp = self.distribute_D_asp(pot.Dresp_asp)

    # Calcualte the GLLB potential and energy 1d
    def add_xc_potential_and_energy_1d(self, v_g):
        w_i = self.coefficients.get_coefficients_1d()
        u2_j = safe_sqr(self.ae.u_j)
        v_g += self.weight * np.dot(w_i, u2_j) / (
            np.dot(self.ae.f_j, u2_j) + self.damp)
        return 0.0  # Response part does not contribute to energy

    def update_potentials(self):
        d('In update response potential')
        if self.fix_potential:
            # Skip the evaluation of the potential so that
            # the existing potential is used. This is relevant for
            # band structure calculation so that the potential
            # does not get updated with the other k-point sampling.
            return

        if self.wfs.kpt_u[0].eps_n is None:
            # This skips the evaluation of the potential so that
            # the existing one is used.
            # This happens typically after reading when occupations
            # haven't been calculated yet and the potential read earlier
            # is used.
            return

        w_kn = self.coefficients.get_coefficients(self.wfs.kpt_u)
        f_kn = [kpt.f_n for kpt in self.wfs.kpt_u]
        if w_kn is not None:
            self.vt_sG[:] = 0.0
            nt_sG = self.gd.zeros(self.nspins)

            for kpt, w_n in zip(self.wfs.kpt_u, w_kn):
                self.wfs.add_to_density_from_k_point_with_occupation(
                    self.vt_sG, kpt, w_n)
                self.wfs.add_to_density_from_k_point(nt_sG, kpt)

            self.wfs.kptband_comm.sum(nt_sG)
            self.wfs.kptband_comm.sum(self.vt_sG)

            if self.wfs.kd.symmetry:
                for nt_G, vt_G in zip(nt_sG, self.vt_sG):
                    self.wfs.kd.symmetry.symmetrize(nt_G, self.gd)
                    self.wfs.kd.symmetry.symmetrize(vt_G, self.gd)

            d('response update D_asp', world.rank, self.Dresp_asp.keys(),
              self.D_asp.keys())
            self.wfs.calculate_atomic_density_matrices_with_occupation(
                self.Dresp_asp, w_kn)
            self.Drespdist_asp = self.distribute_D_asp(self.Dresp_asp)
            d('response update Drespdist_asp', world.rank,
              self.Dresp_asp.keys(), self.D_asp.keys())
            self.wfs.calculate_atomic_density_matrices_with_occupation(
                self.D_asp, f_kn)
            self.Ddist_asp = self.distribute_D_asp(self.D_asp)

            self.vt_sG /= nt_sG + self.damp

        self.density.distribute_and_interpolate(self.vt_sG, self.vt_sg)

    def calculate(self, e_g, n_sg, v_sg):
        self.update_potentials()
        v_sg += self.weight * self.vt_sg

    def distribute_D_asp(self, D_asp):
        return self.density.atomdist.to_work(D_asp)

    def calculate_energy_and_derivatives(self, setup, D_sp, H_sp, a,
                                         addcoredensity=True):
        # Get the XC-correction instance
        c = setup.xc_correction
        ncresp_g = setup.extra_xc_data['core_response'] / self.nspins
        if not addcoredensity:
            ncresp_g[:] = 0.0

        for D_p, dEdD_p, Dresp_p in zip(self.Ddist_asp[a], H_sp,
                                        self.Drespdist_asp[a]):
            D_Lq = np.dot(c.B_pqL.T, D_p)
            n_Lg = np.dot(D_Lq, c.n_qg)  # Construct density
            if addcoredensity:
                n_Lg[0] += c.nc_g * sqrt(4 * pi) / self.nspins
            nt_Lg = np.dot(
                D_Lq, c.nt_qg
            )  # Construct smooth density (_without smooth core density_)

            Dresp_Lq = np.dot(c.B_pqL.T, Dresp_p)
            nresp_Lg = np.dot(Dresp_Lq, c.n_qg)  # Construct 'response density'
            nrespt_Lg = np.dot(
                Dresp_Lq, c.nt_qg
            )  # Construct smooth 'response density' (w/o smooth core)

            for w, Y_L in zip(weight_n, c.Y_nL):
                nt_g = np.dot(Y_L, nt_Lg)
                nrespt_g = np.dot(Y_L, nrespt_Lg)
                x_g = nrespt_g / (nt_g + self.damp)
                dEdD_p -= self.weight * w * np.dot(
                    np.dot(c.B_pqL, Y_L), np.dot(c.nt_qg, x_g * c.rgd.dv_g))

                n_g = np.dot(Y_L, n_Lg)
                nresp_g = np.dot(Y_L, nresp_Lg)
                x_g = (nresp_g + ncresp_g) / (n_g + self.damp)

                dEdD_p += self.weight * w * np.dot(
                    np.dot(c.B_pqL, Y_L), np.dot(c.n_qg, x_g * c.rgd.dv_g))

        return 0.0

    def integrate_sphere(self, a, Dresp_sp, D_sp, Dwf_p, spin=0):
        c = self.wfs.setups[a].xc_correction
        Dresp_p, D_p = Dresp_sp[spin], D_sp[spin]
        D_Lq = np.dot(c.B_pqL.T, D_p)
        n_Lg = np.dot(D_Lq, c.n_qg)  # Construct density
        n_Lg[0] += c.nc_g * sqrt(4 * pi) / len(D_sp)
        nt_Lg = np.dot(D_Lq, c.nt_qg
                       )  # Construct smooth density (without smooth core)
        Dresp_Lq = np.dot(c.B_pqL.T, Dresp_p)  # Construct response
        nresp_Lg = np.dot(Dresp_Lq, c.n_qg)  # Construct 'response density'
        nrespt_Lg = np.dot(
            Dresp_Lq, c.nt_qg
        )  # Construct smooth 'response density' (w/o smooth core)
        Dwf_Lq = np.dot(c.B_pqL.T, Dwf_p)  # Construct lumo wf
        nwf_Lg = np.dot(Dwf_Lq, c.n_qg)
        nwft_Lg = np.dot(Dwf_Lq, c.nt_qg)
        E = 0.0
        for w, Y_L in zip(weight_n, c.Y_nL):
            v = np.dot(Y_L, nwft_Lg) * np.dot(Y_L, nrespt_Lg) / (
                np.dot(Y_L, nt_Lg) + self.damp)
            E -= self.weight * w * np.dot(v, c.rgd.dv_g)
            v = np.dot(Y_L, nwf_Lg) * np.dot(Y_L, nresp_Lg) / (
                np.dot(Y_L, n_Lg) + self.damp)
            E += self.weight * w * np.dot(v, c.rgd.dv_g)
        return E

    def add_smooth_xc_potential_and_energy_1d(self, vt_g):
        w_ln = self.coefficients.get_coefficients_1d(smooth=True)
        v_g = np.zeros(self.ae.N)
        n_g = np.zeros(self.ae.N)
        for w_n, f_n, u_n in zip(w_ln, self.ae.f_ln,
                                 self.ae.s_ln):  # For each angular momentum
            u2_n = safe_sqr(u_n)
            v_g += np.dot(w_n, u2_n)
            n_g += np.dot(f_n, u2_n)

        vt_g += self.weight * v_g / (n_g + self.damp)
        return 0.0  # Response part does not contribute to energy

    def calculate_delta_xc(self, homolumo=None):
        warnings.warn(
            'The function calculate_delta_xc() is deprecated. '
            'Please use calculate_discontinuity_potential() instead. '
            'See documentation on calculating band gap with GLLBSC.',
            DeprecationWarning)

        if homolumo is None:
            # Calculate band gap

            # This always happens, so we don't warn.
            # We should perhaps print it as an ordinary message,
            # but we do not have a file here to which to print.
            # import warnings
            # warnings.warn('Calculating KS-gap directly from the k-points, '
            #               'can be inaccurate.')
            homolumo = self.wfs.get_homo_lumo()
        homo, lumo = homolumo
        dxc_pot = self.calculate_discontinuity_potential(homo * Ha, lumo * Ha)
        self.Dxc_vt_sG = dxc_pot.vt_sG
        self.Dxc_D_asp = dxc_pot.D_asp
        self.Dxc_Dresp_asp = dxc_pot.Dresp_asp

    def calculate_discontinuity_potential(self, homo, lumo):
        r"""Calculate the derivative discontinuity potential.

        This function calculates $`\Delta_{xc}(r)`$ as given by
        Eq. (24) of https://doi.org/10.1103/PhysRevB.82.115106

        Parameters:

        homo:
            homo energy (or energies in spin-polarized case) in eV
        lumo:
            lumo energy (or energies in spin-polarized case) in eV

        Returns:

        dxc_pot: Discontinuity potential
        """
        homolumo = np.asarray((homo, lumo)) / Ha

        dxc_Dresp_asp = self.empty_atomic_matrix()
        dxc_D_asp = self.empty_atomic_matrix()
        for a in self.density.D_asp:
            ni = self.wfs.setups[a].ni
            dxc_Dresp_asp[a] = np.zeros((self.nspins, ni * (ni + 1) // 2))
            dxc_D_asp[a] = np.zeros((self.nspins, ni * (ni + 1) // 2))

        # Calculate new response potential with LUMO reference
        w_kn = self.coefficients.get_coefficients_for_lumo_perturbation(
            self.wfs.kpt_u,
            homolumo=homolumo)
        f_kn = [kpt.f_n for kpt in self.wfs.kpt_u]

        dxc_vt_sG = self.gd.zeros(self.nspins)
        nt_sG = self.gd.zeros(self.nspins)
        for kpt, w_n in zip(self.wfs.kpt_u, w_kn):
            self.wfs.add_to_density_from_k_point_with_occupation(dxc_vt_sG,
                                                                 kpt, w_n)
            self.wfs.add_to_density_from_k_point(nt_sG, kpt)

        self.wfs.kptband_comm.sum(nt_sG)
        self.wfs.kptband_comm.sum(dxc_vt_sG)

        if self.wfs.kd.symmetry:
            for nt_G, dxc_vt_G in zip(nt_sG, dxc_vt_sG):
                self.wfs.kd.symmetry.symmetrize(nt_G, self.gd)
                self.wfs.kd.symmetry.symmetrize(dxc_vt_G, self.gd)

        dxc_vt_sG /= nt_sG + self.damp

        self.wfs.calculate_atomic_density_matrices_with_occupation(
            dxc_Dresp_asp, w_kn)
        self.wfs.calculate_atomic_density_matrices_with_occupation(
            dxc_D_asp, f_kn)
        dxc_pot = ResponsePotential(self, dxc_vt_sG, None,
                                    dxc_D_asp, dxc_Dresp_asp)
        return dxc_pot

    def calculate_delta_xc_perturbation(self):
        warnings.warn(
            'The function calculate_delta_xc_perturbation() is deprecated. '
            'Please use calculate_discontinuity() instead. '
            'See documentation on calculating band gap with GLLBSC.',
            DeprecationWarning)
        dxc_pot = ResponsePotential(self, self.Dxc_vt_sG, None, self.Dxc_D_asp,
                                    self.Dxc_Dresp_asp)
        if self.nspins != 1:
            ret = []
            for spin in range(self.nspins):
                ret.append(self.calculate_discontinuity(dxc_pot, spin=spin))
            return ret
        return self.calculate_discontinuity(dxc_pot)

    def calculate_delta_xc_perturbation_spin(self, s=0):
        warnings.warn(
            'The function calculate_delta_xc_perturbation_spin() is '
            'deprecated. Please use calculate_discontinuity_spin() instead. '
            'See documentation on calculating band gap with GLLBSC.',
            DeprecationWarning)
        dxc_pot = ResponsePotential(self, self.Dxc_vt_sG, None, self.Dxc_D_asp,
                                    self.Dxc_Dresp_asp)
        return self.calculate_discontinuity(dxc_pot, spin=s)

    def calculate_discontinuity(self,
                                dxc_pot: ResponsePotential,
                                spin: int = None):
        r"""Calculate the derivative discontinuity.

        This function evaluates the perturbation theory expression
        for the derivative discontinuity value as given by
        Eq. (25) of https://doi.org/10.1103/PhysRevB.82.115106

        Parameters:

        dxc_pot:
            Discontinuity potential.
        spin:
            Spin value.

        Returns:

        KSgap:
            Kohn-Sham gap in eV
        dxc:
            Derivative discontinuity in eV
        """
        if spin is None:
            if self.nspins != 1:
                raise ValueError('Specify spin for the discontinuity.')
            spin = 0

        # Redistribute discontinuity potential
        if dxc_pot.response is not self:
            dxc_pot.redistribute(self)

        homo, lumo = self.wfs.get_homo_lumo(spin)
        KSgap = lumo - homo

        # Find the lumo-orbital of this spin
        eps_n = self.wfs.kpt_qs[0][spin].eps_n
        assert self.wfs.bd.comm.size == 1
        lumo_n = (eps_n < self.wfs.fermi_level).sum()

        # Calculate the expectation value for all k points, and keep
        # the smallest energy value
        nt_G = self.gd.empty()
        min_energy = np.inf
        for u, kpt in enumerate(self.wfs.kpt_u):
            if kpt.s == spin:
                nt_G[:] = 0.0
                self.wfs.add_orbital_density(nt_G, kpt, lumo_n)
                E = 0.0
                for a in dxc_pot.D_asp:
                    D_sp = dxc_pot.D_asp[a]
                    Dresp_sp = dxc_pot.Dresp_asp[a]
                    P_ni = kpt.P_ani[a]
                    Dwf_p = pack(np.outer(P_ni[lumo_n].T.conj(),
                                          P_ni[lumo_n]).real)
                    E += self.integrate_sphere(a, Dresp_sp, D_sp, Dwf_p,
                                               spin=spin)
                E = self.gd.comm.sum(E)
                E += self.gd.integrate(nt_G * dxc_pot.vt_sG[spin])
                E += kpt.eps_n[lumo_n] - lumo
                min_energy = min(min_energy, E)

        # Take the smallest value over all distributed k points
        dxc = self.wfs.kd.comm.min(min_energy)
        return KSgap * Ha, dxc * Ha

    def calculate_discontinuity_from_average(self,
                                             dxc_pot: ResponsePotential,
                                             spin: int,
                                             testing: bool = False):
        msg = ('This function estimates discontinuity by calculating '
               'the average of the discontinuity potential. '
               'Use only for testing and debugging.')
        if not testing:
            raise RuntimeError(msg)
        else:
            warnings.warn(msg)
        assert self.wfs.world.size == 1

        # Calculate average of lumo reference response potential
        dxc = np.average(dxc_pot.vt_sG[spin])
        return dxc * Ha

    def initialize_from_atomic_orbitals(self, basis_functions):
        # Initialize 'response-density' and density-matrices
        # print('Initializing from atomic orbitals')
        self.Dresp_asp = self.empty_atomic_matrix()
        setups = self.wfs.setups

        for a in self.density.D_asp.keys():
            ni = setups[a].ni
            self.Dresp_asp[a] = np.zeros((self.nspins, ni * (ni + 1) // 2))

        self.D_asp = self.empty_atomic_matrix()
        f_asi = {}
        w_asi = {}

        for a in basis_functions.atom_indices:
            if setups[a].type == 'ghost':
                w_j = [0.0]
            else:
                w_j = setups[a].extra_xc_data['w_j']

            # Basis function coefficients based of response weights
            w_si = setups[a].calculate_initial_occupation_numbers(
                0, False,
                charge=0,
                nspins=self.nspins,
                f_j=w_j)
            # Basis function coefficients based on density
            f_si = setups[a].calculate_initial_occupation_numbers(
                0, False,
                charge=0,
                nspins=self.nspins)

            if a in basis_functions.my_atom_indices:
                self.Dresp_asp[a] = setups[a].initialize_density_matrix(w_si)
                self.D_asp[a] = setups[a].initialize_density_matrix(f_si)

            f_asi[a] = f_si
            w_asi[a] = w_si

        self.Drespdist_asp = self.distribute_D_asp(self.Dresp_asp)
        self.Ddist_asp = self.distribute_D_asp(self.D_asp)
        nt_sG = self.gd.zeros(self.nspins)
        basis_functions.add_to_density(nt_sG, f_asi)
        self.vt_sG = self.gd.zeros(self.nspins)
        basis_functions.add_to_density(self.vt_sG, w_asi)
        # Update vt_sG to correspond atomic response potential. This will be
        # used until occupations and eigenvalues are available.
        self.vt_sG /= nt_sG + self.damp
        self.vt_sg = self.finegd.zeros(self.nspins)
        self.density.distribute_and_interpolate(self.vt_sG, self.vt_sg)

    def get_extra_setup_data(self, extra_data):
        ae = self.ae
        njcore = ae.njcore
        w_ln = self.coefficients.get_coefficients_1d(smooth=True)
        w_j = []
        for w_n in w_ln:
            for w in w_n:
                w_j.append(w)
        extra_data['w_j'] = w_j

        w_j = self.coefficients.get_coefficients_1d()
        x_g = np.dot(w_j[:njcore], safe_sqr(ae.u_j[:njcore]))
        x_g[1:] /= ae.r[1:] ** 2 * 4 * np.pi
        x_g[0] = x_g[1]
        extra_data['core_response'] = x_g

        # For debugging purposes
        w_j = self.coefficients.get_coefficients_1d()
        u2_j = safe_sqr(self.ae.u_j)
        v_g = self.weight * np.dot(w_j, u2_j) / (
            np.dot(self.ae.f_j, u2_j) + self.damp)
        v_g[0] = v_g[1]
        extra_data['all_electron_response'] = v_g

        # Calculate Hardness of spherical atom, for debugging purposes
        l = [np.where(f < 1e-3, e, 1000)
             for f, e in zip(self.ae.f_j, self.ae.e_j)]
        h = [np.where(f > 1e-3, e, -1000)
             for f, e in zip(self.ae.f_j, self.ae.e_j)]
        lumo_e = min(l)
        homo_e = max(h)
        if lumo_e < 999:  # If there is unoccpied orbital
            w_j = self.coefficients.get_coefficients_1d_for_lumo_perturbation()
            v_g = self.weight * np.dot(w_j, u2_j) / (
                np.dot(self.ae.f_j, u2_j) + self.damp)
            e2 = [e + np.dot(u2 * v_g, self.ae.dr)
                  for u2, e in zip(u2_j, self.ae.e_j)]
            lumo_2 = min([np.where(f < 1e-3, e, 1000)
                          for f, e in zip(self.ae.f_j, e2)])
            # print('New lumo eigenvalue:', lumo_2 * 27.2107)
            self.hardness = lumo_2 - homo_e
            # print('Hardness predicted: %10.3f eV' %
            #       (self.hardness * 27.2107))

    def write(self, writer):
        """Writes response specific data."""
        wfs = self.wfs
        kpt_comm = wfs.kd.comm
        gd = wfs.gd

        # Write the pseudodensity on the coarse grid:
        shape = (wfs.nspins,) + tuple(gd.get_size_of_global_array())
        writer.add_array('gllb_pseudo_response_potential', shape)
        if kpt_comm.rank == 0:
            for vt_G in self.vt_sG:
                writer.fill(gd.collect(vt_G) * Ha)

        def pack(X0_asp):
            X_asp = wfs.setups.empty_atomic_matrix(
                wfs.nspins, wfs.atom_partition)
            # XXX some of the provided X0_asp contain strangely duplicated
            # elements.  Take only the minimal set:
            for a in X_asp:
                X_asp[a][:] = X0_asp[a]
            return pack_atomic_matrices(X_asp)

        writer.write(gllb_atomic_density_matrices=pack(self.D_asp))
        writer.write(gllb_atomic_response_matrices=pack(self.Dresp_asp))

    def empty_atomic_matrix(self):
        assert self.wfs.atom_partition is self.density.atom_partition
        return self.wfs.setups.empty_atomic_matrix(self.wfs.nspins,
                                                   self.wfs.atom_partition)

    def read(self, reader):
        r = reader.hamiltonian.xc
        wfs = self.wfs

        d('Reading vt_sG')
        self.gd.distribute(r.gllb_pseudo_response_potential / reader.ha,
                           self.vt_sG)
        self.density.distribute_and_interpolate(self.vt_sG, self.vt_sg)

        def unpack(D_sP):
            return unpack_atomic_matrices(D_sP, wfs.setups)

        # Read atomic density matrices and non-local part of hamiltonian:
        D_asp = unpack(r.gllb_atomic_density_matrices)
        Dresp_asp = unpack(r.gllb_atomic_response_matrices)

        # All density matrices are loaded to all cores
        # First distribute them to match density.D_asp
        self.D_asp = self.empty_atomic_matrix()
        self.Dresp_asp = self.empty_atomic_matrix()
        for a in self.density.D_asp:
            self.D_asp[a][:] = D_asp[a]
            self.Dresp_asp[a][:] = Dresp_asp[a]
        assert len(self.Dresp_asp) == len(self.density.D_asp)

        # Further distributes the density matricies for xc-corrections
        self.Ddist_asp = self.distribute_D_asp(self.D_asp)
        self.Drespdist_asp = self.distribute_D_asp(self.Dresp_asp)

    def heeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeelp(self, olddens):
        # XXX This function should be removed once the deprecated
        # `fixdensity=True` option is removed.
        from gpaw.density import redistribute_array
        self.vt_sg = redistribute_array(self.vt_sg,
                                        olddens.finegd, self.finegd,
                                        self.wfs.nspins, self.wfs.kptband_comm)
        self.vt_sG = redistribute_array(self.vt_sG,
                                        olddens.gd, self.gd,
                                        self.wfs.nspins, self.wfs.kptband_comm)
