from pawpyseed.core.wavefunction import *


class NCLWavefunction(pawpyc.CNCLWavefunction, Wavefunction):
    def __init__(self, struct, pwf, cr, dim, symprec=1e-4, setup_projectors=False):
        """
        Arguments:
                struct (pymatgen.core.Structure): structure that the wavefunction describes
                pwf (pawpyc.PWFPointer): holder class for pswf_t and k-points/k-point weights
                cr (CoreRegion): Contains the pseudopotentials, with projectors and
                        partials waves, for the structure
                dim (pymatgen.io.vasp.outputs.Outcar OR np.ndarry OR list of length 3):
                        Outcar object for reading ngf or the dimensions NG* of the FFT grid
                setup_projectors (bool, False): Whether to set up the core region
                        components of the wavefunctions. Pawpyseed will set up the projectors
                        automatically when they are first needed, so this generally
                        can be left as False.
        Returns:
                Wavefunction object
        """
        self.band_props = pwf.band_props.copy(order="C")
        super(Wavefunction, self).__init__(pwf)
        if not self.ncl:
            raise PAWpyError(
                "Pseudowavefunction is collinear! Call Wavefunction(...) instead"
            )
        self.structure = struct
        self.cr = cr
        self.dim = np.array(dim).astype(np.int32)
        if setup_projectors:
            self.check_c_projectors()

    @staticmethod
    def from_files(
        struct="CONTCAR",
        wavecar="WAVECAR",
        cr="POTCAR",
        vr="vasprun.xml",
        setup_projectors=False,
    ):
        """
        Construct a Wavefunction object from file paths.

        Arguments:
                struct (str): VASP POSCAR or CONTCAR file path
                wavecar (str): VASP WAVECAR file path
                cr (str): VASP POTCAR file path
                vr (str): VASP vasprun file path
                outcar (str): VASP OUTCAR file path
                setup_projectors (bool, False): Whether to set up the core region
                        components of the wavefunctions. Pawpyseed will set up the projectors
                        automatically when they are first needed, so this generally
                        can be left as False.

        Returns:
                Wavefunction object
        """
        vr = Vasprun(vr)
        dim = np.array(
            [vr.parameters["NGX"], vr.parameters["NGY"], vr.parameters["NGZ"]]
        )
        symprec = vr.parameters["SYMPREC"]
        pwf = pawpyc.PWFPointer(wavecar, vr)
        return NCLWavefunction(
            Poscar.from_file(struct).structure,
            pwf,
            CoreRegion(Potcar.from_file(cr)),
            dim,
            symprec,
            setup_projectors,
        )

    @staticmethod
    def from_directory(path, setup_projectors=False):
        """
        Assumes VASP output has the default filenames and is located
        in the directory specificed by path.

        Arguments:
                path (str): VASP output directory
                setup_projectors (bool, False): Whether to set up the core region
                        components of the wavefunctions. Pawpyseed will set up the projectors
                        automatically when they are first needed, so this generally
                        can be left as False.

        Returns:
                Wavefunction object
        """
        filepaths = []
        for d in ["CONTCAR", "WAVECAR", "POTCAR", "vasprun.xml"]:
            filepaths.append(str(os.path.join(path, d)))
        args = filepaths + [setup_projectors]
        return NCLWavefunction.from_files(*args)

    def desymmetrized_copy(self, allkpts=None, weights=None):
        raise NotImplementedError()

    def write_state_realspace(
        self, b, k, s, fileprefix="", dim=None, scale=1, remove_phase=False
    ):
        """
        Writes the real and imaginary parts of a given band to two files,
        prefixed by fileprefix

        Args:
                b (int): band number (0-indexed!)
                k (int): kpoint number (0-indexed!)
                s (int): spin number (0-indexed!)
                fileprefix (string, ""): first part of the file name
                dim (numpy array of 3 ints, None): dimensions of the FFT grid
                scale (scalar, 1): number to multiply the realspace wavefunction by.
                        For example, VASP multiplies charge density by the volume
                        of the structure.
                remove_phase (False): If True, removes the e^(ikr) phase
                        from the wavefunction (this does not necessarily mean
                        the wavefunction is real). This is useful if you want
                        to visualize the wavefunction because the e^(ikr) phase
                        makes the wavefunction non-periodic
        Returns:
                A 3D array (indexed by x,y,z where x,y,z are fractional coordinates)
                        with complex double values for the realspace wavefunction
                The wavefunction is written in two files with z the slow index.
        """
        self.check_c_projectors()
        if dim is not None:
            self.update_dim(np.array(dim))
        filename_base = "%sB%dK%dS%d" % (fileprefix, b, k, s)
        filename1 = "%s_UP_REAL.vasp" % filename_base
        filename2 = "%s_UP_IMAG.vasp" % filename_base
        filename3 = "%s_DOWN_REAL.vasp" % filename_base
        filename4 = "%s_DOWN_IMAG.vasp" % filename_base
        res0, res1 = self._write_realspace_state(
            filename1,
            filename2,
            filename3,
            filename4,
            scale,
            b,
            k,
            s,
            remove_phase=remove_phase,
        )
        self._convert_to_vasp_volumetric(filename1, self.dim)
        self._convert_to_vasp_volumetric(filename2, self.dim)
        self._convert_to_vasp_volumetric(filename3, self.dim)
        self._convert_to_vasp_volumetric(filename4, self.dim)
        return res0, res1

    def write_density_realspace(self, filename="PYAECCAR.vasp", dim=None, scale=1):
        """
        Writes the real and imaginary parts of a given band to two files,
        prefixed by fileprefix

        Args:
                b (int): band number (0-indexed!)
                k (int): kpoint number (0-indexed!)
                s (int): spin number (0-indexed!)
                fileprefix (string, ""): first part of the file name
                dim (numpy array of 3 ints, None): dimensions of the FFT grid
                scale (scalar, 1): number to multiply the realspace wavefunction by.
                        For example, VASP multiplies charge density by the volume
                        of the structure.
        Returns:
                A 3D array (indexed by x,y,z where x,y,z are fractional coordinates)
                        with complex double values for the realspace wavefunction
                The charge density is written with z the slow index.
        """
        self.check_c_projectors()
        if dim is not None:
            self.update_dim(np.array(dim))
        res = self._write_realspace_density(filename, scale)
        self._convert_to_vasp_volumetric(filename, self.dim)
        return res
