#!/usr/bin/env python3
"""Shape-based indicator-diagram comparison (pre/post repair, peer, chamber).

Uses the process-invariant normalised P-theta profile so load/level changes
cancel and only waveform shape remains.  Reports L2 shape distances and saves a
profile figure.  Read-only.
"""
from __future__ import annotations

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

import numpy as np

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

from indicator_features.extract_from_db import (  # noqa: E402
    CYLINDERS,
    SIDES,
    connect,
    fetch_cycles,
    fetch_file_rows,
    point_name,
    sample_per_day,
)
from indicator_features.indicator_features import CYLINDER_STAGE, normalized_profile  # noqa: E402


def collect(connection, unit, cylinders, sides, start, end, per_day, bins):
    profiles = defaultdict(list)
    for cylinder in cylinders:
        for side in sides:
            point = point_name(unit, cylinder, side)
            files = fetch_file_rows(connection, point, start, end)
            picked = sample_per_day(files, per_day)
            if not picked:
                continue
            id_map = {int(r["id"]): r for r in picked}
            for file_id, signal in fetch_cycles(connection, list(id_map), 200).items():
                prof = normalized_profile(signal, cylinder, bins=bins)
                if prof is not None:
                    profiles[(cylinder, side)].append(prof)
    return profiles


def median_profile(items):
    return np.median(np.vstack(items), axis=0) if items else None


STAGE_MEMBERS = {"一级": ("二", "四", "六"), "二级": ("三", "五"), "三级": ("一",)}


def peer_median(medians, cylinder, side):
    peers = STAGE_MEMBERS[CYLINDER_STAGE[cylinder]]
    same = [m for (c, s), m in medians.items() if s == side and c in peers and m is not None]
    return np.median(np.vstack(same), axis=0) if same else None


def main():
    parser = argparse.ArgumentParser(description="功图波形形状对照（修前/修后、同级、双腔）")
    parser.add_argument("--unit", required=True, choices=("7", "8", "9"))
    parser.add_argument("--cylinders", nargs="+", default=list(CYLINDERS), choices=CYLINDERS)
    parser.add_argument("--sides", nargs="+", default=list(SIDES), choices=SIDES)
    parser.add_argument("--pre", nargs=2, required=True)
    parser.add_argument("--post", nargs=2, required=True)
    parser.add_argument("--per-day", type=int, default=6)
    parser.add_argument("--bins", type=int, default=180)
    parser.add_argument("--out-dir", default=str(BACKEND / "cache" / "indicator" / "shape"))
    args = parser.parse_args()

    connection = connect()
    try:
        pre = collect(connection, args.unit, args.cylinders, args.sides, args.pre[0], args.pre[1], args.per_day, args.bins)
        post = collect(connection, args.unit, args.cylinders, args.sides, args.post[0], args.post[1], args.per_day, args.bins)
    finally:
        connection.close()

    pre_med = {k: median_profile(v) for k, v in pre.items() if v}
    post_med = {k: median_profile(v) for k, v in post.items() if v}
    print(f"{args.unit}号机  形状对照  修前 {args.pre[0]}~{args.pre[1]}  修后 {args.post[0]}~{args.post[1]}")
    print(f"{'缸侧':>6}{'样本前/后':>12}{'修前/后 L2':>12}{'peer前':>9}{'peer后':>9}{'双腔前':>9}{'双腔后':>9}")

    for key in sorted(set(pre_med) | set(post_med)):
        cylinder, side = key
        pm, qm = pre_med.get(key), post_med.get(key)
        pp = peer_median(pre_med, cylinder, side)
        qp = peer_median(post_med, cylinder, side)
        peer_pre = np.linalg.norm(pm - pp) if pm is not None and pp is not None else float("nan")
        peer_post = np.linalg.norm(qm - qp) if qm is not None and qp is not None else float("nan")
        other = "轴侧" if side == "盖侧" else "盖侧"
        ch_pre = np.linalg.norm(pm - pre_med[(cylinder, other)]) if pm is not None and (cylinder, other) in pre_med and pre_med[(cylinder, other)] is not None else float("nan")
        ch_post = np.linalg.norm(qm - post_med[(cylinder, other)]) if qm is not None and (cylinder, other) in post_med and post_med[(cylinder, other)] is not None else float("nan")
        cross = np.linalg.norm(pm - qm) if pm is not None and qm is not None else float("nan")
        print(f"{cylinder+side:>6}{len(pre.get(key,[])):>6}/{len(post.get(key,[])):<5}{cross:>12.3f}{peer_pre:>9.3f}{peer_post:>9.3f}{ch_pre:>9.3f}{ch_post:>9.3f}")

    # Save a figure (best effort).
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        out = Path(args.out_dir)
        out.mkdir(parents=True, exist_ok=True)
        keys = [k for k in sorted(set(pre_med) | set(post_med))]
        if keys:
            cols = 3
            rows = (len(keys) + cols - 1) // cols
            fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 3 * rows), squeeze=False)
            for ax, key in zip(axes.ravel(), keys):
                if pre_med.get(key) is not None:
                    ax.plot(pre_med[key], color="crimson", label="pre")
                if post_med.get(key) is not None:
                    ax.plot(post_med[key], color="royalblue", label="post")
                ax.set_title(f"{key[0]}缸{key[1]}", fontsize=10)
                ax.set_ylim(-0.2, 1.4)
                ax.legend(fontsize=7)
            for ax in axes.ravel()[len(keys):]:
                ax.axis("off")
            fig.suptitle(f"{args.unit}号机 归一化功图剖面 pre vs post")
            fig.tight_layout()
            path = out / f"{args.unit}_{args.pre[0]}_{args.post[0]}.png"
            fig.savefig(path, dpi=110)
            print("图:", path)
    except Exception as error:
        print("绘图跳过:", error)


if __name__ == "__main__":
    main()