#!/usr/bin/env python3
"""Build the phase-balanced event-geometry figure for the GRL manuscript."""

from __future__ import annotations

import csv
import importlib.util
import json
import math
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path
from statistics import median

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


OVERLEAF_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = OVERLEAF_ROOT.parent
PHASE_BALANCED_ROOT = PROJECT_ROOT / "odata" / "phase_balanced_20190706_20211113"
PICK_EVAL_SCRIPT = PROJECT_ROOT / "odata" / "evaluate_pick_recall_no_nms.py"
LABEL_JSON = PROJECT_ROOT / "data" / "datax" / "data" / "label" / "annotations_for_continuous_hdf5.json"
WAVEFORM_DB = PROJECT_ROOT / "data" / "datax" / "data" / "index" / "waveform_index.sqlite"
FIG_DIR = OVERLEAF_ROOT / "figures"
FIG_PDF = FIG_DIR / "fig_event_geometry_distribution_polished.pdf"
FIG_PNG = FIG_DIR / "fig_event_geometry_distribution_polished.png"
DATA_CSV = FIG_DIR / "fig_event_geometry_distribution_polished_data.csv"
SUMMARY_JSON = FIG_DIR / "fig_event_geometry_distribution_polished_summary.json"

DAYS = {"20190706", "20211113"}
TP_TOL_S = 1.5
SEARCH_WINDOW_S = 5.0

CONDITIONS = [
    {
        "slug": "moderate",
        "label": "Moderate phase-balanced",
        "threshold": "P>=5/S>=3.73 dB",
        "plot_threshold": r"P$\geq$5/S$\geq$3.73 dB",
        "snr": "p5_sbal",
        "confidence": "conf_p5_sbal",
    },
    {
        "slug": "strict",
        "label": "Strict phase-balanced",
        "threshold": "P>=10/S>=6.01 dB",
        "plot_threshold": r"P$\geq$10/S$\geq$6.01 dB",
        "snr": "p10_sbal",
        "confidence": "conf_p10_sbal",
    },
]

COLORS = {
    "snr_only": "#F58518",
    "confidence_only": "#4C78A8",
    "snr": "#F58518",
    "confidence": "#4C78A8",
}


def load_pick_eval_module():
    spec = importlib.util.spec_from_file_location("pick_eval_no_nms", PICK_EVAL_SCRIPT)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot load {PICK_EVAL_SCRIPT}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def iter_jsonl(path: Path):
    with path.open("r", encoding="utf-8", errors="replace") as handle:
        for line in handle:
            if line.strip():
                yield json.loads(line)


def load_event_summary(condition: str) -> dict:
    path = PHASE_BALANCED_ROOT / "eval_events" / condition / "event_match_summary.json"
    return json.loads(path.read_text(encoding="utf-8"))


def load_true_positive_event_ids(condition: str) -> set[str]:
    path = PHASE_BALANCED_ROOT / "eval_events" / condition / "event_matches.jsonl"
    ids: set[str] = set()
    for row in iter_jsonl(path):
        if row.get("record_type") == "event_true_positive":
            ids.add(str(row["ref_event_id"]))
    return ids


def load_labels(eval_mod) -> list[dict]:
    coverage = eval_mod.Coverage(WAVEFORM_DB)
    labels = []
    for label in eval_mod.iter_label_picks(LABEL_JSON, None, None, coverage):
        day = datetime.fromtimestamp(float(label["time"]), tz=timezone.utc).strftime("%Y%m%d")
        if day in DAYS:
            labels.append(label)
    return labels


def load_auto_pick_index(eval_mod, condition: str):
    indexed = defaultdict(list)
    starts: dict[tuple[str, str], list[float]] = {}
    auto_counts = Counter()
    pick_dir = PHASE_BALANCED_ROOT / "hourly" / condition
    files = sorted(pick_dir.glob("*.phase.jsonl"))
    if not files:
        raise FileNotFoundError(f"No phase files found in {pick_dir}")
    for path in files:
        sub_indexed, _sub_starts, sub_counts = eval_mod.load_auto_picks(path, None, None)
        for key, values in sub_indexed.items():
            indexed[key].extend(values)
        auto_counts.update(sub_counts)
    for key, values in indexed.items():
        values.sort(key=lambda item: item["time"])
        starts[key] = [item["time"] for item in values]
    return indexed, starts, auto_counts


