from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass, field
from functools import cached_property
from operator import attrgetter
from typing import TYPE_CHECKING, Protocol, TypeVar

from cachetools import LRUCache, cachedmethod

if TYPE_CHECKING:
    from collections.abc import Hashable

    from lmfit import Parameters as ParametersLF

    from chemex.containers.data import Data
    from chemex.nmr.spectrometer import Spectrometer
    from chemex.parameters.spin_system import SpinSystem
    from chemex.printers.data import Printer
    from chemex.typing import ArrayBool, ArrayFloat

Self = TypeVar("Self", bound="Profile")


class PulseSequence(Protocol):
    """Defines a protocol for pulse sequence classes."""

    def calculate(self, spectrometer: Spectrometer, data: Data) -> ArrayFloat:
        """Calculate the ArrayFloat given a Spectrometer and Data."""
        ...

    def is_reference(self, metadata: ArrayFloat) -> ArrayBool:
        """Check if the metadata is a reference."""
        ...


class Filterer(Protocol):
    """Defines a protocol for filterer classes."""

    def filter(self, data: Data) -> None:
        """Filter the given Data."""
        ...


@dataclass(order=True)
class Profile:
    """Represents the profile for a specific experiment."""

    data: Data = field(compare=False)
    spectrometer: Spectrometer = field(compare=False)
    pulse_sequence: PulseSequence = field(compare=False)
    name_map: dict[str, str] = field(compare=False)
    printer: Printer = field(compare=False)
    filterer: Filterer | None = field(compare=False, default=None)
    is_scaled: bool = field(compare=False, default=True)
    spin_system: SpinSystem = field(compare=True, init=False)
    cache: LRUCache = field(compare=False, init=False)

    def _cache_key(self, params: ParametersLF) -> tuple[Hashable, ...]:
        return (
            *(params[param_id].value for param_id in self.param_ids),
            self.data.metadata.tostring(),
        )

    def __post_init__(self):
        """Initialize derived attributes."""
        self.spin_system = self.spectrometer.liouvillian.spin_system
        self.data.refs = self.pulse_sequence.is_reference(self.data.metadata)
        self.cache = LRUCache(maxsize=5)

    @cached_property
    def param_ids(self) -> set[str]:
        """Get the set of parameter IDs."""
        return set(self.name_map.values())

    def _get_parameter_values(self, params: ParametersLF) -> dict[str, float]:
        """Get the parameter values from the provided parameters."""
        return {
            local_name: params[param_id].value
            for local_name, param_id in self.name_map.items()
        }

    def update_spectrometer(self, params: ParametersLF) -> None:
        """Update the spectrometer settings."""
        parameter_values = self._get_parameter_values(params)
        self.spectrometer.update(parameter_values)

    def calculate(self, params: ParametersLF) -> ArrayFloat:
        """Calculate and return the ArrayFloat."""
        self.update_spectrometer(params)
        self.data.calc = self.pulse_sequence.calculate(self.spectrometer, self.data)
        if self.is_scaled:
            self.data.scale_calc()
        return self.data.calc

    @cachedmethod(attrgetter("cache"), key=_cache_key)
    def residuals(self, params: ParametersLF) -> ArrayFloat:
        """Calculate and return residuals."""
        residuals = (self.calculate(params) - self.data.exp) / self.data.err
        return residuals[self.data.mask]

    def filter(self, params: ParametersLF) -> None:
        """Apply a filter to the data, if a filterer is available."""
        if self.filterer is not None:
            self.update_spectrometer(params)
            self.filterer.filter(self.data)

    def set_noise(self, value: float):
        """Set the noise value."""
        self.data.err[:] = value

    def prepare_for_simulation(self) -> None:
        """Prepare data for simulation."""
        self.data.prepare_for_simulation()

    def monte_carlo(self: Self) -> Self:
        """Generate a Monte Carlo variant of the profile."""
        profile = deepcopy(self)
        profile.data = profile.data.monte_carlo()
        return profile

    def bootstrap(self: Self) -> Self:
        """Generate a bootstrap variant of the profile."""
        profile = deepcopy(self)
        profile.data = profile.data.bootstrap()
        return profile

    def any_duplicate(self):
        """Check for duplicate data points."""
        return self.data.any_duplicate()

    def __add__(self: Self, other: Self) -> Self:
        """Combine two profiles."""
        profile = deepcopy(self)
        profile.data = self.data + other.data
        return profile

    def __str__(self) -> str:
        """String representation of the Profile."""
        return self.printer.print(str(self.spin_system), self.data)
