"""CPU-only analysis utilities for CCD pipeline.

No pycuda or LegacyCelerisLab dependency — can run with plain python3.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import numpy as np


# ---------------------------------------------------------------------------
# Period detection helpers
# ---------------------------------------------------------------------------

def detect_dominant_frequency(signal: np.ndarray, sample_dt: float) -> Tuple[float, float, float]:
    """Detect dominant frequency via FFT. Returns (f_dom, period, peak_power)."""
    n = len(signal)
    if n < 16:
        return 0.0, 0.0, 0.0
    y = signal - np.mean(signal)
    window = np.hanning(n)
    spec = np.abs(np.fft.rfft(y * window)) ** 2
    freqs = np.fft.rfftfreq(n, d=sample_dt)
    idx = 1 + np.argmax(spec[1:])
    f_dom = float(freqs[idx])
    period = 1.0 / f_dom if f_dom > 0 else 0.0
    return f_dom, period, float(spec[idx])


def detect_cycle_stability(signal: np.ndarray, sample_dt: float) -> Tuple[float, float, list]:
    """Detect cycle lengths using rising zero-crossings. Returns (CV_T, mean_T, lengths)."""
    y = signal - np.mean(signal)
    crossings = np.where((np.sign(y[:-1]) < 0) & (np.sign(y[1:]) > 0))[0]
    if len(crossings) < 2:
        return 0.0, 0.0, []
    cycle_lengths = np.diff(crossings).astype(float) * sample_dt
    if len(cycle_lengths) < 2:
        return 0.0, float(cycle_lengths[0]) if len(cycle_lengths) > 0 else 0.0, cycle_lengths.tolist()
    mean_T = float(np.mean(cycle_lengths))
    std_T = float(np.std(cycle_lengths))
    cv_T = std_T / mean_T if mean_T > 0 else 0.0
    return cv_T, mean_T, cycle_lengths.tolist()


# ---------------------------------------------------------------------------
# Phase resampling
# ---------------------------------------------------------------------------

def phase_resample(data: np.ndarray, cycle_starts: List[int], n_pts: int = 24) -> np.ndarray:
    """Resample a multi-channel signal to uniform phase points per cycle.

    Uses piecewise linear interpolation (no scipy dependency).

    Parameters
    ----------
    data : (T, C) ndarray — multi-channel time series.
    cycle_starts : list of int — indices where each cycle starts.
    n_pts : int — number of phase points per cycle.

    Returns
    -------
    resampled : (n_cycles, n_pts, C) ndarray
    """
    n_cycles = len(cycle_starts) - 1
    if n_cycles < 1:
        raise ValueError("Need at least 2 cycle starts")

    if data.ndim == 1:
        data = data[:, None]
    C = data.shape[1]
    out = np.zeros((n_cycles, n_pts, C), dtype=np.float64)

    for c in range(n_cycles):
        i_start = cycle_starts[c]
        i_end = cycle_starts[c + 1]
        segment = data[i_start:i_end + 1]
        seg_len = len(segment)
        if seg_len < 2:
            continue

        old_idx = np.linspace(0, 1, seg_len)
        new_idx = np.linspace(0, 1, n_pts, endpoint=False)
        for ch in range(C):
            out[c, :, ch] = np.interp(new_idx, old_idx, segment[:, ch])

    return out


# ---------------------------------------------------------------------------
# POD
# ---------------------------------------------------------------------------

def compute_pod(snapshot_matrix: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Compute POD from snapshot matrix using method of snapshots.

    For N << M (e.g. N=96 snapshots, M=2*NX*NY=1.3M), uses temporal
    correlation matrix (N x N) instead of full SVD on (M x N).
    This is ~100x faster and numerically identical.

    Returns (mean_field, modes, singular_values, coefficients).
    """
    mean_field = np.mean(snapshot_matrix, axis=1)
    Q = snapshot_matrix - mean_field[:, None]
    M, N = Q.shape

    if N < M // 10:  # method of snapshots: N << M
        # Temporal correlation: C = Q^T @ Q (N x N)
        C = Q.T.astype(np.float64) @ Q.astype(np.float64)
        s2, V = np.linalg.eigh(C)
        # Sort descending
        idx = np.argsort(s2)[::-1]
        s = np.sqrt(np.maximum(s2[idx], 0.0))
        V = V[:, idx]
        # Spatial modes: U = Q @ V * diag(1/s)
        U = Q.astype(np.float64) @ V
        U[:, s > 1e-12] /= s[s > 1e-12]
        # Coefficients: diag(s) @ V^T
        coefficients = (V * s).T
    else:  # fallback to full SVD for small matrices
        U, s, Vt = np.linalg.svd(Q, full_matrices=False)
        coefficients = np.diag(s) @ Vt

    return mean_field, U, s, coefficients


