"""Math wave LED stream — configurable waveform generator with gradient mapping.

Produces per-LED values from superimposed wave functions (sine, triangle,
sawtooth, square) that vary across LED position and time, then maps the
combined value through a gradient palette for color output.
"""

import math
import threading
import time
from typing import List, Optional

import numpy as np

from ledgrab.core.processing.color_strip_stream import (
    ColorStripStream,
    _compute_gradient_colors,
)
from ledgrab.utils import get_logger
from ledgrab.utils.timer import high_resolution_timer

logger = get_logger(__name__)

TWO_PI = 2.0 * math.pi


# ── Vectorized waveform functions ────────────────────────────────────
# Input x is in radians; output is normalised to [0, 1].


def _wave_sine(x: np.ndarray) -> np.ndarray:
    return 0.5 + 0.5 * np.sin(x)


def _wave_triangle(x: np.ndarray) -> np.ndarray:
    phase = np.mod(x / TWO_PI, 1.0)
    return 1.0 - 2.0 * np.abs(phase - 0.5)


def _wave_sawtooth(x: np.ndarray) -> np.ndarray:
    return np.mod(x / TWO_PI, 1.0)


def _wave_square(x: np.ndarray) -> np.ndarray:
    return (np.sin(x) >= 0).astype(np.float32)


_WAVEFORM_FN = {
    "sine": _wave_sine,
    "triangle": _wave_triangle,
    "sawtooth": _wave_sawtooth,
    "square": _wave_square,
}

MAX_WAVE_LAYERS = 8


