from __future__ import annotations

from typing import Literal, Self

import numpy as np
from numpy.linalg import matrix_power
from pydantic import Field, computed_field, model_validator

from chemex.configuration.base import ExperimentConfiguration, ToBeFitted
from chemex.configuration.conditions import ConditionsWithValidations
from chemex.configuration.data import CestDataSettings
from chemex.configuration.experiment import (
    B1InhomogeneityMixin,
    MFCestSettings,
    normalize_b1_eff_alias,
)
from chemex.configuration.types import B1Field, ChemicalShift, Delay
from chemex.containers.data import Data
from chemex.containers.dataset import load_relaxation_dataset
from chemex.experiments.factories import Creators, factories
from chemex.filterers import CestFilterer
from chemex.nmr.basis import Basis
from chemex.nmr.constants import get_multiplet
from chemex.nmr.spectrometer import Spectrometer
from chemex.parameters.spin_system import SpinSystem
from chemex.parameters.spin_system.nucleus import Nucleus
from chemex.plotters.cest import CestPlotter
from chemex.printers.data import CestPrinter
from chemex.typing import Array

Dataset = list[tuple[SpinSystem, Data]]

EXPERIMENT_NAME = "dcest_15n"

OFFSET_REF = 1e4


class DCest15NSettings(MFCestSettings, B1InhomogeneityMixin):
    """D-CEST 15N experiment settings with DANTE pulse train."""

    name: Literal["dcest_15n"]
    time_t1: Delay = Field(description="CEST relaxation delay (seconds)")
    time_equil: Delay = Field(
        default=0.0,
        description="Equilibration delay at end of CEST period (seconds)",
    )
    carrier: ChemicalShift = Field(description="15N carrier position during CEST (ppm)")
    # Effective B1 for the equivalent continuous-wave CEST irradiation.
    # Accept both 'b1_eff' (preferred) and 'b1_frq' (alias) in configs.
    b1_eff: B1Field = Field(
        ...,
        description="Effective B1 field for DANTE excitation (Hz)",
    )

    _normalize_b1_eff_alias = model_validator(mode="before")(normalize_b1_eff_alias)

    @model_validator(mode="after")
    def _validate_dcest_fields(self) -> Self:
        """Validate D-CEST requires pw90 for hardware B1."""
        if self.pw90 is None:
            msg = (
                "D-CEST experiments require 'pw90' (hardware B1 pulse width). "
                "Both 'pw90' and 'b1_eff' (alias: 'b1_frq') must be specified."
            )
            raise ValueError(msg)
        return self

    # D-CEST convention: the B1 inhomogeneity distribution should be centered
    # on the hardware B1 set by pw90, regardless of whether b1_frq is also
    # provided. We therefore override the nominal B1 used to build the
    # distribution to always derive it from pw90.
    def get_b1_nominal(self) -> float:
        """Get nominal B1 from hardware pw90 for distribution centering."""
        if self.pw90 is None:
            msg = (
                "For D-CEST, 'pw90' must be specified to define the nominal B1 "
                "used for the B1 distribution."
            )
            raise ValueError(msg)
        return 1.0 / (4.0 * float(self.pw90))

    @computed_field
    @property
    def pw_dante(self) -> float:
        """Duration of each DANTE pulse (seconds)."""
        if self.pw90 is None:
            msg = "pw90 must be specified for DANTE pulse calculation"
            raise ValueError(msg)
        return 4.0 * float(self.pw90) * float(self.b1_eff) / float(self.sw)

    @computed_field
    @property
    def tau_dante(self) -> float:
        """Inter-pulse delay in DANTE train (seconds)."""
        return 1.0 / self.sw - self.pw_dante

    @computed_field
    @property
    def ncyc_dante(self) -> int:
        """Number of DANTE pulses in the train."""
        return int(self.time_t1 * self.sw + 0.1)

    @computed_field
    @property
    def start_terms(self) -> list[str]:
        """Starting magnetization terms based on kinetic model."""
        starts = {"2st_hd": ["iz_a"], "4st_hd": ["iz_a", "iz_b"]}
        return starts.get(self.model_name, [f"iz{self.suffix_start}"])

    @computed_field
    @property
    def detection(self) -> str:
        """Detection operator for the experiment."""
        return f"[iz{self.suffix_detect}]"


class DCest15NConfig(
    ExperimentConfiguration[
        DCest15NSettings, ConditionsWithValidations, CestDataSettings
    ],
):
    @property
    def to_be_fitted(self) -> ToBeFitted:
        state = self.experiment.observed_state
        return ToBeFitted(
            rates=["r2_i", f"r1_i_{state}"],
            model_free=[f"tauc_{state}", f"s2_{state}"],
        )


def build_spectrometer(config: DCest15NConfig, spin_system: SpinSystem) -> Spectrometer:
    settings = config.experiment
    conditions = config.conditions

    basis = Basis(type="ixyz", spin_system="nh", model=config.model)
    spectrometer = Spectrometer.from_spin_system(spin_system, basis, conditions)

    spectrometer.carrier_i = settings.carrier

    spectrometer.set_b1_i_inhomogeneity(
        settings.get_b1_nominal(),
        settings.b1_distribution,
    )

    spectrometer.detection = settings.detection

    if "13c" in conditions.label:
        symbol = spin_system.symbols["i"]
        atom = spin_system.atoms["i"]
        spectrometer.jeff_i = get_multiplet(symbol, atom.name)

    return spectrometer


class DCest15NSequence:
    """Sequence for D-CEST 15N experiment."""

    def __init__(self, settings: DCest15NSettings) -> None:
        self.settings = settings

    @staticmethod
    def is_reference(metadata: Array) -> Array:
        return np.abs(metadata) > OFFSET_REF

    def calculate(self, spectrometer: Spectrometer, data: Data) -> Array:
        offsets = data.metadata

        start = spectrometer.get_start_magnetization(
            self.settings.start_terms, Nucleus.N15
        )

        d_eq = (
            spectrometer.delays(self.settings.time_equil)
            if self.settings.time_equil > 0
            else spectrometer.identity
        )

        intensities: dict[float, Array] = {}

        for offset in set(offsets):
            if self.is_reference(offset):
                intensities[offset] = d_eq @ start
                continue

            spectrometer.offset_i = offset

            p_delay = spectrometer.delays(self.settings.tau_dante)
            p_pulse = spectrometer.pulse_i(self.settings.pw_dante, 0.0)

            intensities[offset] = (
                d_eq @ matrix_power(p_delay @ p_pulse, self.settings.ncyc_dante) @ start
            )

        return np.array(
            [spectrometer.detect(intensities[offset]) for offset in offsets],
        )


def register() -> None:
    creators = Creators(
        config_creator=DCest15NConfig,
        spectrometer_creator=build_spectrometer,
        sequence_creator=DCest15NSequence,
        dataset_creator=load_relaxation_dataset,
        filterer_creator=CestFilterer,
        printer_creator=CestPrinter,
        plotter_creator=CestPlotter,
    )
    factories.register(name=EXPERIMENT_NAME, creators=creators)
