"""
Wrappers that provide a minimal interface to kimpy methods and objects

Daniel S. Karls
University of Minnesota
"""

import functools

import kimpy

from .exceptions import KIMModelNotFound, KIMModelInitializationError, KimpyError


def check_call(f, *args):
    """
    Given a function that returns either an integer error code or a
    tuple whose last element is an integer error code, call it with the
    requested arguments.  If a non-zero error code was returned, raise
    an exception.  Otherwise, pass along the rest of the objects
    returned by the function call.
    """

    def _check_error(error, msg):
        if error != 0 and error is not None:
            raise KimpyError('Calling "{}" failed.'.format(msg))

    ret = f(*args)

    if isinstance(ret, int):
        # Only an error code was returned
        _check_error(ret, f.__name__)
    else:
        # An error code plus other variables were returned
        error = ret[-1]
        _check_error(error, f.__name__)

        if len(ret[:-1]) == 1:
            # Pick the single remaining element out of the tuple
            return ret[0]
        else:
            # Return the tuple containing the rest of the elements
            return ret[:-1]


def check_call_wrapper(func):
    @functools.wraps(func)
    def myfunc(*args, **kwargs):
        return check_call(func, *args)

    return myfunc


# kimpy methods wrapped in ``check_error``
collections_create = functools.partial(check_call, kimpy.collections.create)
model_create = functools.partial(check_call, kimpy.model.create)
simulator_model_create = functools.partial(check_call, kimpy.simulator_model.create)
get_species_name = functools.partial(check_call, kimpy.species_name.get_species_name)

# kimpy attributes (here to avoid importing kimpy in higher-level modules)
collection_item_type_portableModel = kimpy.collection_item_type.portableModel


class ModelCollections:
    """
    KIM Portable Models and Simulator Models are installed/managed into
    different "collections".  In order to search through the different
    KIM API model collections on the system, a corresponding object must
    be instantiated.  For more on model collections, see the KIM API's
    install file:
    https://github.com/openkim/kim-api/blob/master/INSTALL
    """

    def __init__(self):
        self.collection = collections_create()

    def __del__(self):
        self.destroy()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, value, traceback):
        self.destroy()

    def get_item_type(self, model_name):
        try:
            model_type = check_call(self.collection.get_item_type, model_name)
        except KimpyError:
            msg = (
                "Could not find model {} installed in any of the KIM API model "
                "collections on this system.  See "
                "https://openkim.org/doc/usage/obtaining-models/ for instructions on "
                "installing models.".format(model_name)
            )
            raise KIMModelNotFound(msg)

        return model_type

    def destroy(self):
        if self.initialized:
            kimpy.collections.destroy(self.collection)
            del self.collection

    @property
    def initialized(self):
        return hasattr(self, "collection")


class PortableModel:
    """ Creates a KIM API Portable Model object and provides a minimal interface to it
    """

    def __init__(self, model_name, debug):
        self.model_name = model_name
        self.debug = debug

        # Create KIM API Model object
        units_accepted, self.kim_model = model_create(
            kimpy.numbering.zeroBased,
            kimpy.length_unit.A,
            kimpy.energy_unit.eV,
            kimpy.charge_unit.e,
            kimpy.temperature_unit.K,
            kimpy.time_unit.ps,
            self.model_name,
        )

        if not units_accepted:
            raise KIMModelInitializationError(
                "Requested units not accepted in kimpy.model.create"
            )

        if self.debug:
            l_unit, e_unit, c_unit, te_unit, ti_unit = self.kim_model.get_units()
            print("Length unit is: {}".format(l_unit))
            print("Energy unit is: {}".format(e_unit))
            print("Charge unit is: {}".format(c_unit))
            print("Temperature unit is: {}".format(te_unit))
            print("Time unit is: {}".format(ti_unit))
            print()

    def __del__(self):
        self.destroy()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, value, traceback):
        self.destroy()

    def get_model_supported_species_and_codes(self):
        """Get all of the supported species for this model and their
        corresponding integer codes that are defined in the KIM API

        Returns
        -------
        species : list of str
            Abbreviated chemical symbols of all species the mmodel
            supports (e.g. ["Mo", "S"])

        codes : list of int
            Integer codes used by the model for each species (order
            corresponds to the order of ``species``)
        """
        species = []
        codes = []
        num_kim_species = kimpy.species_name.get_number_of_species_names()

        for i in range(num_kim_species):
            species_name = get_species_name(i)
            species_support, code = self.get_species_support_and_code(species_name)

            if species_support:
                species.append(str(species_name))
                codes.append(code)

        return species, codes

    @check_call_wrapper
    def compute(self, compute_args_wrapped, release_GIL):
        return self.kim_model.compute(compute_args_wrapped.compute_args, release_GIL)

    @check_call_wrapper
    def get_species_support_and_code(self, species_name):
        return self.kim_model.get_species_support_and_code(species_name)

    def get_influence_distance(self):
        return self.kim_model.get_influence_distance()

    def get_neighbor_list_cutoffs_and_hints(self):
        return self.kim_model.get_neighbor_list_cutoffs_and_hints()

    def compute_arguments_create(self):
        return ComputeArguments(self, self.debug)

    def compute_arguments_destroy(self, compute_args_wrapped):
        compute_args_wrapped.destroy()

    def destroy(self):
        if self.initialized:
            kimpy.model.destroy(self.kim_model)
            del self.kim_model

    @property
    def initialized(self):
        return hasattr(self, "kim_model")


