from __future__ import annotations

import math
from abc import ABC, abstractmethod
from functools import lru_cache

import numpy as np
import pint
import sympy
from attrs import field, frozen, validators
from plotly import graph_objects as go
from plotly.subplots import make_subplots

from xq_pulse.util import unit

DimensionlessLiteral = float | int | pint.Quantity

_TAU_SYMBOL = sympy.Symbol("tau", real=True)
_SYMPY_PARSE_LOCALS = {
    "tau": _TAU_SYMBOL,
    "exp": sympy.exp,
    "sin": sympy.sin,
    "cos": sympy.cos,
    "tan": sympy.tan,
    "sqrt": sympy.sqrt,
    "log": sympy.log,
    "Abs": sympy.Abs,
    "Piecewise": sympy.Piecewise,
    "Max": sympy.Max,
    "Min": sympy.Min,
    "pi": sympy.pi,
    "E": sympy.E,
}
_SYMPY_ALLOWED_FUNCTION_NAMES = frozenset(
    {
        "exp",
        "sin",
        "cos",
        "tan",
        "sqrt",
        "log",
        "Abs",
        "Piecewise",
        "Max",
        "Min",
    }
)


def _ensure_dimensionless_literal(value: DimensionlessLiteral) -> float:
    """Convert a literal to a float while asserting it is dimensionless."""
    if isinstance(value, pint.Quantity):
        assert value.is_compatible_with(unit.dimensionless), f"Envelope value {value} is not dimensionless"
        return value.to(unit.dimensionless).magnitude
    if isinstance(value, (int, float)):
        return float(value)
    raise TypeError(f"Unsupported literal type {type(value)!r} for envelope values")


def _resolve_dimensionless(value: DimensionlessLiteral) -> float:
    return _ensure_dimensionless_literal(value)


def _validate_dimensionless(_, __, value: DimensionlessLiteral) -> None:
    match value:
        case int() | float():
            return
        case pint.Quantity():
            assert value.is_compatible_with(unit.dimensionless), f"Envelope expression {value} is not dimensionless"
        case _:
            raise TypeError(f"Unsupported expression type {type(value)!r} for envelope")


