from __future__ import annotations

import numpy as np

from xq_pulse.pulse.dsl import delay, drive, for_idx, parameter_sweep, pulse_program, sequence
from xq_pulse.pulse.expression import Expression
from xq_pulse.pulse.program import PulseProgram
from xq_pulse.util import Quantity, unit

# AXY8 composite pattern from Qudi: XYXYYXYX
_AXY8_PATTERN: tuple[str, ...] = ("X", "Y", "X", "Y", "Y", "X", "Y", "X")
# KDD5-even composite phases (degrees), Y is X shifted by +90 deg.
_X_PHASES_DEG: tuple[float, ...] = (30.0, 0.0, 90.0, 0.0, 30.0)
_Y_PHASES_DEG: tuple[float, ...] = tuple(phase + 90.0 for phase in _X_PHASES_DEG)

# Cache solved spacing vectors keyed by Fourier coefficients.
_SPACING_CACHE: dict[tuple[float, float, float, float], np.ndarray] = {}
_THETA_CACHE: dict[tuple[float, float, float, float], np.ndarray] = {}


def x90_pulse(frequency: Quantity, amplitude: Quantity, rabi_period: Quantity) -> None:
    drive(duration=rabi_period / 4, frequency=frequency, amplitude=amplitude)


def _kdd5even_residual(theta: np.ndarray, f1: float, f2: float, f3: float, f4: float) -> np.ndarray:
    theta1, theta2, theta3, theta4 = theta
    theta5 = theta2 + theta4 - (theta1 + theta3) + np.pi / 2

    residual = np.zeros(4, dtype=float)
    residual[0] = f1 - 4 / (1 * np.pi) * (
        np.sin(1 * theta1) + np.sin(1 * theta3) + np.sin(1 * theta5) - np.sin(1 * theta2) - np.sin(1 * theta4)
    )
    residual[1] = f2 - 4 / (2 * np.pi) * (
        np.sin(2 * theta1) + np.sin(2 * theta3) + np.sin(2 * theta5) - np.sin(2 * theta2) - np.sin(2 * theta4)
    )
    residual[2] = f3 - 4 / (3 * np.pi) * (
        np.sin(3 * theta1) + np.sin(3 * theta3) + np.sin(3 * theta5) - np.sin(3 * theta2) - np.sin(3 * theta4)
    )
    residual[3] = f4 - 4 / (4 * np.pi) * (
        np.sin(4 * theta1) + np.sin(4 * theta3) + np.sin(4 * theta5) - np.sin(4 * theta2) - np.sin(4 * theta4)
    )
    return residual


def _solve_axy_spacing(f1: float, f2: float, f3: float, f4: float) -> np.ndarray:
    key = (float(f1), float(f2), float(f3), float(f4))
    cached = _SPACING_CACHE.get(key)
    if cached is not None:
        return cached

    def _newton(theta0: np.ndarray) -> np.ndarray:
        theta = theta0.copy()
        for _ in range(40):
            residual = _kdd5even_residual(theta, *key)
            if np.linalg.norm(residual, ord=2) < 1e-11:
                break

            eps = 1e-6
            jacobian = np.zeros((4, 4), dtype=float)
            for col in range(4):
                theta_plus = theta.copy()
                theta_minus = theta.copy()
                theta_plus[col] += eps
                theta_minus[col] -= eps
                jacobian[:, col] = (
                    _kdd5even_residual(theta_plus, *key) - _kdd5even_residual(theta_minus, *key)
                ) / (2.0 * eps)

            try:
                step = np.linalg.solve(jacobian, -residual)
            except np.linalg.LinAlgError:
                return theta

            # Damped Newton step for robustness.
            damping = 1.0
            current_norm = np.linalg.norm(residual, ord=2)
            updated = False
            for _ in range(10):
                candidate = theta + damping * step
                candidate_norm = np.linalg.norm(_kdd5even_residual(candidate, *key), ord=2)
                if candidate_norm < current_norm:
                    theta = candidate
                    updated = True
                    break
                damping *= 0.5
            if not updated:
                theta = theta + 0.1 * step
        return theta

    def _candidate_to_spacings(theta: np.ndarray) -> tuple[np.ndarray, float]:
        theta5 = theta[1] + theta[3] - (theta[0] + theta[2]) + np.pi / 2
        solved = np.append(theta, theta5) / (2 * np.pi)
        # Match Qudi helper: mirror from first three spacings.
        s0, s1, s2 = solved[:3]
        spacings = np.array([s0, s1, s2, 2.0 * s2 - s1, 2.0 * s2 - s0], dtype=float)
        residual_norm = np.linalg.norm(_kdd5even_residual(theta, *key), ord=2)
        return spacings, residual_norm

    seeds: list[np.ndarray] = [
        np.array([0.1 * np.pi, 0.3 * np.pi, 0.6 * np.pi, 0.9 * np.pi], dtype=float),
        np.array([0.08 * np.pi, 0.28 * np.pi, 0.58 * np.pi, 0.88 * np.pi], dtype=float),
        np.array([0.12 * np.pi, 0.32 * np.pi, 0.62 * np.pi, 0.92 * np.pi], dtype=float),
        np.array([0.2 * np.pi, 0.4 * np.pi, 0.7 * np.pi, 1.0 * np.pi], dtype=float),
    ]

    # Continuation seed: nearest cached f1 for same (f2,f3,f4).
    same_tail = [k for k in _THETA_CACHE if k[1:] == key[1:]]
    if same_tail:
        nearest_key = min(same_tail, key=lambda k: abs(k[0] - key[0]))
        seeds.insert(0, _THETA_CACHE[nearest_key])

    best_physical: tuple[np.ndarray, np.ndarray, float] | None = None
    best_any: tuple[np.ndarray, np.ndarray, float] | None = None
    for seed in seeds:
        theta = _newton(seed)
        spacings, residual_norm = _candidate_to_spacings(theta)
        candidate = (theta, spacings, residual_norm)
        if best_any is None or residual_norm < best_any[2]:
            best_any = candidate

        is_physical = np.all(np.diff(spacings) > 0) and np.all(spacings > 0.0) and np.all(spacings < 0.5)
        if is_physical and (best_physical is None or residual_norm < best_physical[2]):
            best_physical = candidate

    if best_physical is None:
        assert best_any is not None, "No AXY solver candidate generated"
        _, bad_spacings, bad_norm = best_any
        raise ValueError(
            "No physical AXY spacing solution found for "
            f"f1={key[0]:.5g}, f2={key[1]:.5g}, f3={key[2]:.5g}, f4={key[3]:.5g}. "
            f"Best residual norm={bad_norm:.3e}, spacings={bad_spacings}."
        )

    theta, spacings, residual_norm = best_physical
    if residual_norm > 1e-7:
        raise ValueError(
            "AXY spacing solver did not converge sufficiently for "
            f"f1={key[0]:.5g}, f2={key[1]:.5g}, f3={key[2]:.5g}, f4={key[3]:.5g}; residual={residual_norm:.3e}"
        )

    _THETA_CACHE[key] = theta
    _SPACING_CACHE[key] = spacings
    return spacings


