#!/usr/bin/env python3
"""A simple, readable before/after figure for one measurement point.

Three panels, all in Chinese:

1. 平均压力 对 曲柄角度（停机前红 / 开机后蓝），高亮关注的角度段；
2. 波动 std 对 曲柄角度（同色）；
3. 关注角度段的 std 随时间变化，停机前红、开机后蓝，并标出事件日。
"""
from __future__ import annotations

import argparse
import sys
from collections import defaultdict
from datetime import datetime
from pathlib import Path

import numpy as np

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

matplotlib.rcParams["font.sans-serif"] = ["Arial Unicode MS", "Heiti TC", "Hiragino Sans GB", "DejaVu Sans"]
matplotlib.rcParams["axes.unicode_minus"] = False

BACKEND = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(BACKEND))

import cn_names as CN  # noqa: E402
from event_bins import collect_after, collect_before, load  # noqa: E402
from run_segments import unit_segments  # noqa: E402


def med(rows, metric, bin_index):
    vals = [r[metric] for r in rows if r["bin"] == bin_index and r[metric] is not None]
    return float(np.median(vals)) if vals else np.nan


def main() -> None:
    parser = argparse.ArgumentParser(description="单点 停机前/开机后 简明对照图")
    parser.add_argument("--csv", required=True)
    parser.add_argument("--unit", required=True)
    parser.add_argument("--date", required=True)
    parser.add_argument("--cylinder", required=True)
    parser.add_argument("--side", required=True)
    parser.add_argument("--highlight", default="", help="关注角度段，如 140-150")
    parser.add_argument("--series-metric", default="std", help="右图随时间画哪个指标（默认 std）")
    parser.add_argument("--pre-days", type=int, default=15)
    parser.add_argument("--post-days", type=int, default=15)
    parser.add_argument("--max-gap", type=int, default=30)
    parser.add_argument("--merge-gap", type=int, default=2)
    parser.add_argument("--out", required=True)
    args = parser.parse_args()

    rows = load(args.csv)
    point = [r for r in rows if r["cylinder"] == args.cylinder and r["side"] == args.side]
    t0 = datetime.fromisoformat(args.date)
    segments = unit_segments(args.unit, args.merge_gap)
    pre, _ = collect_before(point, t0, segments, args.pre_days)
    post, _ = collect_after(point, t0, segments, args.post_days, args.max_gap)
    bins = sorted({r["bin"] for r in point})
    bin_start = {r["bin"]: r["bin_start"] for r in point}
    bin_deg = round(bin_start[bins[1]] - bin_start[bins[0]], 6) if len(bins) > 1 else 10.0
    centers = [bin_start[b] + bin_deg / 2 for b in bins]

    lo = hi = None
    if "-" in args.highlight:
        lo, hi = (float(x) for x in args.highlight.split("-"))
    hi_bins = [b for b in bins if lo is not None and lo <= bin_start[b] < hi]

    def daily_series(rows, bin_index, metric):
        by_day = defaultdict(list)
        for r in rows:
            if r["bin"] == bin_index and r[metric] is not None:
                by_day[r["time"].date()].append(r[metric])
        days = sorted(by_day)
        return days, [float(np.median(by_day[d])) for d in days]

    fig, axes = plt.subplots(1, 3, figsize=(15, 3.8))

    for ax, metric, name in ((axes[0], "mean", "平均压力 (MPa)"), (axes[1], "std", "波动 std (MPa)")):
        ax.plot(centers, [med(pre, metric, b) for b in bins], color="crimson", lw=1.8, label="停机前")
        ax.plot(centers, [med(post, metric, b) for b in bins], color="royalblue", lw=1.8, label="开机后")
        if hi_bins:
            ax.axvspan(lo, hi, color="orange", alpha=0.25)
        ax.set_xlim(0, 360)
        ax.set_xticks([0, 90, 180, 270, 360])
        ax.set_xlabel("曲柄角度 (°)")
        ax.set_title(name)
        ax.grid(alpha=0.25)
    axes[0].legend(fontsize=9)

    ax = axes[2]
    if hi_bins:
        b0 = hi_bins[len(hi_bins) // 2]
        days, values = daily_series(point, b0, args.series_metric)
        for d, v in zip(days, values):
            color = "royalblue" if datetime.combine(d, datetime.min.time()) >= t0 else "crimson"
            ax.plot(d, v, "o", color=color, ms=4)
        ax.axvline(t0, color="k", ls="--", lw=1)
        ax.set_title(f"关注段 {lo:g}-{hi:g}° 的 {CN.metric_col(args.series_metric)} 随时间")
    else:
        ax.set_title("（未指定关注段）")
    ax.set_xlabel("日期")
    ax.grid(alpha=0.25)
    ax.tick_params(axis="x", labelrotation=30)

    fig.suptitle(f"{args.unit}号机 {args.date}  {args.cylinder}缸{args.side}  停机前 vs 开机后")
    fig.tight_layout()
    Path(args.out).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(args.out, dpi=120)
    print("图:", args.out)
    print(f"前段周期 {len({r['wave_file_id'] for r in pre})}，后段周期 {len({r['wave_file_id'] for r in post})}")
    if hi_bins:
        pre_b = [r for r in pre if r["bin"] == b0]
        post_b = [r for r in post if r["bin"] == b0]
        print(f"关注段 {lo:g}-{hi:g}° {CN.metric_col(args.series_metric)}: 前中位 {np.median([r[args.series_metric] for r in pre_b]):.5g} -> 后中位 {np.median([r[args.series_metric] for r in post_b]):.5g}")


if __name__ == "__main__":
    main()