"""Convenience wrappers for the swig interface"""

import logging
import warnings
from typing import Any

import amici
import amici.amici as amici_swig
from amici.amici import (
    _get_ptr,
    AmiciExpData,
    AmiciExpDataVector,
    AmiciModel,
    AmiciSolver,
)
from . import numpy
from .logging import get_logger

logger = get_logger(__name__, log_level=logging.DEBUG)


__all__ = [
    "run_simulation",
    "run_simulations",
    "read_solver_settings_from_hdf5",
    "write_solver_settings_to_hdf5",
    "set_model_settings",
    "get_model_settings",
]


def run_simulation(
    model: AmiciModel,
    solver: AmiciSolver,
    edata: AmiciExpData | None = None,
) -> "numpy.ReturnDataView":
    """
    Convenience wrapper around :py:func:`amici.amici.run_simulation`
    (generated by swig)

    :param model:
        Model instance

    :param solver:
        Solver instance, must be generated from
        :py:meth:`amici.amici.Model.create_solver`

    :param edata:
        ExpData instance (optional)

    :returns:
        ReturnData object with simulation results
    """
    if (
        model.ne > 0
        and solver.get_sensitivity_method()
        == amici_swig.SensitivityMethod.adjoint
        and solver.get_sensitivity_order() == amici_swig.SensitivityOrder.first
    ):
        warnings.warn(
            "Adjoint sensitivity analysis for models with discontinuous right hand sides (events/piecewise functions) has not been thoroughly tested. "
            "Sensitivities might be wrong. Tracked at https://github.com/AMICI-dev/AMICI/issues/18. "
            "Adjoint sensitivity analysis may work if the location of the discontinuity is not parameter-dependent, but we still recommend testing accuracy of gradients.",
            stacklevel=1,
        )

    rdata = amici_swig.run_simulation(
        _get_ptr(solver), _get_ptr(edata), _get_ptr(model)
    )
    _log_simulation(rdata)
    if solver.get_return_data_reporting_mode() == amici.RDataReporting.full:
        _ids_and_names_to_rdata(rdata, model)
    return numpy.ReturnDataView(rdata)


def run_simulations(
    model: AmiciModel,
    solver: AmiciSolver,
    edata_list: AmiciExpDataVector,
    failfast: bool = True,
    num_threads: int = 1,
) -> list["numpy.ReturnDataView"]:
    """
    Convenience wrapper for loops of amici.runAmiciSimulation

    :param model: Model instance
    :param solver: Solver instance, must be generated from Model.getSolver()
    :param edata_list: list of ExpData instances
    :param failfast: returns as soon as an integration failure is encountered
    :param num_threads: number of threads to use (only used if compiled
        with openmp)

    :returns: list of simulation results
    """
    if (
        model.ne > 0
        and solver.get_sensitivity_method()
        == amici_swig.SensitivityMethod.adjoint
        and solver.get_sensitivity_order() == amici_swig.SensitivityOrder.first
    ):
        warnings.warn(
            "Adjoint sensitivity analysis for models with discontinuous right hand sides (events/piecewise functions) has not been thoroughly tested. "
            "Sensitivities might be wrong. Tracked at https://github.com/AMICI-dev/AMICI/issues/18. "
            "Adjoint sensitivity analysis may work if the location of the discontinuity is not parameter-dependent, but we still recommend testing accuracy of gradients.",
            stacklevel=1,
        )

    edata_ptr_vector = amici_swig.ExpDataPtrVector(edata_list)
    rdata_ptr_list = amici_swig.run_simulations(
        _get_ptr(solver),
        edata_ptr_vector,
        _get_ptr(model),
        failfast,
        num_threads,
    )
    for rdata in rdata_ptr_list:
        _log_simulation(rdata)
        if (
            solver.get_return_data_reporting_mode()
            == amici.RDataReporting.full
        ):
            _ids_and_names_to_rdata(rdata, model)

    return [numpy.ReturnDataView(r) for r in rdata_ptr_list]


def read_solver_settings_from_hdf5(
    file: str, solver: AmiciSolver, location: str | None = "solverSettings"
) -> None:
    """
    Convenience wrapper for :py:func:`amici.read_solver_settings_from_hdf5`

    :param file: hdf5 filename
    :param solver: Solver instance to which settings will be transferred
    :param location: location of solver settings in hdf5 file
    """
    amici_swig.read_solver_settings_from_hdf5(file, _get_ptr(solver), location)