class ComputeArguments:
    """
    Creates a KIM API ComputeArguments object from a KIM Portable Model object and
    configures it for ASE.  A ComputeArguments object is associated with a KIM Portable
    Model and is used to inform the KIM API of what the model can compute.  It is also
    used to register the data arrays that allow the KIM API to pass the atomic
    coordinates to the model and retrieve the corresponding energy and forces, etc.
    """

    def __init__(self, kim_model_wrapped, debug):
        self.kim_model_wrapped = kim_model_wrapped
        self.debug = debug

        # Create KIM API ComputeArguments object
        self.compute_args = check_call(
            self.kim_model_wrapped.kim_model.compute_arguments_create
        )

        # Check compute arguments
        kimpy_arg_name = kimpy.compute_argument_name
        num_arguments = kimpy_arg_name.get_number_of_compute_argument_names()
        if self.debug:
            print("Number of compute_args: {}".format(num_arguments))

        for i in range(num_arguments):
            name = check_call(kimpy_arg_name.get_compute_argument_name, i)
            dtype = check_call(kimpy_arg_name.get_compute_argument_data_type, name)

            arg_support = self.get_argument_support_status(name)

            if self.debug:
                print(
                    "Compute Argument name {:21} is of type {:7} and has support "
                    "status {}".format(*[str(x) for x in [name, dtype, arg_support]])
                )

            # See if the model demands that we ask it for anything other than energy and
            # forces.  If so, raise an exception.
            if arg_support == kimpy.support_status.required:
                if (
                    name != kimpy.compute_argument_name.partialEnergy
                    and name != kimpy.compute_argument_name.partialForces
                ):
                    raise KIMModelInitializationError(
                        "Unsupported required ComputeArgument {}".format(name)
                    )

        # Check compute callbacks
        callback_name = kimpy.compute_callback_name
        num_callbacks = callback_name.get_number_of_compute_callback_names()
        if self.debug:
            print()
            print("Number of callbacks: {}".format(num_callbacks))

        for i in range(num_callbacks):
            name = check_call(callback_name.get_compute_callback_name, i)

            support_status = self.get_callback_support_status(name)

            if self.debug:
                print(
                    "Compute callback {:17} has support status {}".format(
                        str(name), support_status
                    )
                )

            # Cannot handle any "required" callbacks
            if support_status == kimpy.support_status.required:
                raise KIMModelInitializationError(
                    "Unsupported required ComputeCallback: {}".format(name)
                )

    @check_call_wrapper
    def set_argument_pointer(self, compute_arg_name, data_object):
        return self.compute_args.set_argument_pointer(compute_arg_name, data_object)

    @check_call_wrapper
    def get_argument_support_status(self, name):
        return self.compute_args.get_argument_support_status(name)

    @check_call_wrapper
    def get_callback_support_status(self, name):
        return self.compute_args.get_callback_support_status(name)

    @check_call_wrapper
    def set_callback(self, compute_callback_name, callback_function, data_object):
        return self.compute_args.set_callback(
            compute_callback_name, callback_function, data_object
        )

    @check_call_wrapper
    def set_callback_pointer(self, compute_callback_name, callback, data_object):
        return self.compute_args.set_callback_pointer(
            compute_callback_name, callback, data_object
        )

    def update(
        self, num_particles, species_code, particle_contributing, coords, energy, forces
    ):
        """ Register model input and output in the kim_model object."""
        compute_arg_name = kimpy.compute_argument_name
        set_argument_pointer = self.set_argument_pointer

        set_argument_pointer(compute_arg_name.numberOfParticles, num_particles)
        set_argument_pointer(compute_arg_name.particleSpeciesCodes, species_code)
        set_argument_pointer(
            compute_arg_name.particleContributing, particle_contributing
        )
        set_argument_pointer(compute_arg_name.coordinates, coords)
        set_argument_pointer(compute_arg_name.partialEnergy, energy)
        set_argument_pointer(compute_arg_name.partialForces, forces)

        if self.debug:
            print("Debug: called update_kim")
            print()

    @check_call_wrapper
    def destroy(self):
        return self.kim_model_wrapped.kim_model.compute_arguments_destroy(
            self.compute_args
        )