def event_support_counts(eval_mod, labels: list[dict], condition: str) -> tuple[dict[str, dict], Counter]:
    indexed, starts, auto_counts = load_auto_pick_index(eval_mod, condition)
    counts = defaultdict(lambda: {"P": 0, "S": 0, "stations": set(), "S_stations": set()})
    for label in labels:
        hit = eval_mod.nearest(
            indexed,
            starts,
            label["station_id"],
            label["phase"],
            label["time"],
            SEARCH_WINDOW_S,
        )
        if hit is None or abs(float(hit["residual_s"])) > TP_TOL_S:
            continue
        event_id = str(label["event_id"])
        phase = str(label["phase"])
        station_id = str(label["station_id"])
        counts[event_id][phase] += 1
        counts[event_id]["stations"].add(station_id)
        if phase == "S":
            counts[event_id]["S_stations"].add(station_id)

    packed = {}
    for event_id, row in counts.items():
        packed[event_id] = {
            "P": int(row["P"]),
            "S": int(row["S"]),
            "stations": len(row["stations"]),
            "S_stations": len(row["S_stations"]),
        }
    return packed, auto_counts


def safe_metric(counts: dict[str, dict], event_id: str, metric: str) -> int:
    return int(counts.get(event_id, {}).get(metric, 0))


def event_rows_for_condition(condition: dict, support: dict[str, dict[str, dict]]) -> tuple[list[dict], dict]:
    snr = condition["snr"]
    confidence = condition["confidence"]
    snr_tp = load_true_positive_event_ids(snr)
    conf_tp = load_true_positive_event_ids(confidence)
    snr_summary = load_event_summary(snr)
    conf_summary = load_event_summary(confidence)

    rows = []
    class_defs = [
        ("snr_only", "SNR-only", sorted(snr_tp - conf_tp), snr, confidence),
        ("confidence_only", "Confidence-only", sorted(conf_tp - snr_tp), confidence, snr),
    ]
    for class_slug, class_label, event_ids, recovered_cond, missed_cond in class_defs:
        for event_id in event_ids:
            for metric in ("P", "S", "stations", "S_stations"):
                recovered = safe_metric(support[recovered_cond], event_id, metric)
                missed = safe_metric(support[missed_cond], event_id, metric)
                rows.append(
                    {
                        "condition_slug": condition["slug"],
                        "condition_label": condition["label"],
                        "threshold": condition["threshold"],
                        "class_slug": class_slug,
                        "class_label": class_label,
                        "event_id": event_id,
                        "metric": metric,
                        "recovered_support": recovered,
                        "missed_support": missed,
                        "difference": recovered - missed,
                    }
                )

    summary = {
        "condition_slug": condition["slug"],
        "condition_label": condition["label"],
        "threshold": condition["threshold"],
        "plot_threshold": condition["plot_threshold"],
        "snr_condition": snr,
        "confidence_condition": confidence,
        "snr_event_metrics": snr_summary["metrics"],
        "confidence_event_metrics": conf_summary["metrics"],
        "snr_event_counts": snr_summary["counts"],
        "confidence_event_counts": conf_summary["counts"],
        "both_tp": len(snr_tp & conf_tp),
        "snr_only_tp": len(snr_tp - conf_tp),
        "confidence_only_tp": len(conf_tp - snr_tp),
    }
    return rows, summary


def describe_differences(rows: list[dict]) -> dict:
    out = {}
    for condition in CONDITIONS:
        cond_rows = [r for r in rows if r["condition_slug"] == condition["slug"]]
        out[condition["slug"]] = {}
        for class_slug in ("snr_only", "confidence_only"):
            out[condition["slug"]][class_slug] = {}
            for metric in ("P", "S", "stations", "S_stations"):
                values = [int(r["difference"]) for r in cond_rows if r["class_slug"] == class_slug and r["metric"] == metric]
                if not values:
                    continue
                out[condition["slug"]][class_slug][metric] = {
                    "n_events": len(values),
                    "median_recovered_minus_missed": median(values),
                    "positive_fraction": sum(v > 0 for v in values) / len(values),
                    "zero_fraction": sum(v == 0 for v in values) / len(values),
                    "negative_fraction": sum(v < 0 for v in values) / len(values),
                    "min": min(values),
                    "max": max(values),
                }
    return out