def _append_composite(
    tau: Expression,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    phases_deg: tuple[float, ...],
    spacings: np.ndarray,
) -> None:
    pi_duration = rabi_period / 2
    two = 2 * unit.dimensionless
    one = 1 * unit.dimensionless
    s0 = float(spacings[0]) * unit.dimensionless
    s1 = float(spacings[1]) * unit.dimensionless
    s2 = float(spacings[2]) * unit.dimensionless
    s3 = float(spacings[3]) * unit.dimensionless
    s4 = float(spacings[4]) * unit.dimensionless

    # Composite pulse centers inside one tau window are at t = 2 * tau * spacing_j.
    idle_durations = [
        tau * two * s0 - pi_duration / 2,
        tau * two * (s1 - s0) - pi_duration,
        tau * two * (s2 - s1) - pi_duration,
        tau * two * (s3 - s2) - pi_duration,
        tau * two * (s4 - s3) - pi_duration,
        tau * (one - two * s4) - pi_duration / 2,
    ]

    for idx, idle_duration in enumerate(idle_durations):
        if isinstance(idle_duration, Quantity):
            assert idle_duration >= 0 * unit.s, (
                "tau is too short for AXY composite pulse spacing and rabi_period "
                f"(idle segment {idx} would be negative: {idle_duration})"
            )
        delay(duration=idle_duration)
        if idx < 5:
            drive(
                duration=pi_duration,
                frequency=frequency,
                amplitude=amplitude,
                phase=phases_deg[idx] * unit.deg,
            )


def axy8_block(
    tau: Expression,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    f1: float = 1.0,
    f2: float = 0.0,
    f3: float = 0.0,
    f4: float = 0.0,
) -> None:
    spacings = _solve_axy_spacing(f1=f1, f2=f2, f3=f3, f4=f4)

    with sequence():
        for pulse_axis in _AXY8_PATTERN:
            phases_deg = _X_PHASES_DEG if pulse_axis == "X" else _Y_PHASES_DEG
            _append_composite(tau, rabi_period, frequency, amplitude, phases_deg, spacings)


def axy8(
    tau: Expression,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    order: int,
    f1: float = 1.0,
    f2: float = 0.0,
    f3: float = 0.0,
    f4: float = 0.0,
) -> None:
    assert int(order) >= 1, "order must be at least 1"

    with sequence():
        x90_pulse(frequency, amplitude, rabi_period)
        with for_idx(1*unit.dimensionless, int(order)*unit.dimensionless, 1*unit.dimensionless) as (_, _loop_param):
            axy8_block(
                tau=tau,
                rabi_period=rabi_period,
                frequency=frequency,
                amplitude=amplitude,
                f1=f1,
                f2=f2,
                f3=f3,
                f4=f4,
            )
        x90_pulse(frequency, amplitude, rabi_period)


def axy8_sequence(
    start: Quantity,
    stop: Quantity,
    step: Quantity,
    rabi_period: Quantity,
    frequency: Quantity,
    amplitude: Quantity,
    order: int,
    f1: float = 1.0,
    f2: float = 0.0,
    f3: float = 0.0,
    f4: float = 0.0,
) -> PulseProgram:
    """Return an AXY8 pulse program that sweeps tau."""
    assert start.is_compatible_with(unit.s), "start must have time units"
    assert stop.is_compatible_with(unit.s), "stop must have time units"
    assert step.is_compatible_with(unit.s), "step must have time units"
    assert rabi_period.is_compatible_with(unit.s), "rabi_period must have time units"
    assert frequency.is_compatible_with(unit.Hz), "frequency must have frequency units"
    assert amplitude.is_compatible_with(unit.T), "amplitude must have magnetic field units"
    assert order >= 1, "order must be at least 1"

    with pulse_program() as program:
        with parameter_sweep(start=start, stop=stop, step=step) as (_, tau):
            axy8(
                tau=tau,
                rabi_period=rabi_period,
                frequency=frequency,
                amplitude=amplitude,
                order=int(order),
                f1=f1,
                f2=f2,
                f3=f3,
                f4=f4,
            )

    return program
