"""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)


# ---------------------------------------------------------------------------
# Weighted STLSQ with Huber-like robust regression
# ---------------------------------------------------------------------------

def _huber_weights(residuals: np.ndarray, c: float = 1.345) -> np.ndarray:
    """Compute Huber-like weights from residuals.

    Args:
        residuals: (T,) residual array.
        c: tuning constant (default 1.345 gives 95% efficiency for Normal errors).

    Returns:
        weights: (T,) weight array in [0, 1].
    """
    s = np.median(np.abs(residuals)) * 1.4826  # robust scale estimate (MAD)
    if s < 1e-12:
        s = 1.0
    r = np.abs(residuals) / s
    w = np.where(r <= c, 1.0, c / r)
    return np.asarray(w, dtype=np.float64)


def fit_sindy_weighted(
    Theta: np.ndarray,
    y: np.ndarray,
    thresholds: Optional[List[float]] = None,
    alpha: float = 1e-4,
    max_iter: int = 25,
    sample_weights: Optional[np.ndarray] = None,
    n_robust_passes: int = 2,
) -> List[dict]:
    """Run SINDy with threshold grid and optional robust weighting.

    Two-stage robust fitting:
    1. First pass: OLS, compute residuals, compute Huber weights
    2. Second pass: weighted STLSQ with Huber weights

    Args:
        Theta: (T, N) feature matrix.
        y: (T,) target.
        thresholds: list of threshold values.
        alpha: ridge regularization.
        max_iter: max STLSQ iterations.
        sample_weights: optional (T,) pre-defined sample weights.
        n_robust_passes: number of robust re-weighting passes (1 = skip).

    Returns:
        results: list of dict per threshold with keys:
            threshold, nz, r2, mae, coef, weights_used
    """
    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

    # Initialize weights
    if sample_weights is not None:
        w = np.asarray(sample_weights, dtype=np.float64).flatten()
        w = w / np.mean(w)  # normalize to mean=1
    else:
        w = np.ones(Theta.shape[0], dtype=np.float64)

    # Robust re-weighting passes
    for _ in range(n_robust_passes - 1):
        # OLS on weighted data
        Theta_w = Theta_s * np.sqrt(w)[:, None]
        y_w = y * np.sqrt(w)
        coef_ols, _, _, _ = np.linalg.lstsq(Theta_w, y_w, rcond=None)
        coef_ols = coef_ols.flatten() / std

        resid = y - Theta @ coef_ols
        w_new = _huber_weights(resid)
        if sample_weights is not None:
            w = w * w_new
        else:
            w = w_new
        w = w / np.mean(w)

    results = []
    for th in thresholds:
        if np.max(w) > 1e-8:
            # Weighted STLSQ: apply weights via sample_weight
            opt = ps.STLSQ(threshold=th, alpha=alpha, max_iter=max_iter)
            opt.fit(Theta_s, y, sample_weight=w)
            coef = np.asarray(opt.coef_, dtype=np.float64).flatten() / std
        else:
            coef = np.zeros(Theta.shape[1], dtype=np.float64)

        y_pred = Theta @ coef
        # Weighted R2
        ssr = float(np.sum(w * (y - y_pred) ** 2))
        sst = float(np.sum(w * (y - np.average(y, weights=w)) ** 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],
            "weights_min": float(np.min(w)),
            "weights_max": float(np.max(w)),
        })

    return results


# ---------------------------------------------------------------------------
# V2 feature matrix builder (configurable feature sets, time normalization)
# ---------------------------------------------------------------------------

def get_feature_matrix_v2(
    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 = False,       # default False for single-scene
    include_cos_sin: bool = False,  # default False (avoid masking linear structure)
    use_time_norm: bool = False,    # if True, use time-normalized deltas (da/dt_c)
    feat_keys: Optional[List[str]] = None,  # custom feat keys (default CORE_FEAT_KEYS_V2)
    dt_c: float = 1.0,             # control interval in T0 units
    n_warmup: int = 2,
    target_forces: Optional[np.ndarray] = None,  # (T, 2) raw lattice target forces
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[str], List[str]]:
    """Build feature matrices with configurable feature sets.

    Uses CORE_FEAT_KEYS_V2 by default (no sin_ua/cos_ua, no mu).
    For Illusion scenes with target forces, provide target_forces and
    feat_keys=ILLUSION_FEAT_KEYS_V2 (or pass None to auto-detect).
    For cross-Re joint fitting, set include_mu=True.
    For time-normalized deltas, set use_time_norm=True and provide dt_c.

    Returns
    -------
    Theta_front : (T-warmup, N_front) NO bias column
    Theta_rear : (T-warmup, N_rear) WITH bias column
    Y : (T-warmup, 3) target actions (physical omega)
    feat_names_front, feat_names_rear
    """
    from .feature_builder import (
        CORE_FEAT_KEYS_V2, TIME_FEAT_KEYS, MU_FEAT_KEYS, ALL_FEAT_KEYS_V2,
        ILLUSION_FEAT_KEYS_V2, ILLUSION_ALL_FEAT_KEYS_V2,
    )

    if feat_keys is None:
        if target_forces is not None:
            # Illusion scenes: include target force features
            feat_keys = ILLUSION_FEAT_KEYS_V2
        elif use_time_norm:
            # Replace discrete da with time-normalized da/dt
            feat_keys = [k for k in CORE_FEAT_KEYS_V2 if not k.startswith("da")]
            feat_keys += TIME_FEAT_KEYS
        else:
            feat_keys = CORE_FEAT_KEYS_V2

    if include_mu:
        # Add mu features if not already in feat_keys
        for mk in MU_FEAT_KEYS:
            if mk not in feat_keys:
                feat_keys = feat_keys + [mk]

    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,
        include_cos_sin=include_cos_sin,
        dt_c=dt_c,
        u0=u0,
        target_forces=target_forces,
    )

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

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

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


# ---------------------------------------------------------------------------
# Derivative-mode feature builders: fit d(alpha)/dt = g(physics_state)
# No action history in input features; Y is time-normalized action derivative.
# ---------------------------------------------------------------------------

def compute_action_deriv(
    actions_phys: np.ndarray,  # (T, 3) physical omega
    dt_c: float,               # control interval in T0 units
    u0: float = U0,
    center_diff: bool = False,  # if True, use (alpha(t) - alpha(t-2))/(2*dt_c)
) -> np.ndarray:
    """Compute time-normalized action derivative d(alpha)/dt.
    
    forward_diff:  d(alpha)/dt ~ (alpha(t) - alpha(t-1)) / dt_c
    center_diff:   d(alpha)/dt ~ (alpha(t) - alpha(t-2)) / (2*dt_c)
    
    Returns (T, 3) array, with first row(s) zero-padded.
    """
    alpha = np.asarray(actions_phys, dtype=np.float64) / u0  # non-dim
    T = alpha.shape[0]
    deriv = np.zeros_like(alpha)
    if center_diff and T >= 3:
        deriv[2:] = (alpha[2:] - alpha[:-2]) / (2.0 * dt_c)
    elif T >= 2:
        deriv[1:] = (alpha[1:] - alpha[:-1]) / dt_c
    return deriv


def get_feature_matrix_deriv(
    sensors: np.ndarray,        # (T, 6)
    forces: np.ndarray,         # (T, 6)
    actions_phys: np.ndarray,   # (T, 3) physical omega
    mu: float,
    u0: float = U0,
    dt_c: float = 1.0,          # control interval in T0 units
    feat_keys: Optional[List[str]] = None,  # default PHYSICS_FEAT_KEYS
    include_mu: bool = False,   # add mu modulation features (for cross-Re joint)
    target_forces: Optional[np.ndarray] = None,  # (T, 2) for Illusion
    n_warmup: int = 2,
    center_diff: bool = False,
    augment_level: int = 0,     # 0=static, 1=+lags, 2=+derivs, 3=both, 4=+action_lag
    output_mode: str = "deriv", # "deriv": predict d(alpha)/dt; "absolute": predict alpha directly
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[str], List[str]]:
    """Build feature matrices for phase-state SINDy.
    
    Input features: physics state with optional temporal context.
    
    Parameters
    ----------
    output_mode : str
        "deriv": Y = d(alpha)/dt (time-normalized). Closed-loop needs integration.
        "absolute": Y = alpha (non-dimensional action). No integration needed.
    
    augment_level : int
        0: PHYSICS_FEAT_KEYS only (static, no memory)
        1: + lag-1 obs features (adds temporal context)
        2: + obs derivative features (adds rate information)
        3: + lag-1 + derivative (full temporal context)
        4: + action lag (aF_lag1 etc., for ablation comparison against old approach)
    """
    from .feature_builder import (
        PHYSICS_FEAT_KEYS, ILLUSION_ERR_KEYS, MU_FEAT_KEYS,
        AUG_LEVEL_1_KEYS, AUG_LEVEL_2_KEYS,
        AUG_LEVEL_3_KEYS, AUG_LEVEL_4_KEYS,
        ILLUSION_AUG_LEVEL_1_KEYS, ILLUSION_AUG_LEVEL_2_KEYS,
        ILLUSION_AUG_LEVEL_3_KEYS, ILLUSION_AUG_LEVEL_4_KEYS,
    )
    
    # Select feature key set based on augment_level
    _AUG_MAP = {0: PHYSICS_FEAT_KEYS, 1: AUG_LEVEL_1_KEYS,
                2: AUG_LEVEL_2_KEYS, 3: AUG_LEVEL_3_KEYS, 4: AUG_LEVEL_4_KEYS}
    _AUG_ILLUSION_MAP = {0: ILLUSION_ERR_KEYS, 1: ILLUSION_AUG_LEVEL_1_KEYS,
                         2: ILLUSION_AUG_LEVEL_2_KEYS, 3: ILLUSION_AUG_LEVEL_3_KEYS,
                         4: ILLUSION_AUG_LEVEL_4_KEYS}
    
    if feat_keys is None:
        if target_forces is not None:
            feat_keys = _AUG_ILLUSION_MAP.get(augment_level, ILLUSION_AUG_LEVEL_3_KEYS)
        else:
            feat_keys = _AUG_MAP.get(augment_level, AUG_LEVEL_3_KEYS)
    
    if include_mu:
        for mk in MU_FEAT_KEYS:
            if mk not in feat_keys:
                feat_keys = feat_keys + [mk]
    
    if include_mu:
        for mk in MU_FEAT_KEYS:
            if mk not in feat_keys:
                feat_keys = feat_keys + [mk]
    
    T = sensors.shape[0]
    # a_prev/a_prev2 only used for computing the DERIVATIVE target Y, not in Theta
    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=False, include_mu=include_mu,
        include_cos_sin=False, dt_c=dt_c, u0=u0,
        target_forces=target_forces,
        sensors_raw=sensors, forces_raw=forces,
    )
    
    Theta_f = build_feature_matrix(sym, feat_keys, add_bias=False)
    Theta_r = build_feature_matrix(sym, feat_keys, add_bias=True)
    
    feat_names_front = get_feature_names(feat_keys, add_bias=False)
    feat_names_rear = get_feature_names(feat_keys, add_bias=True)
    
    # Y: depends on output_mode
    if output_mode == "absolute":
        Y = np.asarray(actions_phys, dtype=np.float64) / u0  # non-dim alpha (absolute action)
    else:
        Y = compute_action_deriv(actions_phys, dt_c, u0=u0, center_diff=center_diff)
    
    return (Theta_f[n_warmup:], Theta_r[n_warmup:],
            Y[n_warmup:],
            feat_names_front, feat_names_rear)