def write_solver_settings_to_hdf5(
    solver: AmiciSolver,
    file: str | object,
    location: str | None = "solverSettings",
) -> None:
    """
    Convenience wrapper for :py:func:`amici.amici.writeSolverSettingsToHDF5`

    :param file: hdf5 filename, can also be an object created by
        :py:func:`amici.amici.createOrOpenForWriting`
    :param solver: Solver instance from which settings will be stored
    :param location: location of solver settings in hdf5 file
    """
    amici_swig.write_solver_settings_to_hdf5(_get_ptr(solver), file, location)


# Values are suffixes of `get[...]` and `set[...]` `amici.Model` methods.
# If either the getter or setter is not named with this pattern, then the value
# is a tuple where the first and second elements are the getter and setter
# methods, respectively.
model_instance_settings = [
    # `set_parameter_{list,scale}` will clear initial state sensitivities, so
    #  `set_parameter_{list,scale}` has to be called first.
    "parameter_list",
    "parameter_scale",  # getter returns a SWIG object
    "add_sigma_residuals",
    "always_check_finite",
    "fixed_parameters",
    "initial_state",
    (
        "get_initial_state_sensitivities",
        "set_unscaled_initial_state_sensitivities",
    ),
    "minimum_sigma_residuals",
    ("n_max_event", "set_n_max_event"),
    "parameters",
    "reinitialization_state_idxs",
    "reinitialize_fixed_parameter_initial_states",
    "state_is_non_negative",
    "steady_state_computation_mode",
    "steady_state_sensitivity_mode",
    ("t0", "set_t0"),
    ("t0_preeq", "set_t0_preeq"),
    "timepoints",
    "steadystate_mask",
]


def get_model_settings(
    model: AmiciModel,
) -> dict[str, Any]:
    """Get model settings that are set independently of the compiled model.

    :param model: The AMICI model instance.

    :returns: Keys are AMICI model attributes, values are attribute values.
    """
    settings = {}
    for setting in model_instance_settings:
        getter = setting[0] if isinstance(setting, tuple) else f"get_{setting}"

        if (
            getter == "get_initial_state"
            and not model.has_custom_initial_state()
        ):
            settings[setting] = []
            continue
        if (
            getter == "get_initial_state_sensitivities"
            and not model.has_custom_initial_state_sensitivities()
        ):
            settings[setting] = []
            continue

        settings[setting] = getattr(model, getter)()
        # TODO `amici.Model.get_parameter_scale` returns a SWIG object instead
        # of a Python list/tuple.
        if setting == "parameter_scale":
            settings[setting] = tuple(settings[setting])
    return settings


def set_model_settings(
    model: AmiciModel,
    settings: dict[str, Any],
) -> None:
    """Set model settings.

    :param model: The AMICI model instance.
    :param settings: Keys are callable attributes (setters) of an AMICI model,
        values are provided to the setters.
    """
    for setting, value in settings.items():
        setter = setting[1] if isinstance(setting, tuple) else f"set_{setting}"
        getattr(model, setter)(value)


def _log_simulation(rdata: amici_swig.ReturnData):
    """Extension warnings to Python logging."""
    amici_severity_to_logging = {
        amici_swig.LogSeverity_debug: logging.DEBUG,
        amici_swig.LogSeverity_warning: logging.WARNING,
        amici_swig.LogSeverity_error: logging.ERROR,
    }
    for msg in rdata.messages:
        condition = f"[{rdata.id}]" if rdata.id else ""
        logger.log(
            amici_severity_to_logging[msg.severity],
            f"{condition}[{msg.identifier}] {msg.message}",
        )


def _ids_and_names_to_rdata(
    rdata: amici_swig.ReturnData, model: amici_swig.Model
):
    """Copy entity IDs and names from a Model to ReturnData."""
    for entity_type in (
        "state",
        "observable",
        "expression",
        "parameter",
        "fixed_parameter",
    ):
        for name_or_id in ("ids", "names"):
            names_or_ids = getattr(model, f"get_{entity_type}_{name_or_id}")()
            setattr(
                rdata,
                f"{entity_type.lower()}_{name_or_id.lower()}",
                names_or_ids,
            )
    rdata.state_ids_solver = model.get_state_ids_solver()
    rdata.state_names_solver = model.get_state_names_solver()
