#!/usr/bin/env python3
"""Stage 02b — Diagnostic plots per run candidate.

Reads 02_runs.json + replays MCAP rel_alt + /mavros/state to produce
4-panel PNG per run (validated + rejected).

Output:
    data/<MISSION>/02_runs_diag/<RUN_ID>.png
    data/<MISSION>/02_runs_diag/index.json
    data/<MISSION>/02_runs_diag/index.html

Usage:
    python3 02b_runs_diag.py --mission 20260505-Lepradet \
        [--ssd-base /mnt/ssd] [--out data/] [--padding-s 120]
"""

import argparse
import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

from mcap.reader import make_reader
from mcap_ros2.decoder import DecoderFactory

TOPIC_REL_ALT = "/mavros/global_position/rel_alt"
TOPIC_STATE = "/mavros/state"
TOPIC_HEADING = "/mavros/global_position/compass_hdg"

AUV_PHYSICAL_MAP = {"AUV010": "AUV210", "AUV012": "AUV212", "AUV013": "AUV213"}
PHYS_TO_MCAP = {v: k for k, v in AUV_PHYSICAL_MAP.items()}


def extract_auv_id(folder_name):
    m = re.search(r"(AUV\d+)$", folder_name)
    return m.group(1) if m else None


def gather_signals_for_auv(bags_root, auv_mcap_id, t_start, t_end):
    """Read rel_alt + state + heading from all bags for this AUV in [t_start,t_end]."""
    alt = []  # (t, v)
    state = []  # (t, armed, mode)
    hdg = []  # (t, v_deg)
    for d in sorted(bags_root.iterdir()):
        if not d.is_dir():
            continue
        if extract_auv_id(d.name) != auv_mcap_id:
            continue
        for mcap_path in sorted(d.glob("*.mcap")):
            try:
                with open(mcap_path, "rb") as fp:
                    reader = make_reader(fp, decoder_factories=[DecoderFactory()])
                    for _schema, channel, message, ros_msg in reader.iter_decoded_messages(
                            topics=[TOPIC_REL_ALT, TOPIC_STATE, TOPIC_HEADING]):
                        ts = message.log_time / 1e9
                        if ts < t_start or ts > t_end:
                            continue
                        if channel.topic == TOPIC_REL_ALT:
                            alt.append((ts, float(ros_msg.data)))
                        elif channel.topic == TOPIC_STATE:
                            state.append((ts, bool(ros_msg.armed), str(ros_msg.mode)))
                        elif channel.topic == TOPIC_HEADING:
                            hdg.append((ts, float(ros_msg.data)))
            except Exception as exc:
                print(f"  [WARN] skip {mcap_path.name}: {exc}", file=sys.stderr)
    alt.sort()
    state.sort()
    hdg.sort()
    return alt, state, hdg