def cumulative_energy(singular_values: np.ndarray) -> np.ndarray:
    e = singular_values ** 2
    return np.cumsum(e) / np.sum(e)


def e95_index(cumulative_energy: np.ndarray) -> int:
    return int(np.searchsorted(cumulative_energy, 0.95) + 1)


# ---------------------------------------------------------------------------
# CCD (reduced, Lyu23-inspired)
# ---------------------------------------------------------------------------

def compute_reduced_ccd(pod_coeffs: np.ndarray, observable: np.ndarray, Q_delay: int = 12) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int]:
    """Compute reduced CCD in POD coefficient space, with clean boundary handling.

    Replaces np.roll-based delay with explicit lagged matrix construction,
    dropping wrapped-around samples instead of zero-padding them.

    Parameters
    ----------
    pod_coeffs : (r, N) ndarray — standardized POD coefficients.
    observable : (m, N) ndarray — standardized observable.
    Q_delay : int — number of delay steps (must be odd for symmetric window).

    Returns
    -------
    W : (r, min(r, m*Q_delay)) — CCD directions.
    sigma : (min_dim,) — singular values.
    R : (m*Q_delay, min_dim) — CCD response modes.
    z : (min_dim, N) — CCD temporal coefficients.
    N_original : int — original sample count.
    N_valid : int — sample count after boundary trimming.
    """
    N = pod_coeffs.shape[1]
    N_original = N
    m = observable.shape[0]

    half = Q_delay // 2
    # Keep only columns where ALL shifts have valid (non-wrapped) indices
    valid_start = half
    valid_end = N - half
    if Q_delay % 2 == 0:
        valid_end = N - half  # symmetric even: one extra on each side
    valid_end = N - (Q_delay - 1) // 2  # general formula
    # For symmetric Q_delay: valid = [half, N - half)
    valid_start = half
    valid_end = N - half
    N_valid = valid_end - valid_start

    if N_valid < Q_delay:
        raise ValueError(f"Too few valid samples ({N_valid}) for Q_delay={Q_delay}. Need at least {Q_delay}.")

    rows = []
    for shift in range(-half, half + 1):
        shifted = observable[:, valid_start + shift:valid_end + shift]
        rows.append(shifted)
    P = np.vstack(rows)  # (m*Q_delay, N_valid)

    A_valid = pod_coeffs[:, valid_start:valid_end]

    # Standardize
    P_mean = np.mean(P, axis=1, keepdims=True)
    P_std = np.std(P, axis=1, keepdims=True) + 1e-12
    P_z = (P - P_mean) / P_std

    A_mean = np.mean(A_valid, axis=1, keepdims=True)
    A_std = np.std(A_valid, axis=1, keepdims=True) + 1e-12
    A_z = (A_valid - A_mean) / A_std

    C = P_z @ A_z.T / (N_valid * np.sqrt(float(Q_delay)))
    R, s, Wt = np.linalg.svd(C, full_matrices=False)
    W = Wt.T
    z = W.T @ A_z
    return W, s, R, z, N_original, N_valid


# ---------------------------------------------------------------------------
# Field stacking helpers
# ---------------------------------------------------------------------------