def write_outputs(rows: list[dict], summary: dict) -> None:
    fieldnames = [
        "condition_slug",
        "condition_label",
        "threshold",
        "class_slug",
        "class_label",
        "event_id",
        "metric",
        "recovered_support",
        "missed_support",
        "difference",
    ]
    with DATA_CSV.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)
    SUMMARY_JSON.write_text(json.dumps(summary, indent=2), encoding="utf-8")


def style_axis(ax):
    ax.spines[["top", "right"]].set_visible(False)
    ax.grid(axis="y", color="#e8e8e8", linewidth=0.8)
    ax.set_axisbelow(True)
    ax.tick_params(labelsize=7)


def panel_label(ax, label: str):
    ax.text(
        0.02,
        0.96,
        label,
        transform=ax.transAxes,
        ha="left",
        va="top",
        fontsize=9,
        fontweight="bold",
        bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.85, "pad": 1.0},
    )


def plot_metric_bars(ax, condition_summary: dict, panel: str):
    metrics = ["precision", "recall"]
    x = np.arange(len(metrics))
    width = 0.35
    snr_values = [condition_summary["snr_event_metrics"][m] for m in metrics]
    conf_values = [condition_summary["confidence_event_metrics"][m] for m in metrics]
    ax.bar(x - width / 2, snr_values, width, label="Phase-balanced SNR", color=COLORS["snr"], edgecolor="black", linewidth=0.3)
    ax.bar(x + width / 2, conf_values, width, label="Top phase probability", color=COLORS["confidence"], edgecolor="black", linewidth=0.3)
    ax.set_xticks(x, ["Precision", "Recall"])
    ax.set_ylim(0, 0.9)
    ax.set_ylabel("Event metric", fontsize=8)
    ax.set_title(condition_summary["plot_threshold"], fontsize=8)
    style_axis(ax)
    panel_label(ax, panel)
    for xpos, value in zip(x - width / 2, snr_values):
        ax.text(xpos, value + 0.025, f"{value:.3f}", ha="center", va="bottom", fontsize=6, rotation=90)
    for xpos, value in zip(x + width / 2, conf_values):
        ax.text(xpos, value + 0.025, f"{value:.3f}", ha="center", va="bottom", fontsize=6, rotation=90)


def plot_difference_hist(
    ax,
    rows: list[dict],
    condition_slug: str,
    metric: str,
    panel: str,
    title: str,
    show_legend: bool = False,
):
    cond_rows = [r for r in rows if r["condition_slug"] == condition_slug and r["metric"] == metric]
    values_by_class = {
        class_slug: [int(r["difference"]) for r in cond_rows if r["class_slug"] == class_slug]
        for class_slug in ("snr_only", "confidence_only")
    }
    all_values = [v for values in values_by_class.values() for v in values]
    xmin = min(-8, min(all_values) if all_values else -1)
    xmax = max(16, max(all_values) if all_values else 1)
    xs = np.arange(xmin, xmax + 1)
    width = 0.38
    offsets = {"snr_only": -width / 2, "confidence_only": width / 2}
    labels = {"snr_only": "SNR-only events", "confidence_only": "Confidence-only events"}
    for class_slug in ("snr_only", "confidence_only"):
        values = values_by_class[class_slug]
        counts = Counter(values)
        denom = len(values) or 1
        heights = np.array([counts.get(int(x), 0) / denom for x in xs], dtype=float)
        ax.bar(
            xs + offsets[class_slug],
            heights,
            width,
            label=labels[class_slug],
            color=COLORS[class_slug],
            alpha=0.86,
            edgecolor="black",
            linewidth=0.2,
        )
    ax.axvline(0, color="#333333", linewidth=0.8)
    ax.set_xlim(xmin - 0.8, xmax + 0.8)
    ax.set_ylim(0, None)
    ax.set_xlabel("Recovered stream minus missed stream", fontsize=8)
    ylabel = "Fraction of discordant events"
    ax.set_ylabel(ylabel, fontsize=8)
    ax.set_title(title, fontsize=8)
    style_axis(ax)
    panel_label(ax, panel)
    text_lines = []
    for class_slug in ("snr_only", "confidence_only"):
        values = values_by_class[class_slug]
        if not values:
            continue
        med = median(values)
        pos = 100.0 * sum(v > 0 for v in values) / len(values)
        text_lines.append(f"{labels[class_slug].replace(' events', '')}: n={len(values)}, med={med:g}, >0={pos:.0f}%")
    ax.text(
        0.98,
        0.95,
        "\n".join(text_lines),
        transform=ax.transAxes,
        ha="right",
        va="top",
        fontsize=6.2,
        bbox={"facecolor": "white", "edgecolor": "#bbbbbb", "alpha": 0.88, "pad": 2.0},
    )
    if show_legend:
        ax.legend(frameon=False, fontsize=7, loc="upper left", bbox_to_anchor=(0.02, 0.82))


