<commit_before>"""
Low-level objects providing an abstraction for the objects involved in the
calculation.
"""
from __future__ import division, print_function

import collections
import os
import abc
import numpy as np
import pymatgen.core.units as units

from pprint import pformat
from monty.design_patterns import singleton
from pymatgen.util.string_utils import is_string
from pymatgen.core.design_patterns import Enum, AttrDict
from pymatgen.core.units import ArrayWithUnit
from pymatgen.serializers.json_coders import MSONable
from pymatgen.symmetry.finder import SymmetryFinder
from pymatgen.core.structure import Structure, Molecule
from pymatgen.io.smartio import read_structure

from .netcdf import structure_from_etsf_file


class AbivarAble(object):
    """
    An AbivarAble object provides a method to_abivars that returns a dictionary
    with the abinit variables.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def to_abivars(self):
        """
        Returns a dictionary with the abinit variables.
        """

    def __str__(self):
        return pformat(self.to_abivars(), indent=1, width=80, depth=None)


@singleton
class MandatoryVariable(object):
    """
    Singleton used to tag mandatory variables, just because I can use
    the cool syntax: variable is MANDATORY!
    """


@singleton
class DefaultVariable(object):
    """
    Singleton used to tag variables that will have the default value
    """

MANDATORY = MandatoryVariable()
DEFAULT = DefaultVariable()


class SpinMode(collections.namedtuple('SpinMode', "mode nsppol nspinor nspden"),
               AbivarAble):
    """
    Different configurations of the electron density as implemented in abinit:
    One can use asspinmode to construct the object via SpinMode.asspinmode
    (string) where string can assume the values:

        - polarized
        - unpolarized
        - afm (anti-ferromagnetic)
        - spinor (non-collinear magnetism)
        - spinor_nomag (non-collinear, no magnetism)
    """
    @classmethod
    def asspinmode(cls, obj):
        """Converts obj into a SpinMode instance"""
        if isinstance(obj, cls):
            return obj
        else:
            # Assume a string with mode
            try:
                return _mode2spinvars[obj]
            except KeyError:
                raise KeyError("Wrong value for spin_mode: %s" % str(obj))

    def to_abivars(self):
        return {
            "nsppol": self.nsppol,
            "nspinor": self.nspinor,
            "nspden": self.nspden,
        }

# An handy Multiton
_mode2spinvars = {
    "unpolarized" : SpinMode("unpolarized" , 1, 1, 1),
    "polarized"   : SpinMode("polarized"   , 2, 1, 2),
    "afm"         : SpinMode("afm"         , 1, 1, 2),
    "spinor"      : SpinMode("spinor"      , 1, 2, 4),
    "spinor_nomag": SpinMode("spinor_nomag", 1, 2, 1),
}


class Smearing(AbivarAble, MSONable):
    """
    Variables defining the smearing technique. The preferred way to instanciate
    a Smearing object is via the class method Smearing.assmearing(string)
    """
    #: Mapping string_mode --> occopt
    _mode2occopt = {
        'nosmearing' : 1,
        'fermi_dirac': 3,
        'marzari4'   : 4,
        'marzari5'   : 5,
        'methfessel' : 6,
        'gaussian'   : 7,
    }

    def __init__(self, occopt, tsmear):
        self.occopt = occopt
        self.tsmear = tsmear

    def __str__(self):
        s = "occopt %d # %s Smearing\n" % (self.occopt, self.mode)
        if self.tsmear:
            s += 'tsmear %s' % self.tsmear
        return s

    def __eq__(self, other):
        return (self.occopt == other.occopt and
                np.allclose(self.tsmear, other.tsmear))

    def __ne__(self, other):
        return not self == other

    def __bool__(self):
        return self.mode != "nosmearing"

    # py2 old version
    __nonzero__ = __bool__

    @classmethod
    def assmearing(cls, obj):
        """
        Constructs an instance of Smearing from obj. Accepts obj in the form:

            * Smearing instance
            * "name:tsmear"  e.g. "gaussian:0.004"  (Hartree units)
            * "name:tsmear units" e.g. "gaussian:0.1 eV"
            * None --> no smearing
        """
        if obj is None:
            return Smearing.nosmearing()

        if isinstance(obj, cls):
            return obj

        # obj is a string
        obj, tsmear = obj.split(":")
        obj.strip()

        if obj == "nosmearing":
            return cls.nosmearing()
        else:
            occopt = cls._mode2occopt[obj]
            try:
                tsmear = float(tsmear)
            except ValueError:
                tsmear, unit = tsmear.split()
                tsmear = units.Energy(float(tsmear), unit).to("Ha")

            return cls(occopt, tsmear)

    @property
    def mode(self):
        for (mode_str, occopt) in self._mode2occopt.items():
            if occopt == self.occopt:
                return mode_str
        raise AttributeError("Unknown occopt %s" % self.occopt)

    @staticmethod
    def nosmearing():
        return Smearing(1, None)

    def to_abivars(self):
        if self.mode == "nosmearing":
            return {}
        else:
            return {"occopt": self.occopt,
                    "tsmear": self.tsmear,}

    @property
    def to_dict(self):
        """json friendly dict representation of Smearing"""
        return {"occopt": self.occopt, "tsmear": self.tsmear,
                "@module": self.__class__.__module__,
                "@class": self.__class__.__name__}

    @staticmethod
    def from_dict(d):
        return Smearing(d["occopt"], d["tsmear"])



class ElectronsAlgorithm(dict, AbivarAble):
    "Variables controlling the SCF/NSCF algorithm."
    # None indicates that we use abinit defaults.
    _default_vars = {
        "iprcell"  : None,
        "iscf"     : None,
        "diemac"   : None,
        "diemix"   : None,
        "diemixmag": None,
        "dielam"   : None,
        "diegap"   : None,
        "dielng"   : None,
        "diecut"   : None,
        "nstep"    : 70,
    }

    def __init__(self, *args, **kwargs):
        super(ElectronsAlgorithm, self).__init__(*args, **kwargs)

        for k in self:
            if k not in self._default_vars:
                raise ValueError("%s: No default value has been provided for "
                                 "key %s" % (self.__class__.__name__, k))

    def to_abivars(self):
        return self.copy()


class Electrons(AbivarAble):
    """
    The electronic degrees of freedom
    """
    def __init__(self,
                 spin_mode = "polarized",
                 smearing  = "fermi_dirac:0.1 eV",
                 algorithm = None,
                 nband     = None,
                 fband     = None,
                 charge    = 0.0,
                 comment   = None,
                 #occupancies = None,
                ):
        """
        Constructor for Electrons object.

        Args:
            comment:
                String comment for Electrons

            charge: float
                    Total charge of the system. Default is 0.
        """
        super(Electrons, self).__init__()

        self.comment = comment

        self.smearing = Smearing.assmearing(smearing)

        self.spin_mode = SpinMode.asspinmode(spin_mode)

        self.nband = nband
        self.fband = fband
        self.charge = charge

        self.algorithm = algorithm

        # FIXME
        if nband is None:
            self.fband = 4

    @property
    def nsppol(self):
        return self.spin_mode.nsppol

    @property
    def nspinor(self):
        return self.spin_mode.nspinor

    @property
    def nspden(self):
        return self.spin_mode.nspden

    #@property
    #def to_dict(self):
    #    "json friendly dict representation"
    #    d = {}
    #    d["@module"] = self.__class__.__module__
    #    d["@class"] = self.__class__.__name__
    #    raise NotImplementedError("")
    #    return d

    #@staticmethod
    #def from_dict(d):
    #    raise NotImplementedError("")

    def to_abivars(self):
        abivars = self.spin_mode.to_abivars()

        abivars.update({
            "nband"  : self.nband,
            "fband"  : self.fband,
            "charge" : self.charge,
        })

        if self.smearing:
            abivars.update(self.smearing.to_abivars())

        if self.algorithm:
            abivars.update(self.algorithm)

        abivars["#comment"] = self.comment
        return abivars


def asabistructure(obj):
    """
    Convert obj into an AbiStructure object. Accepts:

        - AbiStructure instance
        - Subinstances of pymatgen.
        - File paths
    """
    if isinstance(obj, AbiStructure):
        return obj

    if isinstance(obj, Structure):
        # Promote
        return AbiStructure(obj)

    if is_string(obj):
        # Handle file paths.
        if os.path.isfile(obj):

            if obj.endswith(".nc"):
                structure = structure_from_etsf_file(obj)
                #print(structure._sites)
            else:
                structure = read_structure(obj)

            # Promote
            return AbiStructure(structure)

    raise ValueError("Don't know how to convert object %s to an AbiStructure structure" % str(obj))


class AbiStructure(Structure, AbivarAble):
    """Patches the pymatgen structure adding the method to_abivars."""

    @staticmethod
    def asabistructure(obj):
        return asabistructure(obj)

    def __new__(cls, structure):
        new = structure
        new.__class__ = cls
        return new

    def __init__(self, structure):
        pass

    @classmethod
    def boxed_molecule(cls, pseudos, cart_coords, acell=3*(10,)):
        """
        Creates a molecule in a periodic box of lengths acell [Bohr]

        Args:
            pseudos:
                List of pseudopotentials
            cart_coords:
                Cartesian coordinates
            acell:
                Lengths of the box in *Bohr*
        """
        cart_coords = np.atleast_2d(cart_coords)

        molecule = Molecule([p.symbol for p in pseudos], cart_coords)

        l = ArrayWithUnit(acell, "bohr").to("ang")

        structure = molecule.get_boxed_structure(l[0], l[1], l[2])

        return cls(structure)

    @classmethod
    def boxed_atom(cls, pseudo, cart_coords=3*(0,), acell=3*(10,)):
        """
        Creates an atom in a periodic box of lengths acell [Bohr]

        Args:
            pseudo:
                Pseudopotential object.
            cart_coords:
                Cartesian coordinates
            acell:
                Lengths of the box in *Bohr*
        """
        return cls.boxed_molecule([pseudo], cart_coords, acell=acell)

    def to_abivars(self):
        "Returns a dictionary with the abinit variables."
        types_of_specie = self.types_of_specie
        natom = self.num_sites

        znucl_type = [specie.number for specie in types_of_specie]

        znucl_atoms = self.atomic_numbers

        typat = np.zeros(natom, np.int)
        for (atm_idx, site) in enumerate(self):
            typat[atm_idx] = types_of_specie.index(site.specie) + 1

        rprim = ArrayWithUnit(self.lattice.matrix, "ang").to("bohr")
        xred = np.reshape([site.frac_coords for site in self], (-1,3))

        return {
            "acell" : 3 * [1.0],
            "rprim" : rprim,
            "natom" : natom,
            "ntypat": len(types_of_specie),
            "typat" : typat,
            "xred"  : xred,
            "znucl" : znucl_type,
        }


class KSampling(AbivarAble):
    """
    Input variables defining the K-point sampling.
    """
    # Modes supported by the constructor.
    modes = Enum(('monkhorst', 'path', 'automatic',))

    def __init__(self,
                 mode              = "monkhorst",
                 num_kpts          = 0,
                 kpts              = ((1, 1, 1),),
                 kpt_shifts        = (0.5, 0.5, 0.5),
                 use_symmetries    = True,
                 use_time_reversal = True,
                 kpts_weights      = None,
                 labels            = None,
                 chksymbreak       = None,
                 comment           = None,
                 ):
        """
        Highly flexible constructor for KSampling objects.  The flexibility comes
        at the cost of usability and in general, it is recommended that you use
        the default constructor only if you know exactly what you are doing and
        requires the flexibility.  For most usage cases, the object be constructed
        far more easily using the convenience static constructors:

            #. gamma_only
            #. gamma_centered
            #. monkhorst
            #. monkhorst_automatic
            #. path

        and it is recommended that you use those.

        Args:
            mode:
                Mode for generating k-poits. Use one of the KSampling.modes
                enum types.
            num_kpts:
                Number of kpoints if mode is "automatic"
                Number of division for the sampling of the smallest segment if
                mode is "path".
                Not used for the other modes
            kpts:
                Number of divisions. Even when only a single specification is
                required, e.g. in the automatic scheme, the kpts should still
                be specified as a 2D array. e.g., [[20]] or [[2,2,2]].
            kpt_shifts:
                Shifts for Kpoints.
            use_symmetries:
                False if spatial symmetries should not be used to reduced the
                number of independent k-points.
            use_time_reversal:
                False if time-reversal symmetry should not be used to reduced
                the number of independent k-points.
            kpts_weights:
                Optional weights for kpoints. For explicit kpoints.
            labels:
                In path-mode, this should provide a list of labels for each kpt.
            chksymbreak:
                Abinit input variable: check whether the BZ sampling preserves
                the symmetry of the crystal.
            comment:
                String comment for Kpoints

        The default behavior of the constructor is monkhorst.
        """
        if mode not in KSampling.modes:
            raise ValueError("Unknown kpoint mode %s" % mode)

        super(KSampling, self).__init__()

        self.mode = mode
        self.comment = comment

        abivars = {}

        if mode in ("monkhorst",):
            assert num_kpts == 0
            ngkpt  = np.reshape(kpts, (3,)),
            shiftk = np.reshape(kpt_shifts, (-1,3))

            if use_symmetries and use_time_reversal: kptopt = 1
            if not use_symmetries and use_time_reversal: kptopt = 2
            if not use_symmetries and not use_time_reversal: kptopt = 3
            if use_symmetries and not use_time_reversal: kptopt = 4

            abivars.update({
                "ngkpt"      : ngkpt,
                "shiftk"     : shiftk,
                "nshiftk"    : len(shiftk),
                "kptopt"     : kptopt,
                "chksymbreak": chksymbreak,
            })

        elif mode in ("path",):

            if num_kpts <= 0:
                raise ValueError("For Path mode, num_kpts must be specified and >0")

            kptbounds = np.reshape(kpts, (-1,3,))
            #print("in path with kptbound: %s " % kptbounds)

            abivars.update({
                "ndivsm"   : num_kpts,
                "kptbounds": kptbounds,
                "kptopt"   : -len(kptbounds)+1,
                #"labels"  : labels,  # TODO
            })

        elif mode in ("automatic",):
            kpts = np.reshape(kpts, (-1,3))
            if len(kpts) != num_kpts:
                raise ValueError("For Automatic mode, num_kpts must be specified.")

            kptnrm = np.ones(num_kpts)

            abivars.update({
                "kptopt"     : 0,
                "kpt"        : kpts,
                "nkpt"       : num_kpts,
                "kptnrm"     : kptnrm,
                "wtk"        : kpts_weights,  # for iscf/=-2, wtk.
                "chksymbreak": chksymbreak,
            })

        else:
            raise ValueError("Unknown mode %s" % mode)

        self.abivars = abivars
        self.abivars["#comment"] = comment

    @property
    def is_homogeneous(self):
        return self.mode not in ["path",]

    @classmethod
    def gamma_only(cls):
        "Gamma-only sampling"
        return cls(kpt_shifts=(0.0,0.0,0.0), comment="Gamma-only sampling")

    @classmethod
    def gamma_centered(cls, kpts=(1, 1, 1), use_symmetries=True, use_time_reversal=True):
        """
        Convenient static constructor for an automatic Gamma centered Kpoint grid.

        Args:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors.
        use_symmetries:
            False if spatial symmetries should not be used to reduced the number of independent k-points.
        use_time_reversal:
            False if time-reversal symmetry should not be used to reduced the number of independent k-points.

        Returns:
            `KSampling` object.
        """
        return cls(kpts              = [kpts],
                   use_symmetries    = use_symmetries,
                   use_time_reversal = use_time_reversal,
                   comment           = "gamma-centered mode",
                   )

    @classmethod
    def monkhorst(cls, ngkpt, shiftk=(0.5, 0.5, 0.5), chksymbreak=None, use_symmetries=True,
                  use_time_reversal=True, comment=None):
        """
        Convenient static constructor for a Monkhorst-Pack mesh.

        Args:
            ngkpt:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors.
            shiftk:
                Shift to be applied to the kpoints. Defaults to (0,0,0).
            use_symmetries:
                Use spatial symmetries to reduce the number of k-points.
            use_time_reversal:
                Use time-reversal symmetry to reduce the number of k-points.

        Returns:
            `KSampling` object.
        """
        return cls(kpts              = [ngkpt],<commit_msg>added 0 shift to the gamma_centered kpoint sampling for ainit<commit_after>"""
