#!/usr/bin/env python

"""
This module defines the VaspInputSet abstract base class and a concrete
implementation for the parameters used by the Materials Project and the MIT
high throughput project.  The basic concept behind an input set is to specify
a scheme to generate a consistent set of Vasp inputs from a structure
without further user intervention. This ensures comparability across
runs.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Nov 16, 2011"

import os
import abc
import ConfigParser
import json
import re

from pymatgen.io.vaspio.vasp_input import Incar, Poscar, Potcar, Kpoints
from pymatgen.serializers.json_coders import MSONable


class AbstractVaspInputSet(MSONable):
    """
    Abstract base class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files
    (INPUT, KPOINTS, POSCAR and POTCAR) can be generated in an automated
    fashion for any structure.
    """
    __metaclass__ = abc.ABCMeta

    def get_poscar(self, structure):
        """
        Returns Poscar from a structure.
        """
        return Poscar(structure)

    @abc.abstractmethod
    def get_kpoints(self, structure):
        """
        Returns Kpoints from a structure.

        Args:
            structure:
                Structure object

        Returns:
            Kpoints object
        """
        return

    @abc.abstractmethod
    def get_incar(self, structure):
        """
        Returns Incar from a structure.

        Args:
            structure:
                Structure object

        Returns:
            Incar object
        """
        return

    @abc.abstractmethod
    def get_potcar(self, structure):
        """
        Returns Potcar from a structure.

        Args:
            structure:
                Structure object

        Returns:
            Potcar object
        """
        return

    @abc.abstractmethod
    def get_potcar_symbols(self, structure):
        """
        Returns list of POTCAR symbols from a structure.

        Args:
            structure:
                Structure object

        Returns:
            List of POTCAR symbols
        """
        return

    def get_all_vasp_input(self, structure, generate_potcar=True):
        """
        Returns all input files as a dict of {filename: vaspio object}

        Args:
            structure:
                Structure object
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a
                POTCAR, which contains the POTCAR labels but not the actual
                POTCAR. Defaults to True.

        Returns:
            dict of {filename: file_as_string}, e.g., {'INCAR':'EDIFF=1e-4...'}
        """
        d = {'INCAR': self.get_incar(structure),
             'KPOINTS': self.get_kpoints(structure),
             'POSCAR': self.get_poscar(structure)}
        if generate_potcar:
            d['POTCAR'] = self.get_potcar(structure)
        else:
            d['POTCAR.spec'] = "\n".join(self.get_potcar_symbols(structure))
        return d

    def write_input(self, structure, output_dir, make_dir_if_not_present=True):
        """
        Writes a set of VASP input to a directory.

        Args:
            structure:
                Structure object
            output_dir:
                Directory to output the VASP input files
            make_dir_if_not_present:
                Set to True if you want the directory (and the whole path) to
                be created if it is not present.
        """
        if make_dir_if_not_present and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for k, v in self.get_all_vasp_input(structure).items():
            v.write_file(os.path.join(output_dir, k))


class DictVaspInputSet(AbstractVaspInputSet):
    """
    Concrete implementation of VaspInputSet that is initialized from a dict
    settings. This allows arbitrary settings to be input. In general,
    this is rarely used directly unless there is a source of settings in JSON
    format (e.g., from a REST interface). It is typically used by other
    VaspInputSets for initialization.

    Special consideration should be paid to the way the MAGMOM initialization
    for the INCAR is done. The initialization differs depending on the type of
    structure and the configuration settings. The order in which the magmom is
    determined is as follows:

    1. If the site itself has a magmom setting, that is used.
    2. If the species on the site has a spin setting, that is used.
    3. If the species itself has a particular setting in the config file, that
       is used, e.g., Mn3+ may have a different magmom than Mn4+.
    4. Lastly, the element symbol itself is checked in the config file. If
       there are no settings, VASP's default of 0.6 is used.


    """
    def __init__(self, name, config_dict, constrain_total_magmom=False):
        """
        Args:
            name:
                The name in the config file.
            config_dict:
                The config dictionary to use.
            user_incar_settings:
                User INCAR settings. This allows a user to override INCAR
                settings, e.g., setting a different MAGMOM for various elements
                or species.
            constrain_total_magmom:
                Whether to constrain the total magmom (NUPDOWN in INCAR) to be
                the sum of the expected MAGMOM for all species. Defaults to
                False.
        """
        self.name = name
        self.potcar_settings = config_dict["POTCAR"]
        self.kpoints_settings = config_dict['KPOINTS']
        self.incar_settings = config_dict['INCAR']
        self.set_nupdown = constrain_total_magmom

    def get_incar(self, structure):
        incar = Incar()
        comp = structure.composition
        elements = sorted([el for el in comp.elements if comp[el] > 0],
                          key=lambda el: el.X)
        most_electroneg = elements[-1].symbol
        poscar = Poscar(structure)
        for key, setting in self.incar_settings.items():
            if key == "MAGMOM":
                mag = []
                for site in structure:
                    if hasattr(site, 'magmom'):
                        mag.append(site.magmom)
                    elif hasattr(site.specie, 'spin'):
                        mag.append(site.specie.spin)
                    elif str(site.specie) in setting:
                        mag.append(setting.get(str(site.specie)))
                    else:
                        mag.append(setting.get(site.specie.symbol, 0.6))
                incar[key] = mag
            elif key in ('LDAUU', 'LDAUJ', 'LDAUL'):
                if most_electroneg in setting.keys():
                    incar[key] = [setting[most_electroneg].get(sym, 0)
                                  for sym in poscar.site_symbols]
                else:
                    incar[key] = [0] * len(poscar.site_symbols)
            elif key == "EDIFF":
                incar[key] = float(setting) * structure.num_sites
            else:
                incar[key] = setting

        has_u = ("LDAUU" in incar and sum(incar['LDAUU']) > 0)
        if has_u:
            # modify LMAXMIX if LSDA+U and you have d or f electrons
            # note that if the user explicitly sets LMAXMIX in settings it will
            # override this logic.
            if 'LMAXMIX' not in self.incar_settings.keys():
                # contains f-electrons
                if any([el.Z > 56 for el in structure.composition]):
                    incar['LMAXMIX'] = 6
                # contains d-electrons
                elif any([el.Z > 20 for el in structure.composition]):
                    incar['LMAXMIX'] = 4
        else:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]

        if self.set_nupdown:
            nupdown = sum([mag if mag > 0.6 else 0 for mag in incar['MAGMOM']])
            incar['NUPDOWN'] = nupdown

        return incar

    def get_potcar(self, structure):
        return Potcar(self.get_potcar_symbols(structure))

    def get_potcar_symbols(self, structure):
        p = self.get_poscar(structure)
        elements = p.site_symbols
        potcar_symbols = []
        for el in elements:
            potcar_symbols.append(self.potcar_settings[el]
                                  if el in self.potcar_settings else el)
        return potcar_symbols

    def get_kpoints(self, structure):
        """
        Writes out a KPOINTS file using the fully automated grid method. Uses
        Gamma centered meshes  for hexagonal cells and Monk grids otherwise.

        Algorithm:
            Uses a simple approach scaling the number of divisions along each
            reciprocal lattice vector proportional to its length.
        """
        dens = int(self.kpoints_settings['grid_density'])
        return Kpoints.automatic_density(structure, dens)

    def __str__(self):
        return self.name

    def __repr__(self):
        output = [self.name, ""]
        section_names = ['INCAR settings', 'KPOINTS settings',
                         'POTCAR settings']
        count = 0
        for d in [self.incar_settings, self.kpoints_settings,
                  self.potcar_settings]:
            output.append(section_names[count])
            for k, v in d.items():
                output.append("%s = %s" % (k, str(v)))
            output.append("")
            count += 1
        return "\n".join(output)

    @property
    def to_dict(self):
        config_dict = {
            "INCAR": self.incar_settings,
            "KPOINTS": self.kpoints_settings,
            "POTCAR": self.potcar_settings
        }
        return {
            "name": self.name,
            "config_dict": config_dict,
            "constrain_total_magmom": self.set_nupdown,
            "@class": self.__class__.__name__,
            "@module": self.__module__.__name__,
        }

    @staticmethod
    def from_dict(d):
        return DictVaspInputSet(d["name"], d["config_dict"],
                                d["constrain_total_magmom"])


class VaspInputSet(DictVaspInputSet):
    """
    Standard implementation of VaspInputSet that uses a config file to
    initialize settings. See DictVaspInputSet for specific details regarding
    how MAGMOM, LDAU settings are set.
    """
    def __init__(self, name, config_file, user_incar_settings=None,
                 constrain_total_magmom=False):
        """
        Args:
            name:
                The name in the config file.
            config_file:
                The config file to use. If None (the default), a default config
                file containing Materials Project and MIT parameters is used.
            user_incar_settings:
                User INCAR settings. This allows a user to override INCAR
                settings, e.g., setting a different MAGMOM for various elements
                or species.
            constrain_total_magmom:
                Whether to constrain the total magmom (NUPDOWN in INCAR) to be
                the sum of the expected MAGMOM for all species. Defaults to
                False.
        """
        self.name = name
        self.config_file = config_file
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        with open(config_file, "r") as f:
            self._config.readfp(f)

        self.user_incar_settings = user_incar_settings
        potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        incar_settings = dict(self._config.items(self.name + 'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            if key in incar_settings:
                incar_settings[key] = json.loads(incar_settings[key])
        if user_incar_settings:
            incar_settings.update(user_incar_settings)
        DictVaspInputSet.__init__(
            self, name, {"INCAR": incar_settings, "KPOINTS": kpoints_settings,
                         "POTCAR": potcar_settings},
            constrain_total_magmom=constrain_total_magmom)

    @property
    def to_dict(self):
        return {
            "name": self.name,
            "config_file": self.config_file,
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__module__.__name__,
        }

    @staticmethod
    def from_dict(d):
        return VaspInputSet(d["name"], d["config_file"],
                            user_incar_settings=d["user_incar_settings"],
                            constrain_total_magmom=d["constrain_total_magmom"])


class MITVaspInputSet(VaspInputSet):
    """
    Standard implementation of VaspInputSet utilizing parameters in the MIT
    High-throughput project.
    The parameters are chosen specifically for a high-throughput project,
    which means in general pseudopotentials with fewer electrons were chosen.

    Please refer to A Jain, G. Hautier, C. Moore, S. P. Ong, C. Fischer,
    T. Mueller, K. A. Persson, G. Ceder. A high-throughput infrastructure for
    density functional theory calculations. Computational Materials Science,
    2011, 50(8), 2295-2310.
    doi:10.1016/j.commatsci.2011.02.023 for more information.
    """
    def __init__(self, user_incar_settings=None, constrain_total_magmom=False):
        module_dir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(module_dir, "MITVaspInputSet.json")) as f:
            DictVaspInputSet.__init__(
                self, "MIT", json.load(f),
                constrain_total_magmom=constrain_total_magmom)
        self.user_incar_settings = user_incar_settings
        if user_incar_settings:
            self.incar_settings.update(user_incar_settings)

    @property
    def to_dict(self):
        return {
            "name": self.name,
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__class__.__module__,
        }

    @staticmethod
    def from_dict(d):
        return MITVaspInputSet(
            user_incar_settings=d["user_incar_settings"],
            constrain_total_magmom=d["constrain_total_magmom"])


class MITGGAVaspInputSet(VaspInputSet):
    """
    Typical implementation of input set for a GGA run based on MIT parameters.
    """
    def __init__(self, user_incar_settings=None, constrain_total_magmom=False):
        module_dir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(module_dir, "MITVaspInputSet.json")) as f:
            DictVaspInputSet.__init__(
                self, "MaterialsProject GGA", json.load(f),
                constrain_total_magmom=constrain_total_magmom)
        self.user_incar_settings = user_incar_settings
        # INCAR settings to override for GGA runs, since we are basing off a
        # GGA+U inputset
        self.incar_settings['LDAU'] = False
        if 'LDAUU' in self.incar_settings:
            # technically not needed, but clarifies INCAR
            del self.incar_settings['LDAUU']
            del self.incar_settings['LDAUJ']
            del self.incar_settings['LDAUL']
        if user_incar_settings:
            self.incar_settings.update(user_incar_settings)

    @property
    def to_dict(self):
        return {
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__class__.__module__,
        }

    @staticmethod
    def from_dict(d):
        return MITGGAVaspInputSet(
            user_incar_settings=d["user_incar_settings"],
            constrain_total_magmom=d["constrain_total_magmom"])


class MITHSEVaspInputSet(VaspInputSet):
    """
    Typical implementation of input set for a HSE run.
    """
    def __init__(self, user_incar_settings=None, constrain_total_magmom=False):
        module_dir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(module_dir, "MITHSEVaspInputSet.json")) as f:
            DictVaspInputSet.__init__(
                self, "MIT HSE", json.load(f),
                constrain_total_magmom=constrain_total_magmom)
        self.user_incar_settings = user_incar_settings
        if user_incar_settings:
            self.incar_settings.update(user_incar_settings)

    @property
    def to_dict(self):
        return {
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__class__.__module__,
        }

    @staticmethod
    def from_dict(d):
        return MITHSEVaspInputSet(
            user_incar_settings=d["user_incar_settings"],
            constrain_total_magmom=d["constrain_total_magmom"])


class MaterialsProjectVaspInputSet(DictVaspInputSet):
    """
    Implementation of VaspInputSet utilizing parameters in the public
    Materials Project. Typically, the pseudopotentials chosen contain more
    electrons than the MIT parameters, and the k-point grid is ~50% more dense.
    The LDAUU parameters are also different due to the different psps used,
    which result in different fitted values (even though the methodology of
    fitting is exactly the same as the MIT scheme).
    """
    def __init__(self, user_incar_settings=None, constrain_total_magmom=False):
        module_dir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(module_dir, "MPVaspInputSet.json")) as f:
            DictVaspInputSet.__init__(
                self, "MaterialsProject", json.load(f),
                constrain_total_magmom=constrain_total_magmom)
        self.user_incar_settings = user_incar_settings
        if user_incar_settings:
            self.incar_settings.update(user_incar_settings)

    @property
    def to_dict(self):
        return {
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__class__.__module__,
        }

    @staticmethod
    def from_dict(d):
        return MaterialsProjectVaspInputSet(
            user_incar_settings=d["user_incar_settings"],
            constrain_total_magmom=d["constrain_total_magmom"])


class MaterialsProjectGGAVaspInputSet(DictVaspInputSet):
    """
    Same as the MaterialsProjectVaspInput set, but the +U is enforced to be
    turned off.
    """
    def __init__(self, user_incar_settings=None, constrain_total_magmom=False):
        module_dir = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(module_dir, "MPVaspInputSet.json")) as f:
            DictVaspInputSet.__init__(
                self, "MaterialsProject GGA", json.load(f),
                constrain_total_magmom=constrain_total_magmom)

        self.user_incar_settings = user_incar_settings
        # INCAR settings to override for GGA runs, since we are basing off a
        # GGA+U inputset
        self.incar_settings['LDAU'] = False
        if 'LDAUU' in self.incar_settings:
            # technically not needed, but clarifies INCAR
            del self.incar_settings['LDAUU']
            del self.incar_settings['LDAUJ']
            del self.incar_settings['LDAUL']
        if user_incar_settings:
            self.incar_settings.update(user_incar_settings)

    @property
    def to_dict(self):
        return {
            "constrain_total_magmom": self.set_nupdown,
            "user_incar_settings": self.user_incar_settings,
            "@class": self.__class__.__name__,
            "@module": self.__class__.__module__,
        }

    @staticmethod
    def from_dict(d):
        return MaterialsProjectGGAVaspInputSet(
            user_incar_settings=d["user_incar_settings"],
            constrain_total_magmom=d["constrain_total_magmom"])


def batch_write_vasp_input(structures, vasp_input_set, output_dir,
                           make_dir_if_not_present=True, subfolder=None,
                           sanitize=False, include_cif=False):
    """
    Batch write vasp input for a sequence of structures to
    output_dir, following the format output_dir/{group}/{formula}_{number}.

    Args:
        structures:
            Sequence of Structures.
        vasp_input_set:
            pymatgen.io.vaspio_set.VaspInputSet like object that creates
            vasp input files from structures
        output_dir:
            Directory to output files
        make_dir_if_not_present:
            Create the directory if not present. Defaults to True.
        subfolder:
            function to create subdirectory name from structure.
            Defaults to simply "formula_count".
        sanitize:
            Boolean indicating whether to sanitize the structure before
            writing the VASP input files. Sanitized output are generally easier
            for viewing and certain forms of analysis. Defaults to False.
        include_cif:
            Boolean indication whether to output a CIF as well. CIF files are
            generally better supported in visualization programs.
    """
    for i, s in enumerate(structures):
        formula = re.sub("\s+", "", s.formula)
        if subfolder is not None:
            subdir = subfolder(s)
            dirname = os.path.join(output_dir, subdir)
        else:
            dirname = os.path.join(output_dir, '{}_{}'.format(formula, i))
        if sanitize:
            s = s.copy(sanitize=True)
        vasp_input_set.write_input(
            s, dirname, make_dir_if_not_present=make_dir_if_not_present
        )
        if include_cif:
            from pymatgen.io.cifio import CifWriter
            writer = CifWriter(s)
            writer.write_file(os.path.join(dirname, "{}_{}.cif"
                                           .format(formula, i)))
