"""FFT harmonics utilities for Illusion target force reconstruction.

Used by stage_3_validate.py to reconstruct target cylinder force signals
from pre-computed FFT harmonics during closed-loop CFD validation.
"""

from __future__ import annotations

import numpy as np


def analyze_harmonics(states: np.ndarray, n_harmonics: int = 5) -> list[dict]:
    """Extract dominant FFT harmonics independently for each channel."""
    states = np.asarray(states)
    if states.ndim != 2 or states.shape[0] < 2:
        raise ValueError("states must have shape (samples, channels) with at least 2 samples")
    count, channels = states.shape
    result = []
    for channel in range(channels):
        fft = np.fft.rfft(states[:, channel])
        frequencies = np.fft.rfftfreq(count, d=1.0)
        amplitudes = 2.0 * np.abs(fft) / count
        indices = np.argsort(amplitudes[1:])[::-1][:n_harmonics] + 1
        result.append({
            "dc": float(np.real(fft[0]) / count),
            "amps": amplitudes[indices].tolist(),
            "freqs": frequencies[indices].tolist(),
            "phases": np.angle(fft)[indices].tolist(),
        })
    return result


def harmonics_document(harmonics: list[dict], channel_names: list[str], *, schema_version: str) -> dict:
    """Wrap harmonics with an explicit schema while retaining list-loader support."""
    if len(harmonics) != len(channel_names):
        raise ValueError("channel_names length must match harmonics")
    return {"schema_version": schema_version, "channel_names": list(channel_names), "harmonics": harmonics}


def _channels(harmonics: list | dict) -> list:
    if isinstance(harmonics, dict):
        try:
            harmonics = harmonics["harmonics"]
        except KeyError as exc:
            raise ValueError("harmonics document is missing 'harmonics'") from exc
    if not isinstance(harmonics, list):
        raise TypeError("harmonics must be a legacy list or a schema document")
    return harmonics


def gen_target_states_at(t: int | np.ndarray, harmonics: list | dict) -> np.ndarray:
    """Reconstruct channels at one or more phase indices.

    Both historical bare lists and v2 documents containing ``harmonics`` are
    accepted. A scalar phase returns ``(channels,)``; an array returns
    ``(phases, channels)``.
    """
    channels = _channels(harmonics)
    phases = np.atleast_1d(np.asarray(t, dtype=np.float64))
    result = np.zeros((phases.size, len(channels)), dtype=np.float32)
    for index, harmonic in enumerate(channels):
        values = np.full(phases.shape, float(harmonic["dc"]), dtype=np.float64)
        for amplitude, frequency, phase in zip(harmonic["amps"], harmonic["freqs"], harmonic["phases"]):
            values += amplitude * np.cos(2 * np.pi * frequency * phases + phase)
        result[:, index] = values
    return result[0] if np.asarray(t).ndim == 0 else result