Low-level objects providing an abstraction for the objects involved in the
calculation.
"""
from __future__ import division, print_function

import collections
import os
import abc
import numpy as np
import pymatgen.core.units as units

from pprint import pformat
from monty.design_patterns import singleton
from pymatgen.util.string_utils import is_string
from pymatgen.core.design_patterns import Enum, AttrDict
from pymatgen.core.units import ArrayWithUnit
from pymatgen.serializers.json_coders import MSONable
from pymatgen.symmetry.finder import SymmetryFinder
from pymatgen.core.structure import Structure, Molecule
from pymatgen.io.smartio import read_structure

from .netcdf import structure_from_etsf_file


class AbivarAble(object):
    """
    An AbivarAble object provides a method to_abivars that returns a dictionary
    with the abinit variables.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def to_abivars(self):
        """
        Returns a dictionary with the abinit variables.
        """

    def __str__(self):
        return pformat(self.to_abivars(), indent=1, width=80, depth=None)


@singleton
class MandatoryVariable(object):
    """
    Singleton used to tag mandatory variables, just because I can use
    the cool syntax: variable is MANDATORY!
    """


@singleton
class DefaultVariable(object):
    """
    Singleton used to tag variables that will have the default value
    """

MANDATORY = MandatoryVariable()
DEFAULT = DefaultVariable()