class Envelope(ABC):
    """Abstract base class representing a pulse envelope in the normalised time domain."""

    @abstractmethod
    def value(self, tau: np.ndarray) -> np.ndarray:
        """Return the envelope samples for normalised time points in [0, 1]."""
        ...

    def sample(self, *, num_samples: int = 201) -> tuple[np.ndarray, np.ndarray]:
        assert num_samples >= 2, "At least two samples are required to describe an envelope"
        tau = np.linspace(0.0, 1.0, num_samples)
        return tau, self.value(tau)

    def discretize(self, duration: pint.Quantity, dt: pint.Quantity) -> tuple[np.ndarray, np.ndarray, pint.Quantity]:
        """Discretise the envelope over a physical duration."""
        assert duration.is_compatible_with(unit.s), "Duration must be a time quantity"
        assert dt.is_compatible_with(unit.s), "dt must be a time quantity"
        duration_s = duration.to(unit.s).magnitude
        dt_s = dt.to(unit.s).magnitude
        assert dt_s > 0, "Sampling interval must be positive"
        sample_count = max(2, int(math.floor(duration_s / dt_s)) + 1)
        tau = np.linspace(0.0, 1.0, sample_count)
        time_axis = tau * duration_s
        return tau, self.value(tau), time_axis * unit.s

    def plot(
        self,
        *,
        num_samples: int = 2048,
        fft_padding_factor: int = 8,
        fft_significance_threshold: float = 5e-3,
        fft_x_padding_fraction: float = 0.1,
        show: bool = True,
    ) -> go.Figure | None:
        """Plot envelope time samples and FFT magnitude in normalized frequency space."""
        assert num_samples >= 8, f"num_samples must be at least 8 for FFT plotting, got {num_samples}"
        assert fft_padding_factor >= 1, f"fft_padding_factor must be >= 1, got {fft_padding_factor}"
        assert 0.0 < fft_significance_threshold < 1.0, (
            f"fft_significance_threshold must be in (0, 1), got {fft_significance_threshold}"
        )
        assert fft_x_padding_fraction >= 0.0, f"fft_x_padding_fraction must be >= 0, got {fft_x_padding_fraction}"

        tau = np.linspace(0.0, 1.0, num_samples, endpoint=False)
        values = np.asarray(self.value(tau), dtype=float)
        if values.ndim == 0:
            values = np.full_like(tau, float(values), dtype=float)
        else:
            values = np.asarray(np.broadcast_to(values, tau.shape), dtype=float)

        # Add explicit zeros before and after the envelope window in the time plot.
        dt = 1.0 / num_samples
        tau_with_edges = np.concatenate((np.array([-2.0 * dt, -1.0 * dt]), tau, np.array([1.0, 1.0 + dt])))
        values_with_edges = np.concatenate((np.array([0.0, 0.0]), values, np.array([0.0, 0.0])))

        # For FFT, use substantial zero-padding around the pulse to reveal the spectral envelope (e.g. sinc lobes).
        pad_samples = fft_padding_factor * num_samples
        fft_values = np.concatenate((np.zeros(pad_samples), values, np.zeros(pad_samples)))
        frequencies = np.fft.rfftfreq(fft_values.size, d=dt)
        spectrum = np.fft.rfft(fft_values)
        magnitude = np.abs(spectrum) / fft_values.size
        if magnitude.size > 1:
            magnitude[1:] *= 2.0

        # Hide DC in the FFT trace; otherwise it dominates and masks sidelobes.
        fft_frequencies = frequencies[1:]
        fft_magnitude = magnitude[1:]
        if fft_magnitude.size > 0:
            peak = float(np.max(fft_magnitude))
            significant = np.flatnonzero(fft_magnitude >= peak * fft_significance_threshold)
            if significant.size > 0:
                fmax_significant = float(fft_frequencies[int(significant[-1])])
                fmax_plot = min(fmax_significant * (1.0 + fft_x_padding_fraction), float(fft_frequencies[-1]))
            else:
                fmax_plot = float(fft_frequencies[-1])
        else:
            fmax_plot = 1.0
        fmax_plot = max(fmax_plot, float(fft_frequencies[0]) if fft_frequencies.size > 0 else 1.0)

        figure = make_subplots(
            rows=2,
            cols=1,
            vertical_spacing=0.14,
            subplot_titles=("Envelope (Time Domain)", "Envelope FFT Magnitude (Frequency Domain)"),
        )
        figure.add_trace(go.Scatter(x=tau_with_edges, y=values_with_edges, mode='lines', name='Envelope'), row=1, col=1)
        figure.add_trace(go.Scatter(x=fft_frequencies, y=fft_magnitude, mode='lines', name='|FFT|'), row=2, col=1)
        figure.update_xaxes(title_text="Normalized Time τ", row=1, col=1)
        figure.update_xaxes(title_text="Frequency (cycles per pulse duration)", row=2, col=1)
        figure.update_xaxes(range=[0.0, fmax_plot], row=2, col=1)
        figure.update_yaxes(title_text="Amplitude", row=1, col=1)
        figure.update_yaxes(title_text="Magnitude", row=2, col=1)
        figure.update_layout(
            height=700,
            margin=dict(l=60, r=20, t=50, b=50),
            template="plotly_white",
            showlegend=False,
        )

        if show:
            figure.show(config={"scrollZoom": True})
        else:
            return figure


def _parse_symbolic_expression(expression: str) -> sympy.Expr:
    """Parse and validate a symbolic envelope expression in the `tau` variable."""
    expression_text = expression.strip()
    assert expression_text, "Symbolic envelope expression must not be empty"
    parsed = sympy.sympify(
        expression_text,
        locals=_SYMPY_PARSE_LOCALS,
        evaluate=True,
    )
    assert isinstance(parsed, sympy.Expr), f"Expected a SymPy expression, got {type(parsed).__name__}"
    unexpected_symbols = sorted(str(symbol) for symbol in parsed.free_symbols if symbol != _TAU_SYMBOL)
    assert not unexpected_symbols, (
        f"Symbolic envelope expression may only use 'tau'. Found unexpected symbols: {unexpected_symbols}"
    )
    unexpected_functions = sorted(
        function.func.__name__
        for function in parsed.atoms(sympy.Function)
        if function.func.__name__ not in _SYMPY_ALLOWED_FUNCTION_NAMES
    )
    assert not unexpected_functions, (
        f"Symbolic envelope expression contains unsupported functions. Found: {unexpected_functions}"
    )
    return parsed


@lru_cache(maxsize=256)
def _compile_symbolic_expression(expression: str):
    """Compile a symbolic expression string into a NumPy-aware callable."""
    parsed = _parse_symbolic_expression(expression)
    return sympy.lambdify(_TAU_SYMBOL, parsed, modules=("numpy",))


