import os
from typing import Any, ClassVar

import COPASI
from basico import (
    get_reactions,
    get_species,
    load_model,
    run_steadystate,
    run_time_course,
)
from pandas import DataFrame
from process_bigraph import Process, ProcessTypes, Step

from pbest.registry.utils import model_path_resolution


def _set_initial_concentrations(changes: Any, dm: Any) -> None:
    """
    changes: iterable of (species_name, value) pairs
    dm: COPASI DataModel as returned by basico.load_model
    """
    model = dm.getModel()
    if not isinstance(model, COPASI.CModel):
        error_msg = f"Model {model!r} is not a COPASI model."
        raise TypeError(error_msg)

    references = COPASI.ObjectStdVector()

    for name, value in changes:
        species = model.getMetabolite(name)
        if species is None:
            print(f"Species {name} not found in model")
            continue
        if not isinstance(species, COPASI.CMetab):
            error_msg = f"Model {model!r} is not a COPASI model."
            raise TypeError(error_msg)
        species.setInitialConcentration(float(value))
        references.append(species.getInitialConcentrationReference())

    if len(references) > 0:
        model.updateInitialValues(references)


def _get_transient_concentration(name: Any, dm: Any) -> float | None:
    """
    Return the *current* concentration (not initial) of a species.
    """
    model = dm.getModel()
    if not isinstance(model, COPASI.CModel):
        error_msg = f"Model {model!r} is not a COPASI model."
        raise TypeError(error_msg)

    species = model.getMetabolite(name)
    if species is None:
        print(f"Species {name} not found in model")
        return None
    if not isinstance(species, COPASI.CMetab):
        error_msg = f"Model {model!r} is not a COPASI model."
        raise TypeError(error_msg)
    return float(species.getConcentration())


class BaseCopasi:
    cmodel = None
    dm = None
    species_ids: Any
    reaction_ids: Any
    sbml_to_name: dict[Any, Any]
    config: dict[str, Any]

    def interpret_sbml(self) -> None:
        model_source = self.config["model_source"]

        # ---- Load COPASI model ----
        self.dm = load_model(model_path_resolution(model_source))
        if self.dm is None:
            error_msg = (
                f"load_model({model_source!r}) returned None. "
                f"Check that the file exists and is a valid COPASI/SBML model."
            )
            raise RuntimeError(error_msg)

        self.cmodel = self.dm.getModel()

        spec_df = get_species(model=self.dm)

        # External canonical IDs: SBML IDs
        self.species_ids = spec_df["sbml_id"].tolist()

        # Mapping: SBML ID -> COPASI display name (index)
        self.sbml_to_name = {spec_df.loc[name, "sbml_id"]: name for name in spec_df.index}

        rxn_df = get_reactions(model=self.dm)
        # These are typically SBML reaction ids already
        self.reaction_ids = rxn_df.index.tolist()

    def get_concentrations_from_sbml(self) -> dict[str, Any]:
        return {
            "species_concentrations": {
                sbml_id: _get_transient_concentration(
                    name=self.sbml_to_name[sbml_id],  # COPASI name
                    dm=self.dm,
                )
                for sbml_id in self.species_ids
            }
        }