class SimulatorModel:
    """ Creates a KIM API Simulator Model object and provides a minimal
    interface to it.  This is only necessary in this package in order to
    extract any information about a given simulator model because it is
    generally embedded in a shared object.
    """

    def __init__(self, model_name):
        # Create a KIM API Simulator Model object for this model
        self.model_name = model_name
        self.simulator_model = simulator_model_create(self.model_name)

        # Need to close template map in order to access simulator model metadata
        self.simulator_model.close_template_map()

    def __del__(self):
        self.destroy()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, value, traceback):
        self.destroy()

    def destroy(self):
        if self.initialized:
            kimpy.simulator_model.destroy(self.simulator_model)
            del self.simulator_model

    @property
    def simulator_name(self):
        simulator_name, _ = self.simulator_model.get_simulator_name_and_version()
        return simulator_name

    @property
    def num_supported_species(self):
        num_supported_species = self.simulator_model.get_number_of_supported_species()
        if num_supported_species == 0:
            raise KIMModelInitializationError(
                "Unable to determine supported species of simulator model {}.".format(
                    self.model_name
                )
            )
        else:
            return num_supported_species

    @property
    def supported_species(self):
        supported_species = []
        for spec_code in range(self.num_supported_species):
            species = check_call(self.simulator_model.get_supported_species, spec_code)
            supported_species.append(species)

        return tuple(supported_species)

    @property
    def num_metadata_fields(self):
        return self.simulator_model.get_number_of_simulator_fields()

    @property
    def metadata(self):
        sm_metadata_fields = {}
        for field in range(self.num_metadata_fields):
            extent, field_name = check_call(
                self.simulator_model.get_simulator_field_metadata, field
            )
            sm_metadata_fields[field_name] = []
            for ln in range(extent):
                field_line = check_call(
                    self.simulator_model.get_simulator_field_line, field, ln
                )
                sm_metadata_fields[field_name].append(field_line)

        return sm_metadata_fields

    @property
    def supported_units(self):
        try:
            supported_units = self.metadata["units"][0]
        except (KeyError, IndexError):
            raise KIMModelInitializationError(
                "Unable to determine supported units of simulator model {}.".format(
                    self.model_name
                )
            )

        return supported_units

    @property
    def atom_style(self):
        """
        See if a 'model-init' field exists in the SM metadata and, if
        so, whether it contains any entries including an "atom_style"
        command.  This is specific to LAMMPS SMs and is only required
        for using the LAMMPSrun calculator because it uses
        lammps.inputwriter to create a data file.  All other content in
        'model-init', if it exists, is ignored.
        """
        atom_style = None
        for ln in self.metadata.get("model-init", []):
            if ln.find("atom_style") != -1:
                atom_style = ln.split()[1]

        return atom_style

    @property
    def model_defn(self):
        return self.metadata["model-defn"]

    @property
    def initialized(self):
        return hasattr(self, "simulator_model")