"""SINDy fitting utilities: STLSQ threshold grid, feature matrix building.

All features are built using the unified feature_builder module.
"""
from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np

from .feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    get_feature_names, ALL_FEAT_KEYS, U0,
)

# Default thresholds used across all scenes
DEFAULT_THRESHOLDS = [0.0, 0.001, 0.002, 0.005, 0.01, 0.015, 0.02, 0.03, 0.05, 0.1]


def fit_channel(
    Theta: np.ndarray,
    y: np.ndarray,
    thresholds: Optional[List[float]] = None,
    alpha: float = 1e-4,
    max_iter: int = 25,
) -> Tuple[List[dict], dict]:
    """Fit a single channel (one cylinder) with STLSQ threshold grid.

    Returns
    -------
    rows : list of dict per threshold
    best : dict with best threshold entry (highest R2)
    """
    import pysindy as ps

    if thresholds is None:
        thresholds = DEFAULT_THRESHOLDS

    # Normalise features for thresholding stability
    std = np.std(Theta, axis=0)
    std = np.where(std < 1e-8, 1.0, std)
    Theta_s = Theta / std

    best = None
    rows = []
    for th in thresholds:
        opt = ps.STLSQ(threshold=th, alpha=alpha, max_iter=max_iter)
        opt.fit(Theta_s, y)
        coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std
        y_pred = Theta @ coef
        ssr = float(np.sum((y - y_pred) ** 2))
        sst = float(np.sum((y - np.mean(y)) ** 2) + 1e-12)
        r2 = 1.0 - ssr / sst
        mae = float(np.mean(np.abs(y - y_pred)))
        nz = int(np.sum(np.abs(coef) > 1e-8))
        entry = {"threshold": float(th), "nz": nz, "r2": r2, "mae": mae, "coef": coef}
        rows.append(entry)
        if best is None or r2 > best["r2"]:
            best = entry
    return rows, best


def fit_sindy(
    Theta: np.ndarray,
    y: np.ndarray,
    thresholds: Optional[List[float]] = None,
) -> List[dict]:
    """Run SINDy with threshold grid, return results list.

    Each result dict has keys: threshold, nz, r2, mae, coef.
    """
    if thresholds is None:
        thresholds = DEFAULT_THRESHOLDS

    std = np.std(Theta, axis=0)
    std = np.where(std < 1e-8, 1.0, std)
    Theta_s = Theta / std

    results = []
    for th in thresholds:
        import pysindy as ps
        opt = ps.STLSQ(threshold=th, alpha=1e-4, max_iter=25)
        opt.fit(Theta_s, y)
        coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std

        y_pred = Theta @ coef
        ssr = float(np.sum((y - y_pred) ** 2))
        sst = float(np.sum((y - np.mean(y)) ** 2) + 1e-12)
        r2 = 1.0 - ssr / sst
        mae = float(np.mean(np.abs(y - y_pred)))
        nz = int(np.sum(np.abs(coef) > 1e-8))

        results.append({
            "threshold": float(th), "nz": nz, "r2": r2,
            "mae": mae, "coef": [float(c) for c in coef],
        })

    return results


def print_control_law(feature_names: List[str], coef: np.ndarray, channel_label: str = "ch"):
    """Pretty-print a sparse control law."""
    terms = []
    for i, c in enumerate(coef):
        if abs(c) > 1e-8:
            terms.append(f"{c:.6f} * {feature_names[i]}")
    print(f"  {channel_label}:  {' + '.join(terms)}")
    nz = sum(1 for c in coef if abs(c) > 1e-8)
    print(f"       non-zero terms: {nz}")


def get_active_support(
    coef: np.ndarray,
    feat_names: List[str],
    relative_threshold: float = 0.02,
) -> Dict[str, float]:
    """Extract active features from coefficient vector.

    Features with |coef| / max(|coef|) >= relative_threshold are considered active.
    """
    max_c = np.max(np.abs(coef))
    if max_c < 1e-12:
        return {}
    active = {}
    for name, c in zip(feat_names, coef):
        if abs(c) / max_c >= relative_threshold:
            active[name] = float(c)
    return active


def get_feature_matrix_from_data(
    sensors: np.ndarray,    # (T, 6)
    forces: np.ndarray,     # (T, 6)
    actions_phys: np.ndarray,  # (T, 3) physical omega
    mu: float,
    u0: float = U0,
    alpha_mode: bool = False,
    include_mu: bool = True,
    n_warmup: int = 2,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[str], List[str]]:
    """Build feature matrices from raw CFD data.

    Constructs dimensionless features via feature_builder, creates front (no bias)
    and rear (with bias) feature matrices, and returns them aligned with Y.

    Parameters
    ----------
    sensors, forces, actions_phys : raw data arrays.
    mu : 1/Re_D.
    u0 : inlet velocity (lattice units).
    alpha_mode : if True, actions_phys are already nondim alpha.
    include_mu : include mu modulation features.
    n_warmup : number of warmup steps to discard (default 2 for lag/da).

    Returns
    -------
    Theta_front : (T-warmup, N_front) feature matrix, NO bias column
    Theta_rear : (T-warmup, N_rear) feature matrix, WITH bias column
    Y : (T-warmup, 3) target action matrix
    feat_names_front : list of feature names for front
    feat_names_rear : list of feature names for rear
    """
    T = sensors.shape[0]
    a_prev = np.zeros((T, 3), dtype=np.float64)
    a_prev2 = np.zeros((T, 3), dtype=np.float64)
    a_prev[1:] = actions_phys[:-1]
    a_prev2[2:] = actions_phys[:-2]

    dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)
    sym = compute_features(dim, a_prev, a_prev2, mu,
                           alpha_mode=alpha_mode, include_mu=include_mu, u0=u0)

    Theta_f = build_feature_matrix(sym, ALL_FEAT_KEYS, add_bias=False)
    Theta_r = build_feature_matrix(sym, ALL_FEAT_KEYS, add_bias=True)

    feat_names_front = get_feature_names(ALL_FEAT_KEYS, add_bias=False)
    feat_names_rear = get_feature_names(ALL_FEAT_KEYS, add_bias=True)

    return (Theta_f[n_warmup:], Theta_r[n_warmup:],
            actions_phys[n_warmup:],
            feat_names_front, feat_names_rear)