class CopasiUTCStep(Step, BaseCopasi):
    config_schema: ClassVar[dict[str, str]] = {
        "model_source": "string",
        "time": "float",
        "n_points": "integer",
        "output_dir": "string",
    }

    def initialize(self, config: Any = None) -> None:
        self.interpret_sbml()

        # Simulation parameters
        self.interval = float(self.config.get("time", 1.0))
        self.n_points = int(self.config.get("n_points", 2))  # <-- NEW
        if self.n_points < 2:
            error_msg = f"n_points={self.n_points} must be >= 2."
            raise ValueError(error_msg)

        self.intervals = self.n_points - 1  # COPASI requires this

    def initial_state(self) -> dict[str, Any]:
        return self.get_concentrations_from_sbml()

    def inputs(self) -> dict[str, str]:
        return {
            "species_concentrations": "map[float]",
            "species_counts": "map[float]",
        }

    def outputs(self) -> dict[str, str]:
        return {
            "result": "numeric_result",
        }

    def update(self, inputs: Any) -> dict[str, Any]:
        # Apply incoming concentrations
        spec_data = inputs.get("counts", {}) or {}
        changes = [(name, float(value)) for name, value in spec_data.items() if name in self.species_ids]

        if changes:
            _set_initial_concentrations(changes, self.dm)

        # --- Run COPASI time course with intervals = n_points - 1 ---
        tc: DataFrame = run_time_course(
            start_time=0.0,
            duration=self.config["time"],
            intervals=self.intervals,
            update_model=True,
            use_sbml_id=True,
            model=self.dm,
        )
        if "output_dir" in self.config and self.config["output_dir"] != "":
            tc.to_csv(os.path.join(self.config["output_dir"], "results.csv"), index=False)

        # Time series
        time_list = tc.index.to_list()

        result = {
            "time": time_list,
            "columns": [c for c in tc.columns],  # noqa: C416
            "values": tc.values.tolist(),
            # "n_spacial_dimensions": tc.shape,
        }

        return {"result": result}


class CopasiSteadyStateStep(Step, BaseCopasi):
    config_schema: ClassVar[dict[str, str]] = {
        "model_source": "string",
        "time": "float",  # kept for symmetry, not used
    }

    def initialize(self, config: Any = None) -> None:
        self.interpret_sbml()

    # ------------------------------------------------
    # initial state (SBML IDs externally)
    # ------------------------------------------------
    def initial_state(self) -> dict[str, Any]:
        return self.get_concentrations_from_sbml()

    # ------------------------------------------------
    # ports
    # ------------------------------------------------
    def inputs(self) -> dict[str, str]:
        # Externally everything uses SBML IDs
        return {
            "species_concentrations": "map[float]",  # SBML IDs
            "counts": "map[float]",  # SBML IDs
        }

    def outputs(self) -> dict[str, str]:
        # Match TelluriumSteadyStateStep: nested results
        return {
            "results": "any",
        }

    # ------------------------------------------------
    # steady-state update
    # ------------------------------------------------
    def update(self, inputs: Any) -> dict[str, Any]:
        # 1) Prefer counts, otherwise concentrations (keys are SBML IDs)
        spec_data = inputs.get("counts") or inputs.get("concentrations") or {}

        # Convert SBML IDs -> COPASI names for internal set
        changes = []
        for sbml_id, value in spec_data.items():
            name = self.sbml_to_name.get(sbml_id)
            if name is not None:
                changes.append((name, float(value)))

        if changes:
            _set_initial_concentrations(changes, self.dm)

        # 2) Run COPASI steady-state task
        # (use_sbml_id affects task I/O naming, but we read from get_species anyway)
        run_steadystate(
            update_model=True,
            use_sbml_id=True,
            model=self.dm,
        )

        # 3) Read back steady-state species concentrations (SBML IDs externally)
        spec_df = get_species(model=self.dm)
        # spec_df is indexed by COPASI name, with 'sbml_id' and 'concentration' columns
        species_conc_ss = {}
        for name in spec_df.index:
            sbml_id = spec_df.loc[name, "sbml_id"]
            if sbml_id in self.species_ids:
                species_conc_ss[sbml_id] = float(spec_df.loc[name, "concentration"])

        # 4) Steady-state reaction fluxes
        rxn_df = get_reactions(model=self.dm)
        reaction_fluxes_ss = {rid: float(rxn_df.loc[rid, "flux"]) for rid in self.reaction_ids if rid in rxn_df.index}

        # 5) Package as one-point "time series" (t = 0.0) to match Tellurium
        time_list = [0.0]

        species_json = {sid: [val] for sid, val in species_conc_ss.items()}
        flux_json = {rid: [val] for rid, val in reaction_fluxes_ss.items()}

        results = {
            "time": time_list,
            "species_concentrations": species_json,  # SBML IDs as keys
            "fluxes": flux_json,
        }

        return {"results": results}


