from __future__ import annotations

from functools import reduce
from itertools import product
from typing import TYPE_CHECKING

import numpy as np
from cachetools import cached
from cachetools.keys import hashkey
from scipy.linalg import expm

if TYPE_CHECKING:
    # Imports related to type checking and annotations
    from collections.abc import Iterable, Sequence

    from chemex.nmr.constants import Distribution
    from chemex.nmr.liouvillian import LiouvillianIS
    from chemex.typing import ArrayFloat

    DictArrayFloat = dict[str, ArrayFloat]

# A small value used for numerical stability
SMALL_VALUE = 1e-6


def calculate_propagators(
    liouv: ArrayFloat,
    delays: float | Iterable[float],
    dephasing: bool = False,
) -> ArrayFloat:
    """Calculate the propagators for the given delays.

    Parameters
    ----------
    liouv : ArrayFloat
        The Liouvillians of the system
    delays : float | Iterable[float]
        The delays to calculate the propagators for
    dephasing : bool, optional
        Whether to apply dephasing, by default False

    Returns:
    -------
    ArrayFloat
        The propagators
    """
    # Ensure delays is a 1D NumPy array
    delays = np.asarray(delays).flatten()

    if delays.size == 1 and not dephasing:
        # If there's only one delay and no dephasing, calculate the propagator directly
        # using scipy expm function
        return expm(liouv * delays[0])

    # Calculate eigenvalues and eigenvectors of the Liouvillian matrices
    eigenvalues, eigenvectors = np.linalg.eig(liouv)

    if dephasing:
        # If dephasing is requested, adjust eigenvalues
        eigenvalues = np.where(
            abs(eigenvalues.imag) < SMALL_VALUE,
            eigenvalues,
            eigenvalues * 1e9,
        )

    # Calculate the exponentiated eigenvalues for each delay
    evt = np.exp(np.multiply.outer(delays, eigenvalues))

    # Create an empty array for diagonal matrices
    diag = np.empty((*delays.shape, *liouv.shape), dtype=np.complex128)

    # Calculate diagonal matrices for each delay
    for index in product(*(range(x) for x in evt.shape[:-1])):
        diag[index] = np.diag(evt[index])

    # Calculate the propagators using matrix multiplication
    propagators = eigenvectors @ diag @ np.linalg.inv(eigenvectors)

    if propagators.shape[0] == 1:
        # If there's only one propagator, return it with no extra time dimension
        propagators = propagators[0]

    # Return the real part of the propagators
    return propagators.real


def _get_key(liouvillian: LiouvillianIS, *args, **kwargs):
    return hashkey(liouvillian.basis, *args, **kwargs)


@cached(cache={}, key=_get_key)
def _make_perfect180(liouvillian: LiouvillianIS, spin: str) -> ArrayFloat:
    size = liouvillian.size
    identity = np.eye(size).reshape((1, 1, size, size))
    compx, compy, compz = (f"{spin}{axis}" for axis in "xyz")
    perfect180 = {comp: identity.copy() for comp in (compx, compy)}
    for comp in liouvillian.basis.components:
        vect = liouvillian.vectors[comp].ravel()
        if compx in comp or compz in comp:
            perfect180[compy] -= 2 * np.diag(vect)
        if compy in comp or compz in comp:
            perfect180[compx] -= 2 * np.diag(vect)
    p180 = [perfect180[comp] for comp in (compx, compy, compx, compy)]
    return np.array(p180)


@cached(cache={}, key=_get_key)
def _make_perfect90(liouvillian: LiouvillianIS, spin: str) -> ArrayFloat:
    size = liouvillian.size
    zeros = np.zeros((size, size))
    rot = liouvillian.matrices.get(f"b1x_{spin}", zeros)
    return np.asarray(expm(0.25 * rot)).reshape((1, 1, size, size))


@cached(cache={}, key=_get_key)
def _get_phases(liouvillian: LiouvillianIS) -> DictArrayFloat:
    phases = {}
    size = liouvillian.size
    zeros = np.zeros((size, size))
    for spin in "is":
        l_rotz = liouvillian.matrices.get(f"rotz_{spin}", zeros)
        phases[spin] = np.array([expm(n * 0.5 * np.pi * l_rotz) for n in range(4)])
    return phases