def plot_run_diag(run, alt, state, hdg, filter_params, status, reject_reason, out_path):
    """4-panel diagnostic plot."""
    start_e = run["start_epoch"]
    end_e = run["end_epoch"]
    dur = end_e - start_e
    run_id = run["run_id"]

    # Convert epoch → relative seconds from start_e for x-axis readability
    def to_rel(t):
        return t - start_e

    fig, axes = plt.subplots(4, 1, figsize=(12, 11), gridspec_kw={"height_ratios": [3, 1, 2, 2]})
    fig.suptitle(
        f"{run_id} — {status}  |  start={datetime.fromtimestamp(start_e, tz=timezone.utc).strftime('%H:%M:%S UTC')}  duration={dur:.0f}s",
        fontsize=12, fontweight="bold",
        color=("#16a34a" if status == "OK" else "#dc2626"),
    )

    # === Panel A: rel_alt time series ===
    ax = axes[0]
    if alt:
        t_a = [to_rel(t) for t, _ in alt]
        v_a = [v for _, v in alt]
        ax.plot(t_a, v_a, color="#1e40af", lw=1.0, label="rel_alt")
    threshold = filter_params["min_mission_depth_m"]
    ax.axhline(threshold, color="#dc2626", ls="--", lw=1.0,
               label=f"threshold {threshold}m")
    ax.axhline(0, color="#94a3b8", ls=":", lw=0.8, label="surface")
    # Mark run window
    ax.axvspan(0, dur, alpha=0.10, color="#16a34a" if status == "OK" else "#dc2626")
    ax.set_ylabel("rel_alt (m)")
    ax.set_title(
        f"(a) Depth — max={run['max_depth_m']}m  mean={run['mean_depth_m']}m  "
        f"sustained={run['sustained_duration_s']:.0f}s below {threshold}m  "
        f"pct_near_bottom={run['pct_near_bottom']:.1f}%"
    )
    ax.legend(loc="lower right", fontsize=8)
    ax.grid(True, alpha=0.3)
    ax.set_xlim(-30, dur + 30)

    # === Panel B: state armed + mode ===
    ax = axes[1]
    if state:
        t_s = [to_rel(t) for t, _, _ in state]
        armed_y = [1 if a else 0 for _, a, _ in state]
        ax.step(t_s, armed_y, where="post", color="#16a34a", lw=1.2, label="armed")
        # Annotate dominant mode + mode transitions
        prev_mode = None
        for t, _, m in state:
            if m != prev_mode:
                ax.axvline(to_rel(t), color="#7c3aed", lw=0.5, alpha=0.4)
                ax.text(to_rel(t), 0.5, m, rotation=90, fontsize=7,
                        color="#7c3aed", va="center", alpha=0.7)
                prev_mode = m
    ax.set_ylim(-0.2, 1.3)
    ax.set_yticks([0, 1])
    ax.set_yticklabels(["disarmed", "armed"])
    ax.set_title(f"(b) MAVROS state — dominant mode: {run.get('dominant_mode','?')}")
    ax.grid(True, alpha=0.3)
    ax.set_xlim(-30, dur + 30)

    # === Panel C: depth distribution + cumulative time below threshold ===
    ax = axes[2]
    # filter alt to run window
    run_vals = [v for t, v in alt if start_e <= t <= end_e]
    if run_vals:
        ax.hist(run_vals, bins=40, color="#3b82f6", alpha=0.7, edgecolor="white")
        ax.axvline(threshold, color="#dc2626", ls="--", lw=1.0,
                   label=f"threshold {threshold}m")
        n_total = len(run_vals)
        n_below = sum(1 for v in run_vals if v < threshold)
        ax.set_title(
            f"(c) Depth distribution within run  |  {n_below}/{n_total} samples below threshold "
            f"= {100*n_below/n_total:.1f}%"
        )
        ax.legend(loc="upper right", fontsize=8)
    ax.set_xlabel("rel_alt (m)")
    ax.set_ylabel("samples")
    ax.grid(True, alpha=0.3)

    # === Panel D: verdict box ===
    ax = axes[3]
    ax.axis("off")
    color = "#16a34a" if status == "OK" else "#dc2626"
    verdict_text = f"VERDICT: {status}"
    ax.text(0.02, 0.85, verdict_text, fontsize=18, fontweight="bold",
            color=color, transform=ax.transAxes)

    # Build criteria summary
    crit_lines = []
    min_dur = filter_params["min_sustained_duration_s"]
    min_pct = filter_params["min_near_bottom_pct"]
    min_depth = filter_params["min_mission_depth_m"]

    pass_dur = run["sustained_duration_s"] >= min_dur
    pass_pct = run["pct_near_bottom"] >= min_pct
    pass_depth = run["max_depth_m"] < min_depth

    crit_lines.append(
        f"{'OK' if pass_depth else 'KO':3s}  max_depth {run['max_depth_m']}m < {min_depth}m"
    )
    crit_lines.append(
        f"{'OK' if pass_dur else 'KO':3s}  sustained_duration {run['sustained_duration_s']:.0f}s >= {min_dur:.0f}s"
    )
    crit_lines.append(
        f"{'OK' if pass_pct else 'KO':3s}  pct_near_bottom {run['pct_near_bottom']:.1f}% >= {min_pct:.0f}%"
    )
    crit_lines.append(
        f"     duration {run['duration_s']:.0f}s | mode {run.get('dominant_mode','?')} | method {run.get('detection_method','?')}"
    )

    for i, line in enumerate(crit_lines):
        ax.text(0.02, 0.65 - i * 0.13, line, fontsize=10, family="monospace",
                transform=ax.transAxes)

    if reject_reason:
        ax.text(0.02, 0.05, f"REASON: {reject_reason}", fontsize=9,
                color="#dc2626", transform=ax.transAxes, family="monospace")

    axes[-2].set_xlabel("time since run start (s)")
    plt.tight_layout()
    plt.savefig(out_path, dpi=90, bbox_inches="tight")
    plt.close(fig)


