#!/usr/bin/env python3
"""Create a clean data-driven Figure 1 for SNR observability shifts."""

from __future__ import annotations

import csv
import json
import math
from pathlib import Path

import h5py
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D


OVERLEAF_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = OVERLEAF_ROOT.parent
OUTPUTS = PROJECT_ROOT / "outputs"
FIG_DIR = OVERLEAF_ROOT / "figures"

PHASE_RECORDS = OUTPUTS / "snr_transfer_phase_balanced_cache" / "records_train_all.json"
PHASE_SNR = OUTPUTS / "snr_transfer_phase_balanced_cache" / "train_phase_snr_db.json"
DISP_CACHE = OUTPUTS / "disp_snr_transfer_seed20260609" / "ncf_snr_cache.json"
DISP_H5 = PROJECT_ROOT / "data" / "ncf_data" / "ncf_disp_dataset_with_disp_image.h5"

OUT_PDF = FIG_DIR / "fig_observability_real_data_v1.pdf"
OUT_PNG = FIG_DIR / "fig_observability_real_data_v1.png"
OUT_CSV = FIG_DIR / "fig_observability_real_data_v1_data.csv"
OUT_JSON = FIG_DIR / "fig_observability_real_data_v1_summary.json"

RED = "#C44E52"
BLUE = "#2F6DB2"
DARK_RED = "#9D1C1C"
GRAY = "#7A7A7A"
DARK = "#222222"
LIGHT = "#E8E8E8"
FILL_GRAY = "#D5D5D5"
VIOLIN = "#9FC6DF"

PHASE_THRESHOLD = 5.0
STRICT_PHASE_THRESHOLD = 10.0
PERIOD_BINS = [(5.0, 10.0), (10.0, 15.0), (15.0, 25.0), (25.0, 40.0)]


def load_json(path: Path):
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def phase_group(name: str) -> str | None:
    phase = str(name or "")[:1].upper()
    return phase if phase in {"P", "S"} else None


def pick_snr_key(record: dict, idx: int, pick: dict) -> str:
    return f"{record['event']}/{record['station']}/{idx}:{pick['phase']}:{pick['index']}:{pick['source']}"


def finite_float(value, default=float("nan")) -> float:
    try:
        out = float(value)
    except Exception:
        return default
    return out if math.isfinite(out) else default


def read_phase_data() -> tuple[list[dict], list[dict]]:
    records = load_json(PHASE_RECORDS)["records"]
    snr_by_pick = {key: float(value) for key, value in load_json(PHASE_SNR).items()}
    pick_rows: list[dict] = []
    record_rows: list[dict] = []
    for record in records:
        distance = finite_float(record.get("distance_km"))
        if not math.isfinite(distance) or distance <= 0:
            continue
        record_snrs = []
        for idx, pick in enumerate(record.get("phases", [])):
            phase = phase_group(pick.get("phase", ""))
            if phase is None:
                continue
            snr = snr_by_pick.get(pick_snr_key(record, idx, pick), float("nan"))
            if not math.isfinite(snr):
                continue
            record_snrs.append(snr)
            pick_rows.append({"phase": phase, "distance": distance, "snr": snr})
        if record_snrs:
            record_rows.append({"distance": distance, "max_snr": max(record_snrs)})
    return pick_rows, record_rows


def sample_phase_points(points: list[dict], per_phase_bin: int = 30) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(20260627)
    bins = np.array([4, 8, 15, 30, 60, 120, 240, 480, 1000], dtype=float)
    sampled = []
    for phase in ["P", "S"]:
        phase_points = [p for p in points if p["phase"] == phase]
        distances = np.asarray([p["distance"] for p in phase_points], dtype=float)
        for lo, hi in zip(bins[:-1], bins[1:]):
            idx = np.where((distances >= lo) & (distances < hi))[0]
            if idx.size == 0:
                continue
            take = min(per_phase_bin, idx.size)
            chosen = rng.choice(idx, size=take, replace=False)
            sampled.extend(phase_points[i] for i in chosen)
    return {
        "distance": np.asarray([p["distance"] for p in sampled], dtype=float),
        "snr": np.asarray([p["snr"] for p in sampled], dtype=float),
        "phase": np.asarray([p["phase"] for p in sampled]),
    }