def stack_velocity_fields(ux_fields: List[np.ndarray], uy_fields: List[np.ndarray]) -> np.ndarray:
    """Stack list of (ux, uy) field pairs into snapshot matrix (2*nx*ny, N)."""
    snapshots = []
    for ux, uy in zip(ux_fields, uy_fields):
        q = np.concatenate([ux.ravel(), uy.ravel()])
        snapshots.append(q)
    return np.column_stack(snapshots)


def unstack_velocity_modes(modes: np.ndarray, ny: int, nx: int, n_modes: int = 6) -> Tuple[List[np.ndarray], List[np.ndarray]]:
    """Unstack POD/CCD modes back into ux, uy fields."""
    ux_list, uy_list = [], []
    half = nx * ny
    for i in range(min(n_modes, modes.shape[1])):
        mode = modes[:, i]
        ux_list.append(mode[:half].reshape(ny, nx))
        uy_list.append(mode[half:].reshape(ny, nx))
    return ux_list, uy_list


# ---------------------------------------------------------------------------
# Aligned field data loader (new format: fields_aligned.npz + phase_plan.json)
# ---------------------------------------------------------------------------

def load_aligned_fields(scene_name: str) -> dict:
    """Load fields_aligned.npz + telemetry indexed by phase_plan step_indices.

    Transposes fields from (NX, NY) to (NY, NX) to match legacy convention
    so that ravel() produces (y*NX + x) ordering for POD/CCD.

    Parameters
    ----------
    scene_name : str — scene name (e.g. 'target_cylinder_1.0L', 'illusion_1.0L')

    Returns
    -------
    dict with:
        ux, uy : (N_snap, NY, NX) field snapshots
        forces : (N_snap, n_force) or None
        actions : (N_snap, 3) or None
        sensors : (N_snap, 6) or None
        meta : dict with gate, CV_T, f_dom, N_raw_per_cycle
        step_indices : list of int — matching phase_plan
    """
    import json
    import os

    from CCD_analysis.configs import DATA_DIR, SCENES

    if scene_name not in SCENES:
        raise KeyError(f"Unknown scene: {scene_name}")

    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(DATA_DIR, scene_id, scene_name)

    # -- fields_aligned.npz --
    fa_path = os.path.join(data_dir, "fields_aligned.npz")
    if not os.path.isfile(fa_path):
        raise FileNotFoundError(f"{fa_path} not found")

    fd = np.load(fa_path)
    ux_raw = fd["ux"]  # (N, NX, NY)
    uy_raw = fd["uy"]
    fd.close()

    N, NX, NY = ux_raw.shape
    # Transpose to (N, NY, NX) matching legacy convention
    ux = np.ascontiguousarray(ux_raw.transpose(0, 2, 1))
    uy = np.ascontiguousarray(uy_raw.transpose(0, 2, 1))

    # -- phase_plan.json --
    plan_path = os.path.join(DATA_DIR, "resampled", scene_name, "phase_plan.json")
    if not os.path.isfile(plan_path):
        raise FileNotFoundError(f"{plan_path} not found")

    with open(plan_path) as f:
        plan = json.load(f)

    step_indices = list(plan["step_indices"])
    assert len(step_indices) == N, (
        f"{scene_name}: step_indices has {len(step_indices)} entries but "
        f"fields has {N} snapshots"
    )

    # -- telemetry (controlled.npz or sensors.npz) --
    tele_path = None
    for p in [os.path.join(data_dir, "controlled.npz"), os.path.join(data_dir, "sensors.npz")]:
        if os.path.isfile(p):
            tele_path = p
            break
    if tele_path is None:
        raise FileNotFoundError(f"No telemetry (*.npz) found in {data_dir}")

    td = np.load(tele_path)

    result = {
        "ux": ux,
        "uy": uy,
        "forces": None,
        "actions": None,
        "sensors": None,
        "meta": {
            "scene": scene_name,
            "scene_id": scene_id,
            "gate": plan.get("gate", "unknown"),
            "CV_T": plan.get("CV_T"),
            "f_dom": plan.get("f_dom"),
            "N_raw_per_cycle": plan.get("N_raw_per_cycle"),
            "rho_interp": plan.get("rho_interp"),
            "sample_interval": cfg.get("sample_interval"),
        },
        "step_indices": step_indices,
    }

    for key in ["forces", "actions", "sensors"]:
        if key in td:
            full = td[key]
            result[key] = full[step_indices]  # (N_snap, n_channels)

    td.close()
    return result


