"""
A collection of wrappers around commonly employed force fields.
"""
import abc
from enum import Enum

from propertyestimator import unit
from propertyestimator.utils.serialization import TypedBaseModel


class ForceFieldSource(TypedBaseModel):
    """A helper object to define the source of a force field
    and any associated meta data, such as version, file paths,
    or generation options.

    Notes
    -----
    It is likely that this class and classes based off of it will
    not be permanent fixtures of the framework, but rather will
    exist until the force fields can be stored in a uniform format /
    object model.
    """

    @abc.abstractmethod
    def __getstate__(self):
        pass

    @abc.abstractmethod
    def __setstate__(self, state):
        pass


class SmirnoffForceFieldSource(ForceFieldSource):
    """A wrapper around force fields based on the
    `SMIRks Native Open Force Field (SMIRNOFF) specification
    <https://open-forcefield-toolkit.readthedocs.io/en/latest/smirnoff.html>`_.
    """

    def __init__(self, inner_xml=None):
        """Constructs a new SmirnoffForceFieldSource object

        Parameters
        ----------
        inner_xml: str, optional
            A string containing the xml representation of the force field.
        """
        self._inner_xml = inner_xml

    def to_force_field(self):
        """Returns the SMIRNOFF force field created from this source.

        Returns
        -------
        openforcefield.typing.engines.smirnoff.ForceField
            The created force field.
        """
        from openforcefield.typing.engines import smirnoff
        return smirnoff.ForceField(self._inner_xml)

    @classmethod
    def from_object(cls, force_field):
        """Creates a new `SmirnoffForceFieldSource` from an existing
        `ForceField` object

        Notes
        -----
        All cosmetic attributes will be discarded.

        Parameters
        ----------
        force_field: openforcefield.typing.engines.smirnoff.ForceField
            The existing force field.

        Returns
        -------
        SmirnoffForceFieldSource
            The created object.
        """

        return cls(force_field.to_string('XML', True))

    @classmethod
    def from_path(cls, file_path):
        """Creates a new `SmirnoffForceFieldSource` from the file path to a
        `ForceField` object.

        Notes
        -----
        All cosmetic attributes will be discarded.

        Parameters
        ----------
        file_path: str
            The file path to the force field object. This may also be the
            name of a file which can be loaded via an entry point.

        Returns
        -------
        SmirnoffForceFieldSource
            The created object.
        """

        from openforcefield.typing.engines.smirnoff import ForceField

        force_field = ForceField(file_path, allow_cosmetic_attributes=True)
        return cls.from_object(force_field)

    def __getstate__(self):
        return {'inner_xml': self._inner_xml}

    def __setstate__(self, state):
        self._inner_xml = state['inner_xml']


class TLeapForceFieldSource(ForceFieldSource):
    """A wrapper around Amber force fields which may be
    applied via the `tleap` software package.

    Notes
    -----
    Currently this only supports force fields which are installed
    alongside `tleap`.
    """

    @property
    def leap_source(self):
        """list of str: The parameter file which should be sourced by `leap`
            when applying the force field.
        """
        return self._leap_source

    @property
    def cutoff(self):
        """unit.Quantity: The non-bonded interaction cutoff.
        """
        return self._cutoff

    def __init__(self, leap_source='leaprc.gaff2', cutoff=9.0*unit.angstrom):
        """Constructs a new TLeapForceFieldSource object

        Parameters
        ----------
        leap_source: str
            The parameter file which should be sourced by `leap`
            when applying the force field. Currently only
            `'leaprc.gaff'` and `'leaprc.gaff2'` are supported.
        cutoff: unit.Quantity
            The non-bonded interaction cutoff.

        Examples
        --------
        To create a source for the GAFF force field with tip3p water:

        >>> amber_gaff_source = TLeapForceFieldSource('leaprc.gaff')

        To create a source for the GAFF 2 force field with tip3p water:

        >>> amber_gaff_2_source = TLeapForceFieldSource('leaprc.gaff2')
        """

        if leap_source is not None:
            assert leap_source == 'leaprc.gaff2' or leap_source == 'leaprc.gaff'

        self._leap_source = leap_source
        self._cutoff = cutoff

    def __getstate__(self):
        return {
            'leap_source': self._leap_source,
            'cutoff': self._cutoff
        }

    def __setstate__(self, state):
        self._leap_source = state['leap_source']
        self._cutoff = state['cutoff']


class LigParGenForceFieldSource(ForceFieldSource):
    """A wrapper and the OPLSAAM force field which can be applied
    via the `LigParGen server <http://zarbi.chem.yale.edu/ligpargen/>`_.

    References
    ----------
    [1] Potential energy functions for atomic-level simulations of water and organic and
        biomolecular systems. Jorgensen, W. L.; Tirado-Rives, J. Proc. Nat. Acad. Sci.
        USA 2005, 102, 6665-6670
    [2] 1.14*CM1A-LBCC: Localized Bond-Charge Corrected CM1A Charges for Condensed-Phase
        Simulations. Dodda, L. S.; Vilseck, J. Z.; Tirado-Rives, J.; Jorgensen, W. L.
        J. Phys. Chem. B, 2017, 121 (15), pp 3864-3870
    [3] LigParGen web server: An automatic OPLS-AA parameter generator for organic ligands.
        Dodda, L. S.;Cabeza de Vaca, I.; Tirado-Rives, J.; Jorgensen, W. L.
        Nucleic Acids Research, Volume 45, Issue W1, 3 July 2017, Pages W331-W336
    """

    class ChargeModel(Enum):

        CM1A_1_14_LBCC = '1.14*CM1A-LBCC'
        CM1A_1_14 = '1.14*CM1A'

    @property
    def preferred_charge_model(self):
        """ChargeModel: The preferred charge model to apply. In some cases
        the preferred charge model may not be applicable (e.g. 1.14*CM1A-LBCC
        may only be applied to neutral molecules) and so another model may be
        applied in its place.
        """
        return self._preferred_charge_model

    @property
    def cutoff(self):
        """unit.Quantity: The non-bonded interaction cutoff.
        """
        return self._cutoff

    def __init__(self, preferred_charge_model=ChargeModel.CM1A_1_14_LBCC, cutoff=9.0*unit.angstrom):
        """Constructs a new LigParGenForceFieldSource object

        Parameters
        ----------
        preferred_charge_model: ChargeModel
            The preferred charge model to apply. In some cases
            the preferred charge model may not be applicable
            (e.g. 1.14*CM1A-LBCC may only be applied to neutral
            molecules) and so another model may be applied in its
            place.
        cutoff: unit.Quantity
            The non-bonded interaction cutoff.
        """
        self._preferred_charge_model = preferred_charge_model
        self._cutoff = cutoff

    def __getstate__(self):
        return {
            'preferred_charge_model': self._preferred_charge_model,
            'cutoff': self._cutoff
        }

    def __setstate__(self, state):
        self._preferred_charge_model = state['preferred_charge_model']
        self._cutoff = state['cutoff']