HTML_TEMPLATE = """<!doctype html>
<html lang="fr">
<head>
<meta charset="utf-8">
<title>02 Runs Diagnostic — {mission}</title>
<style>
body {{ font-family: -apple-system, system-ui, sans-serif; background: #0f172a; color: #e2e8f0; margin: 0; padding: 20px; }}
h1 {{ color: #38bdf8; margin: 0 0 4px 0; }}
.subtitle {{ color: #94a3b8; margin-bottom: 20px; font-size: 14px; }}
.params {{ background: #1e293b; padding: 10px 14px; border-radius: 6px; margin-bottom: 24px; font-family: monospace; font-size: 12px; color: #cbd5e1; }}
.grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(560px, 1fr)); gap: 20px; }}
.card {{ background: #1e293b; border-radius: 8px; padding: 12px; border-left: 6px solid #475569; }}
.card.ok {{ border-left-color: #16a34a; }}
.card.ko {{ border-left-color: #dc2626; }}
.card h3 {{ margin: 0 0 6px 0; font-size: 16px; }}
.card h3.ok {{ color: #4ade80; }}
.card h3.ko {{ color: #f87171; }}
.card img {{ width: 100%; height: auto; border-radius: 4px; }}
.meta {{ color: #94a3b8; font-size: 12px; margin: 4px 0 8px 0; font-family: monospace; }}
.reason {{ color: #fca5a5; font-size: 12px; margin-top: 6px; font-family: monospace; }}
.summary {{ display: flex; gap: 16px; margin-bottom: 20px; }}
.pill {{ background: #1e293b; padding: 8px 14px; border-radius: 999px; font-size: 13px; }}
.pill.ok {{ color: #4ade80; }}
.pill.ko {{ color: #f87171; }}
</style>
</head>
<body>
<h1>02 Runs Diagnostic — {mission}</h1>
<div class="subtitle">Generated {generated_at} · filter v5 strict</div>
<div class="summary">
  <div class="pill ok">OK: {n_ok}</div>
  <div class="pill ko">Rejected: {n_ko}</div>
  <div class="pill">Total candidates: {n_total}</div>
</div>
<div class="params">{params}</div>
<div class="grid">
{cards}
</div>
</body>
</html>
"""

CARD_TEMPLATE = """<div class="card {cls}">
  <h3 class="{cls}">{run_id} — {status}</h3>
  <div class="meta">duration={duration}s | max_depth={max_depth}m | sustained={sustained}s | pct_near={pct}% | mode={mode}</div>
  <img src="{png}" alt="{run_id}">
  {reason_block}
</div>"""


def build_html(mission, filter_params, generated_at, ok_runs, ko_runs, out_dir):
    cards = []
    for run, status, reason in [(r, "OK", None) for r in ok_runs] + [(r, "REJECTED", r.get("rejected_reason")) for r in ko_runs]:
        cls = "ok" if status == "OK" else "ko"
        reason_block = f'<div class="reason">REASON: {reason}</div>' if reason else ""
        cards.append(CARD_TEMPLATE.format(
            cls=cls, run_id=run["run_id"], status=status,
            duration=f"{run['duration_s']:.0f}",
            max_depth=run["max_depth_m"],
            sustained=f"{run['sustained_duration_s']:.0f}",
            pct=f"{run['pct_near_bottom']:.1f}",
            mode=run.get("dominant_mode", "?"),
            png=f"{run['run_id']}.png",
            reason_block=reason_block,
        ))
    html = HTML_TEMPLATE.format(
        mission=mission,
        generated_at=generated_at,
        params=json.dumps(filter_params, indent=2),
        n_ok=len(ok_runs),
        n_ko=len(ko_runs),
        n_total=len(ok_runs) + len(ko_runs),
        cards="\n".join(cards),
    )
    (out_dir / "index.html").write_text(html, encoding="utf-8")