class SpinMode(collections.namedtuple('SpinMode', "mode nsppol nspinor nspden"),
               AbivarAble):
    """
    Different configurations of the electron density as implemented in abinit:
    One can use asspinmode to construct the object via SpinMode.asspinmode
    (string) where string can assume the values:

        - polarized
        - unpolarized
        - afm (anti-ferromagnetic)
        - spinor (non-collinear magnetism)
        - spinor_nomag (non-collinear, no magnetism)
    """
    @classmethod
    def asspinmode(cls, obj):
        """Converts obj into a SpinMode instance"""
        if isinstance(obj, cls):
            return obj
        else:
            # Assume a string with mode
            try:
                return _mode2spinvars[obj]
            except KeyError:
                raise KeyError("Wrong value for spin_mode: %s" % str(obj))

    def to_abivars(self):
        return {
            "nsppol": self.nsppol,
            "nspinor": self.nspinor,
            "nspden": self.nspden,
        }

# An handy Multiton
_mode2spinvars = {
    "unpolarized" : SpinMode("unpolarized" , 1, 1, 1),
    "polarized"   : SpinMode("polarized"   , 2, 1, 2),
    "afm"         : SpinMode("afm"         , 1, 1, 2),
    "spinor"      : SpinMode("spinor"      , 1, 2, 4),
    "spinor_nomag": SpinMode("spinor_nomag", 1, 2, 1),
}