def make_force_obs(forces: np.ndarray, scene_name: str, mode: str = "fy") -> np.ndarray:
    """Construct force observable from raw forces.

    Parameters
    ----------
    forces : (N, n_channels) ndarray
        target_cylinder: (N, 2) = [Fx, Fy]
        illusion/pinball: (N, 6) = [Fx1,Fy1,Fx2,Fy2,Fx3,Fy3]
    scene_name : str — used to detect target_cylinder vs illusion
    mode : str
        'fy' (primary)  -> SigmaFy (1 channel)
        'fx' (secondary) -> SigmaFx (1 channel)
        'joint' (supplementary) -> [SigmaFx; SigmaFy] (2 channels)

    Returns
    -------
    obs : (C, N) ndarray
    """
    if "target_cylinder" in scene_name:
        fx = forces[:, 0]
        fy = forces[:, 1]
    else:
        fx = forces[:, 0] + forces[:, 2] + forces[:, 4]
        fy = forces[:, 1] + forces[:, 3] + forces[:, 5]

    N = len(fx)
    if mode == "fy":
        return fy.reshape(1, N)
    elif mode == "fx":
        return fx.reshape(1, N)
    elif mode == "joint":
        return np.vstack([fx, fy])
    else:
        raise ValueError(f"Unknown force observable mode: {mode}")


def build_field_matrix(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Stack velocity field snapshots into a snapshot matrix.

    Parameters
    ----------
    ux, uy : (N, NY, NX) ndarray — field snapshots in legacy (NY, NX) order.

    Returns
    -------
    Q : (2 * NX * NY, N) ndarray — snapshot matrix for POD.
    """
    N = ux.shape[0]
    q_list = []
    for s in range(N):
        q_list.append(np.concatenate([ux[s].ravel(), uy[s].ravel()]))
    return np.column_stack(q_list).astype(np.float64)


def project_into_basis(ux: np.ndarray, uy: np.ndarray,
                       modes: np.ndarray, mean_f: np.ndarray) -> np.ndarray:
    """Project velocity fields onto a POD basis.

    Returns POD coefficients array of shape (r, N).
    """
    Q = build_field_matrix(ux, uy)
    return modes.T @ (Q - mean_f[:, None]).astype(np.float64)


# ---------------------------------------------------------------------------
# Harmonics analysis for illusion
# ---------------------------------------------------------------------------

def analyze_harmonics(states: np.ndarray, n_harmonics: int = 5) -> list:
    """FFT harmonics analysis. Returns list of dicts per channel."""
    N, D = states.shape
    result = []
    for d in range(D):
        y = states[:, d]
        fft_coef = np.fft.rfft(y)
        freqs = np.fft.rfftfreq(N, d=1)
        amps = 2.0 * np.abs(fft_coef) / N
        phases = np.angle(fft_coef)
        idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
        harmonics = {
            'dc': float(np.real(fft_coef[0]) / N),
            'amps': amps[idx].tolist(),
            'freqs': freqs[idx].tolist(),
            'phases': phases[idx].tolist(),
        }
        result.append(harmonics)
    return result


def gen_target_states_at(t, harmonics):
    """Reconstruct target observable at step index t from harmonics."""
    t = np.asarray(t)
    D = len(harmonics)
    result = np.zeros((t.size, D), dtype=np.float32)
    for d, h in enumerate(harmonics):
        val = np.full(t.shape, h['dc'], dtype=np.float32)
        for amp, freq, phase in zip(h['amps'], h['freqs'], h['phases']):
            val += amp * np.cos(2 * np.pi * freq * t + phase)
        result[:, d] = val
    if result.shape[0] == 1:
        return result[0]
    return result