class MathWaveColorStripStream(ColorStripStream):
    """Color strip stream generating LED colors from mathematical waves.

    Each LED at normalised position p gets a combined wave value from all
    layers, clamped to [0, 1], then mapped through a gradient LUT to RGB.
    """

    def __init__(self, source):
        self._colors_lock = threading.Lock()
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._fps = 30
        self._frame_time = 1.0 / 30
        self._clock = None
        self._led_count = 1
        self._auto_size = True
        self._gradient_store = None
        self._gradient_lut: Optional[np.ndarray] = None
        self._cached_gradient_id: Optional[str] = None
        self._update_from_source(source)

    def _update_from_source(self, source) -> None:
        from ledgrab.storage.bindable import bfloat

        self._speed = bfloat(getattr(source, "speed", 1.0), 1.0)
        self._gradient_id = getattr(source, "gradient_id", None)
        raw_waves = getattr(source, "waves", None) or []
        self._waves: List[dict] = raw_waves[:MAX_WAVE_LAYERS]
        with self._colors_lock:
            self._colors: Optional[np.ndarray] = None
        # Invalidate LUT so it rebuilds on next frame
        self._cached_gradient_id = None

    def configure(self, device_led_count: int) -> None:
        if self._auto_size and device_led_count > 0:
            new_count = max(self._led_count, device_led_count)
            if new_count != self._led_count:
                self._led_count = new_count

    @property
    def target_fps(self) -> int:
        return self._fps

    @property
    def led_count(self) -> int:
        return self._led_count

    def set_capture_fps(self, fps: int) -> None:
        self._fps = max(1, min(90, fps))
        self._frame_time = 1.0 / self._fps

    def start(self) -> None:
        if self._running:
            return
        self._running = True
        self._thread = threading.Thread(
            target=self._animate_loop,
            name="css-math-wave",
            daemon=True,
        )
        self._thread.start()
        logger.info(
            f"MathWaveColorStripStream started (leds={self._led_count}, "
            f"waves={len(self._waves)})"
        )

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
            if self._thread.is_alive():
                logger.warning("MathWaveColorStripStream thread did not terminate within 5s")
            self._thread = None
        logger.info("MathWaveColorStripStream stopped")

    def get_latest_colors(self) -> Optional[np.ndarray]:
        with self._colors_lock:
            return self._colors

    def update_source(self, source) -> None:
        from ledgrab.storage.color_strip_source import MathWaveColorStripSource

        if isinstance(source, MathWaveColorStripSource):
            prev_led_count = self._led_count if self._auto_size else None
            self._update_from_source(source)
            if prev_led_count and self._auto_size:
                self._led_count = prev_led_count
            logger.info("MathWaveColorStripStream params updated in-place")

    def set_clock(self, clock) -> None:
        self._clock = clock

    def set_gradient_store(self, gradient_store) -> None:
        self._gradient_store = gradient_store

    # ── Gradient LUT ───────────────────────────────────────────────

    def _ensure_gradient_lut(self) -> np.ndarray:
        """Build or return a 256-entry RGB lookup table from the current gradient."""
        gid = self._gradient_id
        if gid == self._cached_gradient_id and self._gradient_lut is not None:
            return self._gradient_lut

        self._cached_gradient_id = gid
        stops = None
        if gid and self._gradient_store:
            stops = self._gradient_store.resolve_stops(gid)
        if not stops:
            # Fallback: rainbow-ish default
            stops = [
                {"position": 0.0, "color": [255, 0, 0]},
                {"position": 0.25, "color": [255, 255, 0]},
                {"position": 0.5, "color": [0, 255, 0]},
                {"position": 0.75, "color": [0, 255, 255]},
                {"position": 1.0, "color": [0, 0, 255]},
            ]
        self._gradient_lut = _compute_gradient_colors(stops, 256)
        return self._gradient_lut

    # ── Animation loop ─────────────────────────────────────────────

    def _animate_loop(self) -> None:
        _pool_n = 0
        _buf_a = _buf_b = None
        _use_a = True
        _positions: Optional[np.ndarray] = None  # normalised LED positions [0,1]

        try:
            with high_resolution_timer():
                while self._running:
                    wall_start = time.perf_counter()
                    frame_time = self._frame_time
                    try:
                        clock = self._clock
                        if clock:
                            if not clock.is_running:
                                time.sleep(0.1)
                                continue
                            t = clock.get_time()
                            speed = clock.speed * self.resolve("speed", self._speed)
                        else:
                            t = wall_start
                            speed = self.resolve("speed", self._speed)

                        n = self._led_count
                        if n != _pool_n:
                            _pool_n = n
                            _buf_a = np.empty((n, 3), dtype=np.uint8)
                            _buf_b = np.empty((n, 3), dtype=np.uint8)
                            _positions = (
                                np.linspace(0.0, 1.0, n, dtype=np.float32)
                                if n > 1
                                else np.array([0.0], dtype=np.float32)
                            )

                        buf = _buf_a if _use_a else _buf_b
                        _use_a = not _use_a

                        self._render_waves(buf, n, _positions, t, speed)

                        with self._colors_lock:
                            self._colors = buf
                    except Exception as e:
                        logger.error(f"MathWaveColorStripStream animation error: {e}")

                    elapsed = time.perf_counter() - wall_start
                    time.sleep(max(frame_time - elapsed, 0.001))
        except Exception as e:
            logger.error(f"Fatal MathWaveColorStripStream loop error: {e}", exc_info=True)
        finally:
            self._running = False

    # ── Render ─────────────────────────────────────────────────────

    def _render_waves(
        self,
        buf: np.ndarray,
        n: int,
        positions: np.ndarray,
        t: float,
        speed: float,
    ) -> None:
        """Compute superimposed wave values and map through gradient LUT."""
        combined = np.zeros(n, dtype=np.float32)

        for wave in self._waves:
            waveform = wave.get("waveform", "sine")
            fn = _WAVEFORM_FN.get(waveform, _wave_sine)
            freq = float(wave.get("frequency", 1.0))
            amp = float(wave.get("amplitude", 1.0))
            phase = float(wave.get("phase", 0.0))
            offset = float(wave.get("offset", 0.0))

            # x = 2*pi * frequency * (position + speed*t) + phase
            x = TWO_PI * freq * (positions + speed * t) + phase
            combined += amp * fn(x) + offset

        # Clamp to [0, 1]
        np.clip(combined, 0.0, 1.0, out=combined)

        # Map through gradient LUT (256-entry)
        lut = self._ensure_gradient_lut()
        indices = (combined * 255.0).astype(np.uint8)
        buf[:] = lut[indices]