class Smearing(AbivarAble, MSONable):
    """
    Variables defining the smearing technique. The preferred way to instanciate
    a Smearing object is via the class method Smearing.assmearing(string)
    """
    #: Mapping string_mode --> occopt
    _mode2occopt = {
        'nosmearing' : 1,
        'fermi_dirac': 3,
        'marzari4'   : 4,
        'marzari5'   : 5,
        'methfessel' : 6,
        'gaussian'   : 7,
    }

    def __init__(self, occopt, tsmear):
        self.occopt = occopt
        self.tsmear = tsmear

    def __str__(self):
        s = "occopt %d # %s Smearing\n" % (self.occopt, self.mode)
        if self.tsmear:
            s += 'tsmear %s' % self.tsmear
        return s

    def __eq__(self, other):
        return (self.occopt == other.occopt and
                np.allclose(self.tsmear, other.tsmear))

    def __ne__(self, other):
        return not self == other

    def __bool__(self):
        return self.mode != "nosmearing"

    # py2 old version
    __nonzero__ = __bool__

    @classmethod
    def assmearing(cls, obj):
        """
        Constructs an instance of Smearing from obj. Accepts obj in the form:

            * Smearing instance
            * "name:tsmear"  e.g. "gaussian:0.004"  (Hartree units)
            * "name:tsmear units" e.g. "gaussian:0.1 eV"
            * None --> no smearing
        """
        if obj is None:
            return Smearing.nosmearing()

        if isinstance(obj, cls):
            return obj

        # obj is a string
        obj, tsmear = obj.split(":")
        obj.strip()

        if obj == "nosmearing":
            return cls.nosmearing()
        else:
            occopt = cls._mode2occopt[obj]
            try:
                tsmear = float(tsmear)
            except ValueError:
                tsmear, unit = tsmear.split()
                tsmear = units.Energy(float(tsmear), unit).to("Ha")

            return cls(occopt, tsmear)

    @property
    def mode(self):
        for (mode_str, occopt) in self._mode2occopt.items():
            if occopt == self.occopt:
                return mode_str
        raise AttributeError("Unknown occopt %s" % self.occopt)

    @staticmethod
    def nosmearing():
        return Smearing(1, None)

    def to_abivars(self):
        if self.mode == "nosmearing":
            return {}
        else:
            return {"occopt": self.occopt,
                    "tsmear": self.tsmear,}

    @property
    def to_dict(self):
        """json friendly dict representation of Smearing"""
        return {"occopt": self.occopt, "tsmear": self.tsmear,
                "@module": self.__class__.__module__,
                "@class": self.__class__.__name__}

    @staticmethod
    def from_dict(d):
        return Smearing(d["occopt"], d["tsmear"])