def main():
    parser = argparse.ArgumentParser(description="Stage 02b — runs diagnostic plots")
    parser.add_argument("--mission", required=True)
    parser.add_argument("--ssd-base", default="/mnt/ssd")
    parser.add_argument("--out", default="data/")
    parser.add_argument("--padding-s", type=float, default=120.0,
                        help="padding (s) avant/après run pour contexte")
    args = parser.parse_args()

    mission_dir = Path(args.out) / args.mission
    runs_json = mission_dir / "02_runs.json"
    if not runs_json.exists():
        print(f"[ERROR] {runs_json} missing — run 02_mission_run_detect first", file=sys.stderr)
        sys.exit(1)

    bags_root = Path(args.ssd_base) / args.mission / "raw_data" / "logs" / "SUB" / "bag"
    if not bags_root.exists():
        print(f"[ERROR] bags dir not found: {bags_root}", file=sys.stderr)
        sys.exit(1)

    data = json.loads(runs_json.read_text())
    filter_params = data["filter_params"]
    out_dir = mission_dir / "02_runs_diag"
    out_dir.mkdir(parents=True, exist_ok=True)

    ok_runs = data["all_runs_sorted"]
    ko_runs = data["all_runs_rejected"]
    print(f"[stage02b] {len(ok_runs)} OK + {len(ko_runs)} rejected = {len(ok_runs)+len(ko_runs)} runs to plot")

    # Group by AUV → 1 read per AUV bag set covering all runs
    runs_by_auv = {}
    for run in ok_runs + ko_runs:
        # run_id "AUV210_run_00" → AUV210
        phys = run["run_id"].split("_")[0]
        runs_by_auv.setdefault(phys, []).append(run)

    pad = args.padding_s
    for phys_id, runs in runs_by_auv.items():
        mcap_id = PHYS_TO_MCAP.get(phys_id, phys_id)
        t_start = min(r["start_epoch"] for r in runs) - pad
        t_end = max(r["end_epoch"] for r in runs) + pad
        print(f"  [{phys_id}] reading {mcap_id} bags [{t_start:.0f}..{t_end:.0f}] for {len(runs)} runs")
        alt, state, hdg = gather_signals_for_auv(bags_root, mcap_id, t_start, t_end)
        print(f"    {len(alt)} alt pts, {len(state)} state pts, {len(hdg)} hdg pts")

        for run in runs:
            r_start = run["start_epoch"] - pad
            r_end = run["end_epoch"] + pad
            alt_r = [(t, v) for t, v in alt if r_start <= t <= r_end]
            state_r = [(t, a, m) for t, a, m in state if r_start <= t <= r_end]
            hdg_r = [(t, v) for t, v in hdg if r_start <= t <= r_end]
            status = "REJECTED" if run in ko_runs else "OK"
            reason = run.get("rejected_reason") if status == "REJECTED" else None
            out_png = out_dir / f"{run['run_id']}.png"
            plot_run_diag(run, alt_r, state_r, hdg_r, filter_params, status, reason, out_png)
            print(f"    [{status}] {run['run_id']}.png written")

    # Build index.html + index.json
    generated_at = datetime.now(timezone.utc).isoformat()
    index = {
        "mission": args.mission,
        "generated_at": generated_at,
        "filter_params": filter_params,
        "ok_runs": [r["run_id"] for r in ok_runs],
        "rejected_runs": [r["run_id"] for r in ko_runs],
    }
    (out_dir / "index.json").write_text(json.dumps(index, indent=2))
    build_html(args.mission, filter_params, generated_at, ok_runs, ko_runs, out_dir)
    print(f"[stage02b] Done: {out_dir}/index.html")


if __name__ == "__main__":
    main()