def read_dispersion_snr_by_period(max_per_bin: int = 6000) -> tuple[list[np.ndarray], list[dict]]:
    raw_values = [[] for _ in PERIOD_BINS]
    cache = load_json(DISP_CACHE)
    with h5py.File(DISP_H5, "r") as h5:
        paths = h5["paths"]
        for key, row in cache.items():
            if row.get("split") != "train":
                continue
            snr = finite_float(row.get("snr_db"))
            if not math.isfinite(snr) or key not in paths:
                continue
            group = paths[key]
            try:
                periods = np.asarray(group["disp_periods"][()], dtype=float)
                velocities = np.asarray(group["disp_velocity"][()], dtype=float)
                mask = np.asarray(group["disp_mask"][()], dtype=bool)
            except Exception:
                continue
            valid = mask & np.isfinite(periods) & np.isfinite(velocities) & (periods > 0)
            for period in periods[valid]:
                for i, (lo, hi) in enumerate(PERIOD_BINS):
                    in_bin = (lo <= period < hi) or (i == len(PERIOD_BINS) - 1 and period <= hi)
                    if in_bin:
                        raw_values[i].append(snr)
                        break

    rng = np.random.default_rng(20260627)
    plot_values = []
    summaries = []
    for (lo, hi), values in zip(PERIOD_BINS, raw_values):
        arr = np.asarray(values, dtype=float)
        arr = arr[np.isfinite(arr)]
        if arr.size > max_per_bin:
            plot_arr = rng.choice(arr, size=max_per_bin, replace=False)
        else:
            plot_arr = arr
        plot_values.append(plot_arr)
        summaries.append(
            {
                "period_left": lo,
                "period_right": hi,
                "n": int(arr.size),
                "q05": float(np.percentile(arr, 5)),
                "q25": float(np.percentile(arr, 25)),
                "median": float(np.percentile(arr, 50)),
                "q75": float(np.percentile(arr, 75)),
                "q95": float(np.percentile(arr, 95)),
                "fraction_ge_5db": float(np.mean(arr >= PHASE_THRESHOLD)),
            }
        )
    return plot_values, summaries


def style_axis(ax: plt.Axes, grid: bool = False) -> None:
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    for side in ["left", "bottom"]:
        ax.spines[side].set_color("#8A8A8A")
        ax.spines[side].set_linewidth(0.65)
    ax.tick_params(labelsize=6.5, width=0.6, length=2.6, color="#666666")
    if grid:
        ax.grid(color=LIGHT, lw=0.45, alpha=0.9)
        ax.set_axisbelow(True)


def add_header(ax: plt.Axes, label: str, title: str) -> None:
    ax.axis("off")
    ax.text(0.00, 0.66, label, fontsize=12.8, fontweight="bold", color=DARK, ha="left", va="center")
    ax.text(0.18, 0.68, title, fontsize=8.2, fontweight="bold", color=DARK, ha="left", va="center")


def set_log_distance_axis(ax: plt.Axes) -> None:
    ax.set_xscale("log")
    ax.set_xlim(4, 1000)
    ax.set_xticks([10, 100, 1000])
    ax.get_xaxis().set_major_formatter(mpl.ticker.ScalarFormatter())
    ax.xaxis.set_minor_locator(mpl.ticker.NullLocator())


def draw_phase_scatter(ax: plt.Axes, points: dict[str, np.ndarray], threshold: bool = False) -> None:
    distance = points["distance"]
    snr = points["snr"]
    phase = points["phase"]
    kept = snr >= PHASE_THRESHOLD

    if threshold:
        for phase_name, color in [("P", RED), ("S", BLUE)]:
            removed_mask = (phase == phase_name) & ~kept
            ax.scatter(
                distance[removed_mask],
                np.clip(snr[removed_mask], -2, 22),
                s=13,
                facecolors="none",
                edgecolors=color,
                alpha=0.42,
                linewidth=0.55,
                rasterized=True,
            )
        for phase_name, color in [("P", RED), ("S", BLUE)]:
            kept_mask = (phase == phase_name) & kept
            ax.scatter(
                distance[kept_mask],
                np.clip(snr[kept_mask], -2, 22),
                s=13,
                color=color,
                alpha=0.86,
                linewidth=0,
                rasterized=True,
            )
        ax.axhline(PHASE_THRESHOLD, color="#B00000", linestyle=(0, (4, 2)), linewidth=1.0)
        ax.text(4.7, PHASE_THRESHOLD + 0.75, "SNR = 5 dB", color="#B00000", fontsize=7.1, fontweight="bold")
        ax.text(520, PHASE_THRESHOLD + 2.2, "Kept", color="#B00000", fontsize=7.4, fontweight="bold", ha="center")
        ax.text(520, PHASE_THRESHOLD - 2.0, "Removed", color=GRAY, fontsize=7.0, ha="center")
        ax.set_yticks([0, PHASE_THRESHOLD, 15])
        ax.set_yticklabels(["Low", "5", "High"])
    else:
        for phase_name, color, label in [("P", RED, "P"), ("S", BLUE, "S")]:
            mask = phase == phase_name
            ax.scatter(
                distance[mask],
                np.clip(snr[mask], -2, 22),
                s=10,
                color=color,
                alpha=0.72,
                linewidth=0,
                rasterized=True,
                label=label,
            )
        ax.set_yticks([0, 10, 20])

    set_log_distance_axis(ax)
    ax.set_ylim(-2.5, 22)
    ax.set_xlabel("Epicentral distance (km)", fontsize=7.0)
    ax.set_ylabel("Pick SNR (dB)", fontsize=7.0)
    style_axis(ax, grid=True)