class ElectronsAlgorithm(dict, AbivarAble):
    "Variables controlling the SCF/NSCF algorithm."
    # None indicates that we use abinit defaults.
    _default_vars = {
        "iprcell"  : None,
        "iscf"     : None,
        "diemac"   : None,
        "diemix"   : None,
        "diemixmag": None,
        "dielam"   : None,
        "diegap"   : None,
        "dielng"   : None,
        "diecut"   : None,
        "nstep"    : 70,
    }

    def __init__(self, *args, **kwargs):
        super(ElectronsAlgorithm, self).__init__(*args, **kwargs)

        for k in self:
            if k not in self._default_vars:
                raise ValueError("%s: No default value has been provided for "
                                 "key %s" % (self.__class__.__name__, k))

    def to_abivars(self):
        return self.copy()


class Electrons(AbivarAble):
    """
    The electronic degrees of freedom
    """
    def __init__(self,
                 spin_mode = "polarized",
                 smearing  = "fermi_dirac:0.1 eV",
                 algorithm = None,
                 nband     = None,
                 fband     = None,
                 charge    = 0.0,
                 comment   = None,
                 #occupancies = None,
                ):
        """
        Constructor for Electrons object.

        Args:
            comment:
                String comment for Electrons

            charge: float
                    Total charge of the system. Default is 0.
        """
        super(Electrons, self).__init__()

        self.comment = comment

        self.smearing = Smearing.assmearing(smearing)

        self.spin_mode = SpinMode.asspinmode(spin_mode)

        self.nband = nband
        self.fband = fband
        self.charge = charge

        self.algorithm = algorithm

        # FIXME
        if nband is None:
            self.fband = 4

    @property
    def nsppol(self):
        return self.spin_mode.nsppol

    @property
    def nspinor(self):
        return self.spin_mode.nspinor

    @property
    def nspden(self):
        return self.spin_mode.nspden

    #@property
    #def to_dict(self):
    #    "json friendly dict representation"
    #    d = {}
    #    d["@module"] = self.__class__.__module__
    #    d["@class"] = self.__class__.__name__
    #    raise NotImplementedError("")
    #    return d

    #@staticmethod
    #def from_dict(d):
    #    raise NotImplementedError("")

    def to_abivars(self):
        abivars = self.spin_mode.to_abivars()

        abivars.update({
            "nband"  : self.nband,
            "fband"  : self.fband,
            "charge" : self.charge,
        })

        if self.smearing:
            abivars.update(self.smearing.to_abivars())

        if self.algorithm:
            abivars.update(self.algorithm)

        abivars["#comment"] = self.comment
        return abivars