def make_figure(rows: list[dict], summaries: dict) -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8,
            "axes.labelsize": 8,
            "xtick.labelsize": 7,
            "ytick.labelsize": 7,
            "legend.fontsize": 7,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )
    fig, axes = plt.subplots(
        2,
        2,
        figsize=(7.2, 4.95),
        dpi=300,
        constrained_layout=True,
    )
    moderate = summaries["moderate"]
    strict = summaries["strict"]
    moderate_pr = (
        f"{moderate['plot_threshold']}, matched S arrivals\n"
        f"P/R: SNR {moderate['snr_event_metrics']['precision']:.3f}/{moderate['snr_event_metrics']['recall']:.3f}; "
        f"prob. {moderate['confidence_event_metrics']['precision']:.3f}/{moderate['confidence_event_metrics']['recall']:.3f}"
    )
    strict_pr = (
        f"{strict['plot_threshold']}, matched S arrivals\n"
        f"P/R: SNR {strict['snr_event_metrics']['precision']:.3f}/{strict['snr_event_metrics']['recall']:.3f}; "
        f"prob. {strict['confidence_event_metrics']['precision']:.3f}/{strict['confidence_event_metrics']['recall']:.3f}"
    )
    plot_difference_hist(axes[0, 0], rows, "moderate", "S", "A", moderate_pr, show_legend=True)
    plot_difference_hist(axes[0, 1], rows, "strict", "S", "B", strict_pr, show_legend=True)
    plot_difference_hist(axes[1, 0], rows, "moderate", "S_stations", "C", "Moderate budget, S-contributing stations")
    plot_difference_hist(axes[1, 1], rows, "strict", "S_stations", "D", "Strict budget, S-contributing stations")
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIG_PDF, bbox_inches="tight")
    fig.savefig(FIG_PNG, dpi=300, bbox_inches="tight", facecolor="white")
    plt.close(fig)


def main() -> None:
    eval_mod = load_pick_eval_module()
    labels = load_labels(eval_mod)
    needed_conditions = sorted({cond["snr"] for cond in CONDITIONS} | {cond["confidence"] for cond in CONDITIONS})
    support = {}
    auto_counts = {}
    for condition in needed_conditions:
        support[condition], auto_counts[condition] = event_support_counts(eval_mod, labels, condition)

    rows = []
    summaries = {}
    for condition in CONDITIONS:
        condition_rows, condition_summary = event_rows_for_condition(condition, support)
        rows.extend(condition_rows)
        summaries[condition["slug"]] = condition_summary

    full_summary = {
        "source_root": str(PHASE_BALANCED_ROOT),
        "days": sorted(DAYS),
        "tp_tolerance_s": TP_TOL_S,
        "search_window_s": SEARCH_WINDOW_S,
        "event_matching": "20 km epicentral distance / 3 s origin time",
        "auto_pick_counts_by_condition": {key: dict(value) for key, value in auto_counts.items()},
        "conditions": summaries,
        "difference_summary": describe_differences(rows),
    }
    write_outputs(rows, full_summary)
    make_figure(rows, summaries)
    print(json.dumps(full_summary, indent=2))


if __name__ == "__main__":
    main()