def draw_dispersion_violin(ax: plt.Axes, values: list[np.ndarray]) -> None:
    clipped = [np.clip(arr, -8, 16) for arr in values]
    parts = ax.violinplot(clipped, positions=np.arange(1, len(values) + 1), widths=0.72, showextrema=False)
    for body in parts["bodies"]:
        body.set_facecolor(VIOLIN)
        body.set_edgecolor("none")
        body.set_alpha(0.55)
    for x, arr in enumerate(values, start=1):
        q05, q25, q50, q75, q95 = np.percentile(arr, [5, 25, 50, 75, 95])
        q05, q25, q50, q75, q95 = np.clip([q05, q25, q50, q75, q95], -8, 16)
        ax.plot([x, x], [q05, q95], color="#1F77B4", lw=1.15)
        ax.plot([x - 0.10, x + 0.10], [q05, q05], color="#1F77B4", lw=1.15)
        ax.plot([x - 0.10, x + 0.10], [q95, q95], color="#1F77B4", lw=1.15)
        ax.plot([x - 0.17, x + 0.17], [q25, q25], color="#1F77B4", lw=0.75, alpha=0.75)
        ax.plot([x - 0.17, x + 0.17], [q75, q75], color="#1F77B4", lw=0.75, alpha=0.75)
        ax.scatter([x], [q50], s=9, color="#1F77B4", zorder=3)
    labels = [f"{int(lo)}-{int(hi)}" for lo, hi in PERIOD_BINS]
    ax.set_xticks(np.arange(1, len(values) + 1))
    ax.set_xticklabels(labels)
    ax.set_ylim(-8, 16)
    ax.set_ylabel("Dispersion SNR (dB)", fontsize=7.0)
    ax.set_xlabel("Period bin (s)", fontsize=7.0)
    style_axis(ax, grid=True)


def distance_histogram(distances: np.ndarray, mask: np.ndarray, bins: np.ndarray) -> np.ndarray:
    counts, _ = np.histogram(distances[mask], bins=bins)
    total = counts.sum()
    return counts / total if total else counts.astype(float)