def asabistructure(obj):
    """
    Convert obj into an AbiStructure object. Accepts:

        - AbiStructure instance
        - Subinstances of pymatgen.
        - File paths
    """
    if isinstance(obj, AbiStructure):
        return obj

    if isinstance(obj, Structure):
        # Promote
        return AbiStructure(obj)

    if is_string(obj):
        # Handle file paths.
        if os.path.isfile(obj):

            if obj.endswith(".nc"):
                structure = structure_from_etsf_file(obj)
                #print(structure._sites)
            else:
                structure = read_structure(obj)

            # Promote
            return AbiStructure(structure)

    raise ValueError("Don't know how to convert object %s to an AbiStructure structure" % str(obj))


class AbiStructure(Structure, AbivarAble):
    """Patches the pymatgen structure adding the method to_abivars."""

    @staticmethod
    def asabistructure(obj):
        return asabistructure(obj)

    def __new__(cls, structure):
        new = structure
        new.__class__ = cls
        return new

    def __init__(self, structure):
        pass

    @classmethod
    def boxed_molecule(cls, pseudos, cart_coords, acell=3*(10,)):
        """
        Creates a molecule in a periodic box of lengths acell [Bohr]

        Args:
            pseudos:
                List of pseudopotentials
            cart_coords:
                Cartesian coordinates
            acell:
                Lengths of the box in *Bohr*
        """
        cart_coords = np.atleast_2d(cart_coords)

        molecule = Molecule([p.symbol for p in pseudos], cart_coords)

        l = ArrayWithUnit(acell, "bohr").to("ang")

        structure = molecule.get_boxed_structure(l[0], l[1], l[2])

        return cls(structure)

    @classmethod
    def boxed_atom(cls, pseudo, cart_coords=3*(0,), acell=3*(10,)):
        """
        Creates an atom in a periodic box of lengths acell [Bohr]

        Args:
            pseudo:
                Pseudopotential object.
            cart_coords:
                Cartesian coordinates
            acell:
                Lengths of the box in *Bohr*
        """
        return cls.boxed_molecule([pseudo], cart_coords, acell=acell)

    def to_abivars(self):
        "Returns a dictionary with the abinit variables."
        types_of_specie = self.types_of_specie
        natom = self.num_sites

        znucl_type = [specie.number for specie in types_of_specie]

        znucl_atoms = self.atomic_numbers

        typat = np.zeros(natom, np.int)
        for (atm_idx, site) in enumerate(self):
            typat[atm_idx] = types_of_specie.index(site.specie) + 1

        rprim = ArrayWithUnit(self.lattice.matrix, "ang").to("bohr")
        xred = np.reshape([site.frac_coords for site in self], (-1,3))

        return {
            "acell" : 3 * [1.0],
            "rprim" : rprim,
            "natom" : natom,
            "ntypat": len(types_of_specie),
            "typat" : typat,
            "xred"  : xred,
            "znucl" : znucl_type,
        }


