"""High-level Python API for QuEst (Quenching Estimator).

This module provides a small, GUI-independent interface around the core
`DonorDecay` simulation so that you can use QuEst from scripts and
Jupyter notebooks.

Typical usage
-------------

>>> from pathlib import Path
>>> from quest.api import load_project, simulate_project
>>> project = load_project(Path("my_project.quest.json"))
>>> result = simulate_project(project)
>>> result.time.shape == result.donor_counts.shape
True

The API is intentionally minimal and focused on a single labeled donor
(with optional FRET acceptor). The underlying project schema is the same
as the JSON written by the QuEst GUI.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, MutableMapping, Optional, Sequence, Tuple

import json

import numpy as np

from .lib.tools.dye_diffusion.dye_diffusion import DonorDecay


@dataclass
class DecaySimulationResult:
    """Container for a single QuEst decay simulation.

    Attributes
    ----------
    time:
        1D array of bin centers in nanoseconds.
    donor_counts:
        1D array of donor-only decay counts per bin.
    fret_counts:
        1D array of donor+FRET counts per bin, or ``None`` if FRET was
        not enabled in the project.
    quantum_yield_donor:
        Simulated quantum yield of the donor-only decay.
    quantum_yield_fret:
        Quantum yield in the presence of FRET, or ``None`` if FRET was
        not simulated.
    collisions_fraction:
        Fraction of time steps in which the donor was in contact with a
        quencher (0–1).
    fret_efficiency:
        Effective FRET efficiency (0–1) estimated from the simulation,
        or ``None`` if FRET was not simulated.
    project:
        The project dictionary used to run the simulation. This is
        returned unchanged to make it easy to inspect the settings.
    """

    time: np.ndarray
    donor_counts: np.ndarray
    fret_counts: Optional[np.ndarray]
    quantum_yield_donor: float
    quantum_yield_fret: Optional[float]
    collisions_fraction: float
    fret_efficiency: Optional[float]
    project: Mapping[str, Any]


def load_project(path: Path | str) -> MutableMapping[str, Any]:
    """Load a QuEst project JSON file.

    The returned mapping follows the schema produced by
    :meth:`quest.lib.tools.dye_diffusion.TransientDecayGenerator.to_project_dict`,
    including the ``"fret"`` block with ``"dyes"`` and ``"R0_matrix"``.

    Parameters
    ----------
    path:
        Path to a ``.quest.json`` (or plain ``.json``) project file.

    Returns
    -------
    dict
        A mutable dictionary with all project settings.
    """

    p = Path(path)
    with p.open("r", encoding="utf8") as fp:
        return json.load(fp)


def _build_donor_from_project(
    project: Mapping[str, Any],
    pdb_path: Optional[Path | str] = None,
) -> DonorDecay:
    """Create and configure a :class:`DonorDecay` from a project dict.

    This helper does not run any simulation. It only instantiates the
    core model and wires parameters from the project JSON into
    :class:`DonorDecay` attributes.
    """

    p = project

    attachment = p.get("attachment") or {}

    fret_block = p.get("fret") or {}
    dyes = fret_block.get("dyes") or []
    if not dyes:
        raise KeyError("Project must define at least one dye with a diffusion coefficient 'D'.")
    donor_dye = dyes[0]
    donor_D = donor_dye.get("D")
    if donor_D is None:
        raise KeyError("Donor dye entry must define diffusion coefficient 'D'.")

    d = DonorDecay(
        tau0=p["tau0"],
        kQ=p["kQ"],
        nph=p["n_photons"],
        verbose=False,
        auto_update=False,
        pdb=None,
        attachment_residue=attachment.get("residue"),
        attachment_atom=attachment.get("atom"),
        attachment_chain=attachment.get("chain"),
        dg=p["dg"],
        sticky_mode=p.get("sticky_mode", "quencher"),
        save_avs=p.get("save_avs", False),
        D=donor_D,
        slow_fact=p["slow_fact"],
        critical_distance=p["critical_distance"],
        output_file=p.get("output_file", "out"),
        av_parameter=p["av_parameter"],
        quencher=p["quencher"],
        t_max=p["t_max"],
        t_step=p["t_step"],
        slow_radius=p["slow_radius"],
        all_quencher_atoms=p.get("all_quencher_atoms", True),
    )

    structure_path = pdb_path or p.get("pdb")
    if structure_path is None:
        raise ValueError("Project must define a 'pdb' path or pdb_path must be given.")
    d.structure = str(structure_path)

    fret = p.get("fret") or {}
    if fret.get("enabled"):
        d.fret_enabled = True
        dyes: Sequence[Mapping[str, Any]] = fret.get("dyes") or []
        if len(dyes) > 1:
            acceptor = dyes[1]
            acc_att = acceptor.get("attachment") or {}
            d.acceptor_chain = acc_att.get("chain")
            d.acceptor_residue = acc_att.get("residue")
            d.acceptor_atom = acc_att.get("atom")
            av_param_acc = acceptor.get("av_parameter")
            if av_param_acc is not None:
                d.fret_av_parameter = av_param_acc

        R0_matrix: Sequence[Sequence[Optional[float]]] = fret.get("R0_matrix") or []
        if R0_matrix and len(R0_matrix[0]) > 1:
            r01 = R0_matrix[0][1]
            if r01 is not None:
                d.fret_R0 = float(r01)

    return d


def simulate_project(
    project: Mapping[str, Any],
    *,
    pdb_path: Optional[Path | str] = None,
    nbins: Optional[int] = None,
    tac_range: Tuple[float, float] = (0.0, 50.0),
) -> DecaySimulationResult:
    """Run a full QuEst simulation from a project dictionary.

    This function performs the following steps:

    1. Construct a :class:`DonorDecay` instance from the project.
    2. Run :meth:`DonorDecay.update_all` to compute the donor-only
       trajectory and decay.
    3. If FRET is enabled in the project, run the additional donor+FRET
       decay.
    4. Convert the photon arrival times into histogrammed decay curves.

    Parameters
    ----------
    project:
        Project settings dictionary, typically created by
        :func:`load_project` or by the QuEst GUI.
    pdb_path:
        Optional override for the structure file. If given, this path is
        used instead of ``project["pdb"]``.
    nbins:
        Number of histogram bins for the decay. If omitted, the value of
        ``project["n_bins"]`` is used, or 4096 as a fallback.
    tac_range:
        Time range (start, stop) in nanoseconds for the decay histogram.

    Returns
    -------
    DecaySimulationResult
        Object containing the time axis, donor-only decay, optional
        donor+FRET decay, quantum yields, collision fraction and FRET
        efficiency estimate.
    """

    if nbins is None:
        nbins = int(project.get("n_bins", 4096))

    model = _build_donor_from_project(project, pdb_path=pdb_path)

    model.update_all(verbose=False)

    edges, donor_counts = model.get_histogram(nbins=nbins, range=tac_range)
    time = 0.5 * (edges[:-1] + edges[1:])

    fret_counts: Optional[np.ndarray] = None
    qy_fret: Optional[float] = None
    fret_eff: Optional[float] = None

    if getattr(model, "fret_enabled", False):
        xf, yf = model.get_histogram_fret(nbins=nbins, range=tac_range)
        if xf is not None and yf is not None:
            fret_counts = yf
            qy_fret = model.quantum_yield_fret
            fret_eff = model.fret_efficiency

    return DecaySimulationResult(
        time=time,
        donor_counts=donor_counts,
        fret_counts=fret_counts,
        quantum_yield_donor=model.quantum_yield,
        quantum_yield_fret=qy_fret,
        collisions_fraction=model.collisions,
        fret_efficiency=fret_eff,
        project=project,
    )