def draw_distance_statistics(ax_dist: plt.Axes, ax_ret: plt.Axes, records: list[dict]) -> tuple[dict, list[dict]]:
    distances = np.asarray([row["distance"] for row in records], dtype=float)
    snr = np.asarray([row["max_snr"] for row in records], dtype=float)
    bins = np.geomspace(4, 1000, 22)
    centers = np.sqrt(bins[:-1] * bins[1:])
    full_mask = np.ones_like(snr, dtype=bool)
    med_mask = snr >= PHASE_THRESHOLD
    strict_mask = snr >= STRICT_PHASE_THRESHOLD

    full_fraction = distance_histogram(distances, full_mask, bins)
    med_fraction = distance_histogram(distances, med_mask, bins)
    strict_fraction = distance_histogram(distances, strict_mask, bins)
    ax_dist.stairs(full_fraction, bins, fill=True, color=FILL_GRAY, alpha=0.65, label="Unfiltered")
    ax_dist.stairs(med_fraction, bins, color=RED, lw=1.15, label="SNR >= 5 dB")
    ax_dist.stairs(strict_fraction, bins, color=DARK_RED, lw=1.15, label="SNR >= 10 dB")
    ax_dist.set_xscale("log")
    ax_dist.set_xlim(4, 1000)
    ax_dist.set_xticks([10, 100, 1000])
    ax_dist.get_xaxis().set_major_formatter(mpl.ticker.ScalarFormatter())
    ax_dist.xaxis.set_minor_locator(mpl.ticker.NullLocator())
    ax_dist.set_ylabel("Fraction of\nrecords", fontsize=7.0)
    ax_dist.legend(frameon=False, fontsize=5.9, loc="upper right", handlelength=1.3)
    style_axis(ax_dist, grid=False)

    full_counts, _ = np.histogram(distances, bins=bins)
    retained_rows = []
    for mask, color, label in [(med_mask, RED, "SNR >= 5 dB"), (strict_mask, DARK_RED, "SNR >= 10 dB")]:
        kept_counts, _ = np.histogram(distances[mask], bins=bins)
        with np.errstate(divide="ignore", invalid="ignore"):
            retained = np.where(full_counts > 0, kept_counts / full_counts, np.nan)
        ax_ret.plot(centers, retained, color=color, lw=1.25, marker="o", ms=2.2, label=label)
        for lo, hi, frac, full_n, kept_n in zip(bins[:-1], bins[1:], retained, full_counts, kept_counts):
            retained_rows.append(
                {
                    "condition": label,
                    "bin_left": float(lo),
                    "bin_right": float(hi),
                    "retained_fraction": float(frac) if math.isfinite(frac) else float("nan"),
                    "full_count": int(full_n),
                    "kept_count": int(kept_n),
                }
            )
    ax_ret.set_xscale("log")
    ax_ret.set_xlim(4, 1000)
    ax_ret.set_ylim(0, 1.03)
    ax_ret.set_xticks([10, 100, 1000])
    ax_ret.get_xaxis().set_major_formatter(mpl.ticker.ScalarFormatter())
    ax_ret.xaxis.set_minor_locator(mpl.ticker.NullLocator())
    ax_ret.set_yticks([0, 0.5, 1.0])
    ax_ret.set_xlabel("Epicentral distance (km)", fontsize=7.0)
    ax_ret.set_ylabel("Retained\nfraction", fontsize=7.0)
    style_axis(ax_ret, grid=False)

    summary = {
        "record_count": int(distances.size),
        "retained_record_fraction_5db": float(np.mean(med_mask)),
        "retained_record_fraction_10db": float(np.mean(strict_mask)),
        "median_distance_unfiltered_km": float(np.median(distances)),
        "median_distance_5db_km": float(np.median(distances[med_mask])),
        "median_distance_10db_km": float(np.median(distances[strict_mask])),
    }
    hist_rows = []
    for condition, fraction in [
        ("unfiltered", full_fraction),
        ("snr_ge_5db", med_fraction),
        ("snr_ge_10db", strict_fraction),
    ]:
        for lo, hi, frac in zip(bins[:-1], bins[1:], fraction):
            hist_rows.append({"condition": condition, "bin_left": float(lo), "bin_right": float(hi), "fraction": float(frac)})
    return summary, hist_rows + retained_rows


def write_exports(
    phase_points: dict[str, np.ndarray],
    all_pick_count: int,
    disp_summaries: list[dict],
    distance_summary: dict,
    distance_rows: list[dict],
) -> None:
    rows = []
    for distance, snr, phase in zip(phase_points["distance"], phase_points["snr"], phase_points["phase"]):
        rows.append(
            {
                "panel": "phase_scatter",
                "condition": "sampled",
                "metric": f"{phase}_pick_snr_db",
                "x": float(distance),
                "y": float(snr),
                "value": "",
            }
        )
    for summary in disp_summaries:
        label = f"{int(summary['period_left'])}-{int(summary['period_right'])}s"
        for key in ["n", "q05", "q25", "median", "q75", "q95", "fraction_ge_5db"]:
            rows.append(
                {
                    "panel": "dispersion_period_snr",
                    "condition": label,
                    "metric": key,
                    "x": "",
                    "y": "",
                    "value": summary[key],
                }
            )
    for row in distance_rows:
        rows.append(
            {
                "panel": "distance_coverage",
                "condition": row.get("condition", ""),
                "metric": "fraction" if "fraction" in row else "retained_fraction",
                "x": row["bin_left"],
                "y": row["bin_right"],
                "value": row.get("fraction", row.get("retained_fraction")),
            }
        )
    with OUT_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["panel", "condition", "metric", "x", "y", "value"], lineterminator="\n")
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "phase_pick_count": int(all_pick_count),
        "phase_points_sampled": int(len(phase_points["distance"])),
        "phase_thresholds_db": [PHASE_THRESHOLD, STRICT_PHASE_THRESHOLD],
        "dispersion_period_bins": disp_summaries,
        "distance_summary": distance_summary,
    }
    with OUT_JSON.open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)