class KSampling(AbivarAble):
    """
    Input variables defining the K-point sampling.
    """
    # Modes supported by the constructor.
    modes = Enum(('monkhorst', 'path', 'automatic',))

    def __init__(self,
                 mode              = "monkhorst",
                 num_kpts          = 0,
                 kpts              = ((1, 1, 1),),
                 kpt_shifts        = (0.5, 0.5, 0.5),
                 use_symmetries    = True,
                 use_time_reversal = True,
                 kpts_weights      = None,
                 labels            = None,
                 chksymbreak       = None,
                 comment           = None,
                 ):
        """
        Highly flexible constructor for KSampling objects.  The flexibility comes
        at the cost of usability and in general, it is recommended that you use
        the default constructor only if you know exactly what you are doing and
        requires the flexibility.  For most usage cases, the object be constructed
        far more easily using the convenience static constructors:

            #. gamma_only
            #. gamma_centered
            #. monkhorst
            #. monkhorst_automatic
            #. path

        and it is recommended that you use those.

        Args:
            mode:
                Mode for generating k-poits. Use one of the KSampling.modes
                enum types.
            num_kpts:
                Number of kpoints if mode is "automatic"
                Number of division for the sampling of the smallest segment if
                mode is "path".
                Not used for the other modes
            kpts:
                Number of divisions. Even when only a single specification is
                required, e.g. in the automatic scheme, the kpts should still
                be specified as a 2D array. e.g., [[20]] or [[2,2,2]].
            kpt_shifts:
                Shifts for Kpoints.
            use_symmetries:
                False if spatial symmetries should not be used to reduced the
                number of independent k-points.
            use_time_reversal:
                False if time-reversal symmetry should not be used to reduced
                the number of independent k-points.
            kpts_weights:
                Optional weights for kpoints. For explicit kpoints.
            labels:
                In path-mode, this should provide a list of labels for each kpt.
            chksymbreak:
                Abinit input variable: check whether the BZ sampling preserves
                the symmetry of the crystal.
            comment:
                String comment for Kpoints

        The default behavior of the constructor is monkhorst.
        """
        if mode not in KSampling.modes:
            raise ValueError("Unknown kpoint mode %s" % mode)

        super(KSampling, self).__init__()

        self.mode = mode
        self.comment = comment

        abivars = {}

        if mode in ("monkhorst",):
            assert num_kpts == 0
            ngkpt  = np.reshape(kpts, (3,)),
            shiftk = np.reshape(kpt_shifts, (-1,3))

            if use_symmetries and use_time_reversal: kptopt = 1
            if not use_symmetries and use_time_reversal: kptopt = 2
            if not use_symmetries and not use_time_reversal: kptopt = 3
            if use_symmetries and not use_time_reversal: kptopt = 4

            abivars.update({
                "ngkpt"      : ngkpt,
                "shiftk"     : shiftk,
                "nshiftk"    : len(shiftk),
                "kptopt"     : kptopt,
                "chksymbreak": chksymbreak,
            })

        elif mode in ("path",):

            if num_kpts <= 0:
                raise ValueError("For Path mode, num_kpts must be specified and >0")

            kptbounds = np.reshape(kpts, (-1,3,))
            #print("in path with kptbound: %s " % kptbounds)

            abivars.update({
                "ndivsm"   : num_kpts,
                "kptbounds": kptbounds,
                "kptopt"   : -len(kptbounds)+1,
                #"labels"  : labels,  # TODO
            })

        elif mode in ("automatic",):
            kpts = np.reshape(kpts, (-1,3))
            if len(kpts) != num_kpts:
                raise ValueError("For Automatic mode, num_kpts must be specified.")

            kptnrm = np.ones(num_kpts)

            abivars.update({
                "kptopt"     : 0,
                "kpt"        : kpts,
                "nkpt"       : num_kpts,
                "kptnrm"     : kptnrm,
                "wtk"        : kpts_weights,  # for iscf/=-2, wtk.
                "chksymbreak": chksymbreak,
            })

        else:
            raise ValueError("Unknown mode %s" % mode)

        self.abivars = abivars
        self.abivars["#comment"] = comment

    @property
    def is_homogeneous(self):
        return self.mode not in ["path",]

    @classmethod
    def gamma_only(cls):
        "Gamma-only sampling"
        return cls(kpt_shifts=(0.0,0.0,0.0), comment="Gamma-only sampling")

    @classmethod
    def gamma_centered(cls, kpts=(1, 1, 1), use_symmetries=True, use_time_reversal=True):
        """
        Convenient static constructor for an automatic Gamma centered Kpoint grid.

        Args:
            kpts:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors.
        use_symmetries:
            False if spatial symmetries should not be used to reduced the number of independent k-points.
        use_time_reversal:
            False if time-reversal symmetry should not be used to reduced the number of independent k-points.

        Returns:
            `KSampling` object.
        """
        return cls(kpts              = [kpts],
                   kpt_shifts        = (0.0, 0.0, 0.0),
                   use_symmetries    = use_symmetries,
                   use_time_reversal = use_time_reversal,
                   comment           = "gamma-centered mode",
                   )

    @classmethod
    def monkhorst(cls, ngkpt, shiftk=(0.5, 0.5, 0.5), chksymbreak=None, use_symmetries=True,
                  use_time_reversal=True, comment=None):
        """
        Convenient static constructor for a Monkhorst-Pack mesh.

        Args:
            ngkpt:
                Subdivisions N_1, N_2 and N_3 along reciprocal lattice vectors.
            shiftk:
                Shift to be applied to the kpoints. Defaults to (0,0,0).
            use_symmetries:
                Use spatial symmetries to reduce the number of k-points.
            use_time_reversal:
                Use time-reversal symmetry to reduce the number of k-points.

        Returns:
            `KSampling` object.
        """
        return cls(kpts              = [ngkpt],
{'fraction_non_alphanumeric': 0.08992277572329781, 'fraction_numerical': 0.006825103328257559, 'mean_word_length': 1.7298842042755345, 'pattern_counts': {'":': 36, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2957586', 'n_tokens_mistral': 11280, 'n_tokens_neox': 10646, 'n_words': 2986}