"""
Offline baseline comparison: HalfSpaceTrees (online) vs Isolation Forest (batch) vs Z-Score (statistical).

Run from the project root:
    python scripts/compare_baselines.py

The script generates the full 800-observation synthetic stream (same seed as the live demo),
trains each method on Phase A alone, scores Phase C, and prints a comparison table.
It also measures drift detection latency for ADWIN.
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

# Make project root importable when run as a script
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import numpy as np
from sklearn.ensemble import IsolationForest
from sklearn.metrics import f1_score, precision_score, recall_score, roc_auc_score

from src.config import settings
from src.detector.anomaly import AnomalyDetector
from src.detector.drift import DriftDetector
from src.pipeline.runner import _drift_signal
from src.stream.generator import _generate_sync


def _features(obs_list):
    return np.array([[o.temperature, o.pressure, o.vibration] for o in obs_list])


def _binary_labels(obs_list):
    return np.array([0 if o.label == "normal" else 1 for o in obs_list])


def run_comparison(seed: int = 42, verbose: bool = True) -> dict:
    cfg = settings.stream
    cfg_d = settings.detector
    cfg_dr = settings.drift

    observations = list(
        _generate_sync(
            phase_a_length=cfg.phase_a_length,
            phase_b_length=cfg.phase_b_length,
            phase_c_length=cfg.phase_c_length,
            drift_magnitude=cfg.drift_magnitude,
            anomaly_rate=cfg.anomaly_rate,
            point_ratio=cfg.point_ratio,
            seed=seed,
        )
    )

    pa_end = cfg.phase_a_length
    pb_end = pa_end + cfg.phase_b_length
    phase_a = observations[:pa_end]
    phase_b = observations[pa_end:pb_end]
    phase_c = observations[pb_end:]

    X_a = _features(phase_a)
    X_b = _features(phase_b)
    X_c = _features(phase_c)
    y_c = _binary_labels(phase_c)

    results = {}

    # ── 1. Z-Score (statistical, fit on Phase A) ─────────────────────────────────
    mu  = X_a.mean(axis=0)
    std = X_a.std(axis=0) + 1e-9
    t0  = time.perf_counter()
    z_scores = np.abs((X_c - mu) / std).max(axis=1)
    z_preds  = (z_scores > 2.5).astype(int)
    z_time   = (time.perf_counter() - t0) * 1000
    results["Z-Score (max |z|>2.5)"] = _metrics(y_c, z_preds, z_scores, z_time, batch=True)

    # ── 2. Isolation Forest (batch, train on Phase A) ────────────────────────────
    t0  = time.perf_counter()
    iso = IsolationForest(
        n_estimators=100,
        contamination=cfg.anomaly_rate,
        random_state=seed,
    )
    iso.fit(X_a)
    # score_samples returns negative anomaly scores; negate for "higher = more anomalous"
    iso_scores_raw = -iso.score_samples(X_c)
    iso_preds      = (iso.predict(X_c) == -1).astype(int)
    iso_time       = (time.perf_counter() - t0) * 1000
    results["Isolation Forest (batch)"] = _metrics(y_c, iso_preds, iso_scores_raw, iso_time, batch=True)

    # ── 3. HalfSpaceTrees + ADWIN (online, streaming) ───────────────────────────
    detector = AnomalyDetector(
        threshold=cfg_d.threshold,
        n_trees=cfg_d.n_trees,
        height=cfg_d.height,
        window_size=cfg_d.window_size,
        seed=cfg_d.seed,
    )
    drift_det = DriftDetector(delta=cfg_dr.delta, grace_period=cfg_dr.grace_period)

    hst_scores: list[float] = []
    hst_preds:  list[int]   = []
    latencies:  list[float] = []

    t0 = time.perf_counter()
    _frozen = False
    for i, obs in enumerate(observations):
        t_obs = time.perf_counter()
        if obs.phase == "A":
            detector.learn_scaler(obs)
        elif not _frozen:
            detector.freeze_baseline()
            _frozen = True
        score = detector.score(obs)
        drift_det.update(_drift_signal(obs))
        latencies.append((time.perf_counter() - t_obs) * 1000)

        if i >= pb_end:
            hst_scores.append(score)
            hst_preds.append(1 if score > cfg_d.threshold else 0)

    hst_time = (time.perf_counter() - t0) * 1000
    results["Mahalanobis + ADWIN (online)"] = _metrics(
        y_c, np.array(hst_preds), np.array(hst_scores), hst_time, batch=False
    )
    results["Mahalanobis + ADWIN (online)"]["mean_latency_us"] = np.mean(latencies) * 1000

    # ── Drift detection latency ─────────────────────────────────────────────────
    # Re-run to find first ADWIN fire in Phase B
    detector2  = AnomalyDetector(threshold=cfg_d.threshold, n_trees=cfg_d.n_trees, height=cfg_d.height, window_size=cfg_d.window_size, seed=cfg_d.seed)
    drift_det2 = DriftDetector(delta=cfg_dr.delta, grace_period=cfg_dr.grace_period)
    first_drift_obs = None
    _frozen2 = False

    for i, obs in enumerate(observations):
        if obs.phase == "A":
            detector2.learn_scaler(obs)
        elif not _frozen2:
            detector2.freeze_baseline()
            _frozen2 = True
        detector2.score(obs)
        drift_det2.update(_drift_signal(obs))
        if drift_det2.drift_detected and first_drift_obs is None:
            first_drift_obs = i

    drift_latency = (first_drift_obs - pa_end) if first_drift_obs is not None else None

    if verbose:
        _print_report(results, y_c, phase_a, phase_b, phase_c, drift_latency, cfg)

    return results


def _metrics(y_true, y_pred, scores, elapsed_ms, *, batch: bool) -> dict:
    # Guard against all-zero predictions for AUC
    try:
        auc = roc_auc_score(y_true, scores)
    except ValueError:
        auc = float("nan")
    return {
        "precision": precision_score(y_true, y_pred, zero_division=0),
        "recall":    recall_score(y_true, y_pred, zero_division=0),
        "f1":        f1_score(y_true, y_pred, zero_division=0),
        "roc_auc":   auc,
        "elapsed_ms": elapsed_ms,
        "batch": batch,
    }


def _print_report(results, y_c, phase_a, phase_b, phase_c, drift_latency, cfg):
    sep = "-" * 72

    print()
    print("  Offline Baseline Comparison — Real-Time Anomaly Detection")
    print(sep)
    print(f"  Stream:  {len(phase_a)} Phase-A | {len(phase_b)} Phase-B | {len(phase_c)} Phase-C")
    print(f"  Anomalies in Phase C: {int(y_c.sum())} / {len(y_c)}  ({100*y_c.mean():.1f}%)")
    print()
    print(f"  {'Method':<38}  {'Prec':>6}  {'Recall':>6}  {'F1':>6}  {'AUC':>6}  {'Mode'}")
    print(sep)

    for name, m in results.items():
        mode = "batch" if m["batch"] else "online"
        auc  = f"{m['roc_auc']:.3f}" if not (m['roc_auc'] != m['roc_auc']) else " n/a "
        print(
            f"  {name:<38}  {m['precision']:>6.3f}  {m['recall']:>6.3f}  {m['f1']:>6.3f}  {auc:>6}  {mode}"
        )

    print(sep)
    print()
    print("  Latency")
    print(sep)
    for name, m in results.items():
        if not m["batch"]:
            lat = m.get("mean_latency_us")
            if lat is not None:
                print(f"  {name:<38}  {lat:.2f} µs / observation (mean)")
        else:
            total = m["elapsed_ms"]
            per   = total / len(y_c)
            print(f"  {name:<38}  {per:.3f} ms / observation  ({total:.1f} ms total, amortised)")

    print()
    print("  Drift detection")
    print(sep)
    if drift_latency is not None:
        print(f"  ADWIN first fired at observation {cfg.phase_a_length + drift_latency}")
        print(f"  → {drift_latency} observations into Phase B ({100*drift_latency/cfg.phase_b_length:.1f}% of drift phase)")
    else:
        print("  ADWIN did not fire during Phase B (try lowering delta).")
    print()
    print("  Notes")
    print(sep)
    print("  Z-Score checks each sensor independently, so it misses contextual anomalies")
    print("  where no single sensor looks extreme but the combination is unusual (e.g.,")
    print("  high temp + low pressure when they normally move together).")
    print("  Mahalanobis uses the full covariance matrix and catches both types.")
    print("  The baseline is fit from Phase A data and stays frozen after that.")
    print()


if __name__ == "__main__":
    run_comparison()