from fireworks import Firework
from pymatgen.io.vasp.sets import MPStaticSet

from atomate.common.firetasks.glue_tasks import CreateFolder, PassCalcLocs
from atomate.vasp.firetasks import CheckBandgap, CopyVaspOutputs, ModifyIncar, RunVaspCustodian, VaspToDb
from atomate.vasp.fireworks import StaticInterpolateFW, StaticFW


class LcalcpolFW(Firework):
    def __init__(self, structure, name="static dipole moment", static_name="static", vasp_cmd="vasp",
                 vasp_input_set=None, vasp_input_set_params=None, db_file=None, parents=None,
                 gap_threshold=0.010, interpolate=False, start=None, end=None, this_image=0, nimages=5, **kwargs):
        """
        Static Firework that calculates the dipole moment of structure or a single interpolated image of structures
        output by two calculations specified by PassCalcLoc names start and end.

        This Firework uses three steps to calculate the dipole moment.
        1. A StaticFW or StaticInterpolatFW calculation is performed to compute the band gap of the structure.
        2. Because VASP cannot compute the dipole moment of metallic structures, CheckBandgap checks that the structure
        has a band gap greater than gap_threshold. If the structure has a band gap less than gap_threshold,
        the Firework defuses. Otherwise,
        3. a polarization calculation (LCALCPOL=TRUE) is performed to calculate the dipole moment.

        If interpolate is equal to True, the keywords start and end are PassCalcLoc names used to locate the output
        structures (CONTCAR) of two previous calculations. These two structures must have the same number and
        type of atoms given in identical order in the structures. These structures are used to create nimages structures
        interpolations.

        this_image specifies which of the interpolated structures to use in this Firework. For example in the
        ferroelectric workflow, two structures that can be distorted into one another (a
        high-symmetry nonpolar structure and a low-symmetry polar structure) are relaxed using OptimizeFW. Then the
        dipole moment is calculated for nimages intermediate structures generated by interpolating between these two
        relaxed structures.

        Args:
            structure (Structure): Input structure. For an interpolation, this is a dummy structure. See interpolate arg
                description.
            name (str): Name for the polarization FireWork.
            static_name (str): Name for the SCF run to be used in PassCalcLoc if copy_vasp_outputs != True.
            vasp_cmd (str): Command to run vasp.
            vasp_input_set (str): string name for the VASP input set (e.g., "MITMDVaspInputSet").
            vasp_input_set_params (dict): Dict of vasp_input_set_kwargs.
            db_file (str): Path to file specifying db credentials.
            parents (Firework): Parents of this particular Firework. FW or list of FWS.
            gap_threshold: Band gap cutoff for determining whether polarization calculation will proceed from
                SCF band gap.
            interpolate (bool): use structure created by interpolating CONTCAR structure from calculations indicated by
                start and end PassCalcLoc names. If interpolate is False, start and end do not need to be set.
            start (str): PassCalcLoc name of StaticFW or RelaxFW run of starting structure
            end (str): PassCalcLoc name of StaticFW or RelaxFW run of ending structure
            this_image (int): which interpolation to use for this run of the nimage interpolations.
            nimages (int): number of interpolations between CONTCAR structures from calculations indicated by start and
                end args.
        """

        t = []

        # Ensure that LWAVE is set to true so we can use WAVECAR for polarization calculation.
        vasp_input_set_params = vasp_input_set_params or {}
        vasp_input_set_params = vasp_input_set_params.copy()

        if not vasp_input_set_params.get('user_incar_settings', None):
            vasp_input_set_params.update({'user_incar_settings': {}})
        vasp_input_set_params['user_incar_settings'].update({'LWAVE': True})

        vasp_input_set = vasp_input_set or "MPStaticSet"

        if interpolate:
            static = StaticInterpolateFW(structure, start, end, name=static_name, vasp_input_set=vasp_input_set,
                                         vasp_input_set_params=vasp_input_set_params, vasp_cmd=vasp_cmd,
                                         db_file=db_file, parents=parents, this_image=this_image,
                                         nimages=nimages, **kwargs)
        else:
            vasp_input_set = MPStaticSet(structure, **vasp_input_set_params)
            static = StaticFW(structure, name=static_name, vasp_input_set=vasp_input_set,
                              vasp_input_set_params=vasp_input_set_params, vasp_cmd=vasp_cmd,
                              db_file=db_file, parents=parents, **kwargs)
        t.extend(static.tasks)

        # Defuse workflow if bandgap is less than gap_threshold.
        t.append(CheckBandgap(min_gap=gap_threshold))

        # Create new directory and move to that directory to perform polarization calculation
        t.append(CreateFolder(folder_name="polarization",change_dir=True))

        # Copy VASP Outputs from static calculation
        t.append(CopyVaspOutputs(calc_loc=static_name,
                                 additional_files=["CHGCAR", "WAVECAR"],
                                 contcar_to_poscar=True))

        t.extend([
            # WriteVaspStaticFromPrev(prev_calc_dir="../",other_params={'lcalcpol':True}),
            ModifyIncar(incar_update={'lcalcpol':True}),
            RunVaspCustodian(vasp_cmd=vasp_cmd),
            PassCalcLocs(name=name),
            VaspToDb(db_file=db_file,
                     additional_fields={"task_label": name})])

        # Note, Outcar must have read_lcalcpol method for polarization information to be processed.
        # ...assuming VaspDrone will automatically assimilate all properties of the Outcar.

        super(LcalcpolFW, self).__init__(t, parents=parents, name="{}-{}".format(
            structure.composition.reduced_formula,
            name), **kwargs)