@frozen
class SquareEnvelope(Envelope):
    level: DimensionlessLiteral = field(
        default=1 * unit.dimensionless,
        converter=lambda value: value * unit.dimensionless if isinstance(value, (int, float)) else value,
        repr=str,
        validator=_validate_dimensionless,
    )

    def value(self, tau: np.ndarray) -> np.ndarray:
        _ = tau  # Square envelope ignores position, keep signature symmetric
        level = _resolve_dimensionless(self.level)
        return np.full_like(tau, fill_value=level, dtype=float)


@frozen
class EnvelopeSegment:
    end: DimensionlessLiteral = field(
        converter=lambda value: value * unit.dimensionless if isinstance(value, (int, float)) else value,
        validator=_validate_dimensionless,
        repr=str,
    )
    value: DimensionlessLiteral = field(
        converter=lambda value: value * unit.dimensionless if isinstance(value, (int, float)) else value,
        validator=_validate_dimensionless,
        repr=str,
    )


@frozen
class PiecewiseEnvelope(Envelope):
    segments: tuple[EnvelopeSegment, ...] = field()

    def __attrs_post_init__(self) -> None:
        assert len(self.segments) > 0, "Piecewise envelope requires at least one segment"
        literal_segments = all(isinstance(segment.end, pint.Quantity) for segment in self.segments)
        if literal_segments:
            ends = [_ensure_dimensionless_literal(segment.end) for segment in self.segments]
            assert ends[0] > 0.0, "First segment in piecewise envelope must end after tau=0"
            assert all(l1 < l2 for l1, l2 in zip(ends, ends[1:])), "Segment ends must be strictly increasing"
            assert math.isclose(ends[-1], 1.0, rel_tol=1e-9), "Final segment must end at tau=1"

    def value(self, tau: np.ndarray) -> np.ndarray:
        ends = np.array([_resolve_dimensionless(segment.end) for segment in self.segments], dtype=float)
        values = np.array([_resolve_dimensionless(segment.value) for segment in self.segments], dtype=float)
        result = np.empty_like(tau, dtype=float)
        indices = np.searchsorted(ends, tau, side="right")
        indices = np.clip(indices, 0, len(values) - 1)
        result[:] = values[indices]
        return result


@frozen
class SymbolicEnvelope(Envelope):
    """Envelope defined by a SymPy expression over normalized time variable `tau`."""

    expression: str = field(validator=validators.instance_of(str))

    def __attrs_post_init__(self) -> None:
        _parse_symbolic_expression(self.expression)

    def value(self, tau: np.ndarray) -> np.ndarray:
        tau_array = np.asarray(tau, dtype=float)
        values = _compile_symbolic_expression(self.expression)(tau_array)
        values_array = np.asarray(values, dtype=float)
        if values_array.ndim == 0:
            return np.full_like(tau_array, fill_value=float(values_array), dtype=float)
        return np.asarray(np.broadcast_to(values_array, tau_array.shape), dtype=float)


def symbolic_envelope(*, expression: str) -> SymbolicEnvelope:
    """Build a symbolic envelope from an expression string in `tau`."""
    return SymbolicEnvelope(expression=expression)


def gaussian_envelope(
    *,
    amplitude: DimensionlessLiteral = 1 * unit.dimensionless,
    center: DimensionlessLiteral = 0.5 * unit.dimensionless,
    sigma: DimensionlessLiteral = 0.125 * unit.dimensionless,
    offset: DimensionlessLiteral = 0 * unit.dimensionless,
) -> SymbolicEnvelope:
    """
    Build a Gaussian envelope as a symbolic expression in normalized time.

    value(tau) = offset + amplitude * exp(-0.5 * ((tau - center) / sigma)^2)
    """

    amplitude_literal = _ensure_dimensionless_literal(amplitude)
    center_literal = _ensure_dimensionless_literal(center)
    sigma_literal = _ensure_dimensionless_literal(sigma)
    offset_literal = _ensure_dimensionless_literal(offset)
    assert sigma_literal > 0.0, f"gaussian_envelope sigma must be positive. Got {sigma_literal}"

    expression = sympy.Float(offset_literal) + sympy.Float(amplitude_literal) * sympy.exp(
        -sympy.Float(0.5) * ((_TAU_SYMBOL - sympy.Float(center_literal)) / sympy.Float(sigma_literal)) ** 2
    )
    return SymbolicEnvelope(expression=sympy.sstr(expression))


def sin_envelope() -> SymbolicEnvelope:
    """
    Build a smooth sinusoidal envelope with range [0, 1] and zero slope at tau=0 and tau=1.

    value(tau) = (cos(2 * pi * tau - pi) + 1) / 2
    """
    expression = (sympy.cos(2 * sympy.pi * _TAU_SYMBOL - sympy.pi) + 1) / 2
    return SymbolicEnvelope(expression=sympy.sstr(expression))