class CopasiUTCProcess(Process, BaseCopasi):
    config_schema: ClassVar[dict[str, str]] = {
        "model_source": "string",
        "time": "float",
        "intervals": "integer",
    }

    def initialize(self, config: Any = None) -> None:
        self.interpret_sbml()

        # ---- Sim parameters ----
        self.time = float(self.config.get("time", 1.0))
        self.intervals = int(self.config.get("intervals", 10))

    # -----------------------------------------------------------------
    # initial state
    # -----------------------------------------------------------------
    def initial_state(self) -> dict[str, Any]:
        return self.get_concentrations_from_sbml()

    # -----------------------------------------------------------------
    # I/O schema
    # -----------------------------------------------------------------
    def inputs(self) -> dict[str, str]:
        return {
            "species_concentrations": "map[float]",  # SBML IDs
            "species_counts": "map[float]",  # SBML IDs
        }

    def outputs(self) -> dict[str, str]:
        return {
            "species_concentrations": "map[float]",  # SBML IDs
            "fluxes": "map[float]",
            "time": "list[float]",
        }

    # -----------------------------------------------------------------
    # update
    # -----------------------------------------------------------------
    def update(self, inputs: dict[str, Any], interval: float) -> dict[str, Any]:
        # --- 1) Determine incoming species map (SBML IDs)
        incoming = inputs.get("species_counts") or inputs.get("species_concentrations") or {}

        # Convert SBML IDs → COPASI names for internal setting
        changes = []
        for sbml_id, value in incoming.items():
            name = self.sbml_to_name.get(sbml_id)
            if name is not None:
                changes.append((name, float(value)))

        if changes:
            _set_initial_concentrations(changes, self.dm)

        # --- 2) Run time course with SBML-ID columns ----
        tc = run_time_course(
            start_time=0.0,
            duration=interval,
            intervals=self.intervals,
            update_model=True,
            use_sbml_id=True,  # <-- critical
            model=self.dm,
        )

        # Extract time points
        time = tc["Time"].tolist() if "Time" in tc.columns else []

        # --- 3) Read back final state: export SBML IDs ----
        species_concentrations = {
            sbml_id: _get_transient_concentration(name=self.sbml_to_name[sbml_id], dm=self.dm)
            for sbml_id in self.species_ids
        }

        # --- 4) Reaction fluxes (COPASI reaction IDs already match SBML IDs) ----
        rxn_df = get_reactions(model=self.dm)
        reaction_fluxes = {rxn_id: float(rxn_df.loc[rxn_id, "flux"]) for rxn_id in self.reaction_ids}

        return {
            "species_concentrations": species_concentrations,
            "fluxes": reaction_fluxes,
            "time": time,
        }


def run_copasi_utc(core: Any) -> None:
    copasi_process = CopasiUTCStep(
        {
            "model_source": "models/BIOMD0000000012_url.xml",  # represillator model
            "time": 10.0,
            "n_points": 5,
        },
        core=core,
    )

    initial_state = copasi_process.initial_state()

    print(f"Initial state: {initial_state}")

    results = copasi_process.update(initial_state)

    print(f"Results: {results}")


def run_copasi_ss(core: Any) -> None:
    copasi_process = CopasiSteadyStateStep(
        {
            "model_source": "models/BIOMD0000000012_url.xml",  # represillator model
        },
        core=core,
    )

    initial_state = copasi_process.initial_state()

    print(f"Initial state: {initial_state}")

    results = copasi_process.update(initial_state)

    print(f"Results: {results}")


if __name__ == "__main__":
    core = ProcessTypes()
    core.register_process("copasi_utc", CopasiUTCStep)
    core.register_process("copasi_ss", CopasiSteadyStateStep)
    core.register_process("copasi_process", CopasiUTCProcess)

    run_copasi_utc(core=core)
    run_copasi_ss(core=core)