def make_figure() -> None:
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    pick_rows, record_rows = read_phase_data()
    phase_points = sample_phase_points(pick_rows)
    disp_values, disp_summaries = read_dispersion_snr_by_period()

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 7,
            "axes.linewidth": 0.65,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )

    fig = plt.figure(figsize=(7.25, 4.08), dpi=300)
    outer = fig.add_gridspec(
        1,
        3,
        width_ratios=[1.04, 1.22, 1.04],
        left=0.055,
        right=0.985,
        top=0.875,
        bottom=0.12,
        wspace=0.28,
    )
    panel_a = outer[0, 0].subgridspec(3, 1, height_ratios=[0.18, 1.0, 0.88], hspace=0.72)
    panel_b = outer[0, 1].subgridspec(2, 1, height_ratios=[0.18, 1.0], hspace=0.20)
    panel_c = outer[0, 2].subgridspec(3, 1, height_ratios=[0.18, 0.82, 0.70], hspace=0.34)

    ax_a_head = fig.add_subplot(panel_a[0])
    ax_a_phase = fig.add_subplot(panel_a[1])
    ax_a_disp = fig.add_subplot(panel_a[2])
    add_header(ax_a_head, "A", "Unfiltered observations")
    draw_phase_scatter(ax_a_phase, phase_points, threshold=False)
    ax_a_phase.set_title("Phase picks", loc="left", fontsize=7.1, pad=1, fontweight="bold")
    ax_a_phase.legend(frameon=False, fontsize=6.0, loc="upper right", handletextpad=0.2, borderpad=0.1)
    draw_dispersion_violin(ax_a_disp, disp_values)
    ax_a_disp.set_title("Dispersion by period", loc="left", fontsize=7.1, pad=1, fontweight="bold")

    ax_b_head = fig.add_subplot(panel_b[0])
    ax_b = fig.add_subplot(panel_b[1])
    add_header(ax_b_head, "B", "Hard SNR cutoff")
    draw_phase_scatter(ax_b, phase_points, threshold=True)
    ax_b.set_title("Example phase-pick threshold", loc="left", fontsize=7.4, pad=2, fontweight="bold")
    legend_handles = [
        Line2D([0], [0], marker="o", color="none", markerfacecolor=RED, markeredgecolor=RED, markersize=4.5, label="P"),
        Line2D([0], [0], marker="o", color="none", markerfacecolor=BLUE, markeredgecolor=BLUE, markersize=4.5, label="S"),
        Line2D([0], [0], marker="o", color="none", markerfacecolor="none", markeredgecolor=GRAY, markersize=4.5, label="removed"),
    ]
    ax_b.legend(handles=legend_handles, frameon=False, fontsize=6.0, loc="upper right", handletextpad=0.3)

    ax_c_head = fig.add_subplot(panel_c[0])
    ax_c_dist = fig.add_subplot(panel_c[1])
    ax_c_ret = fig.add_subplot(panel_c[2], sharex=ax_c_dist)
    add_header(ax_c_head, "C", "Distance coverage changes")
    distance_summary, distance_rows = draw_distance_statistics(ax_c_dist, ax_c_ret, record_rows)
    ax_c_dist.set_title("Phase-picking records", loc="left", fontsize=7.2, pad=2, fontweight="bold")
    plt.setp(ax_c_dist.get_xticklabels(), visible=False)

    fig.suptitle("Hard SNR thresholds reshape seismic observability", fontsize=10.3, fontweight="bold", y=0.98)
    fig.savefig(OUT_PDF, bbox_inches="tight")
    fig.savefig(OUT_PNG, bbox_inches="tight", dpi=300)
    write_exports(phase_points, len(pick_rows), disp_summaries, distance_summary, distance_rows)
    print(f"Wrote {OUT_PDF}")
    print(f"Wrote {OUT_PNG}")
    print(f"Wrote {OUT_CSV}")
    print(f"Wrote {OUT_JSON}")


if __name__ == "__main__":
    make_figure()