"""Diagnostic baseline: classical regressors on simple spectral features, with
the same held-out split as train.py, to show what the spectrum identifies in
principle independently of the network.

    python3 ml/baseline.py --data DATASET_DIR [--val-every 8]

Features per clip, from the cached spectrogram: the mean log spectrum over
the whole clip, over the first 0.5 s, and over the last 1.5 s (1024 bins each,
pooled 4:1 to 256), and the frequencies of the first 8 spectral peaks of the
whole-clip mean. Models: ridge regression (alpha by cross-validation) and
gradient boosting, one per target, on standardized log targets.
"""

import argparse
import math
import os
import sys

import numpy as np
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.linear_model import RidgeCV
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from features import FS, N_FFT, HOP, build_cache  # noqa: E402
from train import TARGETS, DERIVED, metrics  # noqa: E402

DERIVED = dict(DERIVED, **{"kappa/area": lambda d: d["h"] + 0.5 * (d["E"] - d["rho"]) - d["Lx"] - d["Ly"]})


def peak_freqs(spec_db, n=8, min_prom=6.0):
    """Frequencies [Hz] of the first n peaks of a mean spectrum in dB, by increasing frequency; missing ones repeat the last."""
    from scipy.signal import find_peaks
    pk, _ = find_peaks(spec_db, prominence=min_prom)
    f = pk * FS / N_FFT
    f = list(f[:n])
    while len(f) < n:
        f.append(f[-1] if f else 0.0)
    return np.log(np.array(f) + 1.0)


def features_of(x):
    """x: (1024, T) in [0, 1] (dB / 80 + 1)."""
    db = (x.astype(np.float32) - 1) * 80
    fps = FS / HOP
    whole = db.mean(1)
    early = db[:, : int(0.5 * fps)].mean(1)
    late = db[:, int(1.5 * fps):].mean(1)
    pool = lambda v: v.reshape(-1, 4).mean(1)
    return np.concatenate([pool(whole), pool(early), pool(late), pool(early - late), peak_freqs(whole)])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--val-every", type=int, default=8)
    ap.add_argument("--max-p2p", type=float, default=None, help="keep only clips whose initial shape has a peak-to-peak height below this [m]")
    args = ap.parse_args()
    rows, feats = build_cache(args.data, os.path.join(args.data, "ml", "cache"))
    if args.max_p2p is not None:
        keep = [i for i, r in enumerate(rows) if r["peakToPeak"] < args.max_p2p]
        rows, feats = [rows[i] for i in keep], feats[keep]
    X = np.stack([features_of(f) for f in feats])
    y = np.stack([[math.log(r[t]) for t in TARGETS] for r in rows]).astype(np.float64)
    idx = np.array([r["index"] for r in rows])
    val = idx % args.val_every == 0
    tr = ~val
    print(f"{len(rows)} clips: {tr.sum()} train, {val.sum()} val; {X.shape[1]} features")
    mu = y[tr].mean(0)
    for name, make in [("ridge", lambda: make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-2, 4, 25)))),
                       ("gbdt", lambda: HistGradientBoostingRegressor(max_iter=400, learning_rate=0.05, max_leaf_nodes=15, l2_regularization=1.0))]:
        pred = np.zeros((val.sum(), len(TARGETS)))
        for i, t in enumerate(TARGETS):
            m = make().fit(X[tr], y[tr, i])
            pred[:, i] = m.predict(X[val])
        d_pred = {t: pred[:, i] for i, t in enumerate(TARGETS)}
        d_true = {t: y[val, i] for i, t in enumerate(TARGETS)}
        d_mean = {t: np.full(val.sum(), mu[i]) for i, t in enumerate(TARGETS)}
        print(f"--- {name}")
        for t in TARGETS:
            r = metrics(d_pred[t], d_true[t], d_mean[t])
            print(f"{t:>10}: rel RMS {r['rel_rms_pct']:5.1f}% (baseline {r['baseline_rel_rms_pct']:5.1f}%) R^2 {r['r2']:5.2f}")
        for t, f in DERIVED.items():
            r = metrics(f(d_pred), f(d_true), f(d_mean))
            print(f"{t:>10}: rel RMS {r['rel_rms_pct']:5.1f}% (baseline {r['baseline_rel_rms_pct']:5.1f}%) R^2 {r['r2']:5.2f}")


if __name__ == "__main__":
    main()