class Spectrometer:
    def _add_phases(self, propagator: ArrayFloat, spin: str = "i") -> ArrayFloat:
        phases = self._phases[spin]
        return np.array([phases[i] @ propagator @ phases[-i] for i in range(4)])

    def __init__(self, liouvillian: LiouvillianIS) -> None:
        self.liouvillian = liouvillian
        size = self.liouvillian.size
        self.identity = np.eye(self.liouvillian.size).reshape((1, 1, size, size))
        self._phases = _get_phases(liouvillian)
        self.perfect90_i = self._add_phases(_make_perfect90(liouvillian, "i"))
        self.perfect180_i = _make_perfect180(liouvillian, "i")
        self.perfect180_s = _make_perfect180(liouvillian, "s")
        self.zfilter = self.keep(
            self.identity,
            components={"ie", "se", "iz", "sz", "2izsz"}
            & set(liouvillian.basis.components),
        )
        self.calculate_i_flag = True
        self.calculate_s_flag = True
        self._pw90_i = 0.0
        self._pw90_s = 0.0
        self._p90_i = np.array(0.0)
        self._p180_i = np.array(0.0)
        self._p240_i = np.array(0.0)
        self._p180_s = np.array(0.0)

    def keep(self, magnetization: ArrayFloat, components: Iterable[str]) -> ArrayFloat:
        return self.liouvillian.keep(magnetization, components)

    def update(self, par_values: dict[str, float]) -> None:
        self.liouvillian.update(par_values)
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def par_values(self) -> dict[str, float]:
        return self.liouvillian.par_values

    @property
    def carrier_i(self) -> float:
        return self.liouvillian.carrier_i

    @carrier_i.setter
    def carrier_i(self, value: float):
        self.liouvillian.carrier_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def carrier_s(self) -> float:
        return self.liouvillian.carrier_s

    @carrier_s.setter
    def carrier_s(self, value: float):
        self.liouvillian.carrier_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def offset_i(self) -> float:
        return self.liouvillian.offset_i

    @offset_i.setter
    def offset_i(self, value: float):
        self.liouvillian.offset_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def offset_s(self) -> float:
        return self.liouvillian.offset_s

    @offset_s.setter
    def offset_s(self, value: float):
        self.liouvillian.offset_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i(self) -> float:
        return self.liouvillian.b1_i

    @b1_i.setter
    def b1_i(self, value: float):
        self._pw90_i = 1.0 / (4.0 * value) if value else 0.0
        self.liouvillian.b1_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_s(self) -> float:
        return self.liouvillian.b1_s

    @b1_s.setter
    def b1_s(self, value: float):
        self._pw90_s = 1.0 / (4.0 * value) if value else 0.0
        self.liouvillian.b1_s = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i_inh_scale(self) -> float:
        return self.liouvillian.b1_i_inh_scale

    @b1_i_inh_scale.setter
    def b1_i_inh_scale(self, value: float):
        self.liouvillian.b1_i_inh_scale = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def b1_i_inh_res(self) -> int:
        return self.liouvillian.b1_i_inh_res

    @b1_i_inh_res.setter
    def b1_i_inh_res(self, value: int):
        self.liouvillian.b1_i_inh_res = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    @property
    def jeff_i(self) -> Distribution:
        return self.liouvillian.jeff_i

    @jeff_i.setter
    def jeff_i(self, value: Distribution):
        self.liouvillian.jeff_i = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    def get_equilibrium(self) -> ArrayFloat:
        return self.liouvillian.get_equilibrium()

    def get_start_magnetization(
        self,
        terms: Iterable[str],
        atom: str = "h",
    ) -> ArrayFloat:
        return self.liouvillian.get_start_magnetization(terms=terms, atom=atom)

    @property
    def detection(self) -> str:
        return self.liouvillian.detection

    @detection.setter
    def detection(self, value: str):
        self.liouvillian.detection = value

    def detect(self, magnetization: ArrayFloat) -> float:
        return self.liouvillian.detect(magnetization)

    @property
    def gradient_dephasing(self) -> float:
        return self.liouvillian.gradient_dephasing

    @gradient_dephasing.setter
    def gradient_dephasing(self, value: float):
        self.liouvillian.gradient_dephasing = value
        self.calculate_i_flag = True
        self.calculate_s_flag = True

    def delays(self, times: float | Iterable[float]) -> ArrayFloat:
        return calculate_propagators(self.liouvillian.l_free, times)

    def pulse_i(
        self,
        times: float | Iterable[float],
        phase: float,
        scale: float = 1.0,
    ) -> ArrayFloat:
        dephased = self.b1_i_inh_scale == np.inf
        rad = phase * np.pi * 0.5
        liouv = (
            self.liouvillian.l_free
            + scale * np.cos(rad) * self.liouvillian.l_b1x_i
            + scale * np.sin(rad) * self.liouvillian.l_b1y_i
        )
        return calculate_propagators(liouv, times, dephased)

    def pulse_s(
        self,
        times: float | Iterable[float],
        phase: float,
        scale: float = 1.0,
    ) -> ArrayFloat:
        rad = phase * np.pi * 0.5
        liouv = (
            self.liouvillian.l_free
            + scale * np.cos(rad) * self.liouvillian.l_b1x_s
            + scale * np.sin(rad) * self.liouvillian.l_b1y_s
        )
        return calculate_propagators(liouv, times)

    def pulse_is(
        self,
        times: float | Iterable[float],
        phase_i: float,
        phase_s: float,
    ) -> ArrayFloat:
        dephased = self.b1_i_inh_scale == np.inf
        liouv = (
            self.liouvillian.l_free
            + np.cos(phase_i * np.pi * 0.5) * self.liouvillian.l_b1x_i
            + np.sin(phase_i * np.pi * 0.5) * self.liouvillian.l_b1y_i
            + np.cos(phase_s * np.pi * 0.5) * self.liouvillian.l_b1x_s
            + np.sin(phase_s * np.pi * 0.5) * self.liouvillian.l_b1y_s
        )
        return calculate_propagators(liouv, times, dephased)

    def shaped_pulse_i(
        self,
        pw: float,
        amplitudes: Sequence[float],
        phases: Iterable[float],
    ) -> ArrayFloat:
        time = pw / len(amplitudes)
        pairs = list(zip(amplitudes, phases, strict=True))
        pulses = {
            (amp, ph): self.pulse_i(time, ph, scale=amp) for amp, ph in set(pairs)
        }
        base = reduce(np.matmul, (pulses[pair] for pair in reversed(pairs)))
        return self._add_phases(base, "i")

    def _calculate_base_pulses_i(self) -> None:
        if self.calculate_i_flag:
            pws = np.array([1.0, 2.0, 8.0 / 3.0]) * self._pw90_i
            base = self.pulse_i(pws, 0.0)
            pulses = self._add_phases(base, "i")
            self._p90_i, self._p180_i, self._p240_i = pulses.swapaxes(0, 1)

    @property
    def p90_i(self) -> ArrayFloat:
        self._calculate_base_pulses_i()
        return self._p90_i

    @property
    def p180_i(self) -> ArrayFloat:
        self._calculate_base_pulses_i()
        return self._p180_i

    @property
    def p240_i(self) -> ArrayFloat:
        self._calculate_base_pulses_i()
        return self._p240_i

    @property
    def p9018090_i_1(self) -> ArrayFloat:
        return self.p90_i[[3, 0, 1, 2]] @ self.p180_i @ self.p90_i[[3, 0, 1, 2]]

    @property
    def p9018090_i_2(self) -> ArrayFloat:
        return self.p90_i[[1, 2, 3, 0]] @ self.p180_i @ self.p90_i[[1, 2, 3, 0]]

    @property
    def p9024090_i_1(self) -> ArrayFloat:
        return self.p90_i[[3, 0, 1, 2]] @ self.p240_i @ self.p90_i[[3, 0, 1, 2]]

    @property
    def p9024090_i_2(self) -> ArrayFloat:
        return self.p90_i[[1, 2, 3, 0]] @ self.p240_i @ self.p90_i[[1, 2, 3, 0]]

    def _calculate_base_pulses_s(self) -> None:
        if self.calculate_s_flag:
            pws = np.array([2.0]) * self._pw90_s
            base = self.pulse_s(pws, 0.0)
            pulses = self._add_phases(base, "s")
            self._p180_s = pulses.swapaxes(0, 1)

    @property
    def p180_s(self) -> ArrayFloat:
        self._calculate_base_pulses_s()
        return self._p180_s

    def p9024090_nh(self, reverse: bool = False) -> ArrayFloat:
        ph_n = 1 if reverse else 3
        ph_h = 3 if reverse else 1
        pw240i, pw9024090i = np.array([8.0, 14.0]) * self._pw90_i / 3.0
        pw240s, pw9024090s = np.array([8.0, 14.0]) * self._pw90_s / 3.0
        t0, t1, t2, t3 = 0.5 * np.diff(
            np.sort([pw240i, pw240s, pw9024090i, pw9024090s]),
            prepend=0.0,
        )
        p0 = self.pulse_is(2.0 * t0, 0, 0)
        if pw9024090i <= pw9024090s:
            p1 = self.pulse_is(t1, ph_n, 0)
            p2 = (
                self.pulse_is(t2, ph_n, ph_h)
                if pw9024090i > pw240s
                else self.pulse_s(t2, 0)
            )
            p3 = self.pulse_s(t3, ph_h)
        else:
            p1 = self.pulse_is(t1, 0, ph_h)
            p2 = (
                self.pulse_is(t2, ph_n, ph_h)
                if pw9024090s > pw240i
                else self.pulse_i(t2, 0)
            )
            p3 = self.pulse_i(t3, ph_n)
        pw9024090is_xx = p3 @ p2 @ p1 @ p0 @ p1 @ p2 @ p3
        return self._add_phases(self._add_phases(pw9024090is_xx, "s"), "i")

    @property
    def p9024090_nh_1(self) -> ArrayFloat:
        return self.p9024090_nh()

    @property
    def p9024090_nh_2(self) -> ArrayFloat:
        return self.p9024090_nh(reverse=True)

    def calculate_shifts(self) -> ArrayFloat:
        liouv = self.liouvillian.l_free.reshape(
            (self.liouvillian.size, self.liouvillian.size),
        )
        return np.linalg.eigvals(liouv).imag

    def offsets_to_ppms(self, offsets: ArrayFloat) -> ArrayFloat:
        return self.liouvillian.offsets_to_ppms(offsets)

    def ppms_to_offsets(self, ppms: ArrayFloat | float) -> ArrayFloat | float:
        return self.liouvillian.ppms_to_offsets(ppms)
