#!/usr/bin/env python3
"""Generate public, derived figures from the private historical CSV archive."""

from __future__ import annotations

import argparse
import csv
import math
import shutil
from collections import defaultdict
from pathlib import Path
from typing import Iterable

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np


BATTERY_SOURCE = Path("10_电池续航测试/battery_vbin_B4.csv")
DRIFT_SOURCE = Path("05_漂移测试/drift_949534_20260530_105446.csv")
JOINT_SOURCE = Path("09_仿真与真机对比验证/processed/real_dynamic_metrics.csv")
SIMULATION_SOURCE = Path("09_仿真与真机对比验证/processed/sim_bo_trials.csv")

BLUE = "#0072B2"
ORANGE = "#E69F00"
GREEN = "#009E73"
RED = "#D55E00"
PURPLE = "#CC79A7"
GRAY = "#6B7280"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Create derived historical evaluation tables and figures."
    )
    parser.add_argument(
        "--data-root",
        type=Path,
        required=True,
        help="Root of the private IMU test-data archive.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path("."),
        help="Repository root for analysis and docs output (default: current directory).",
    )
    return parser.parse_args()


def read_rows(path: Path) -> list[dict[str, str]]:
    if not path.is_file():
        raise FileNotFoundError(f"Required source CSV was not found: {path}")
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        return list(csv.DictReader(handle))


def write_rows(path: Path, fieldnames: list[str], rows: Iterable[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames, lineterminator="\n")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def configure_plotting() -> None:
    plt.rcParams.update(
        {
            "axes.spines.top": False,
            "axes.spines.right": False,
            "axes.grid": True,
            "grid.alpha": 0.22,
            "grid.linewidth": 0.7,
            "axes.labelsize": 10,
            "axes.titlesize": 11,
            "legend.fontsize": 9,
            "font.size": 10,
            "savefig.facecolor": "white",
            "figure.facecolor": "white",
        }
    )


def save_figure(fig: plt.Figure, stem: str, figure_dir: Path, docs_figure_dir: Path) -> None:
    figure_dir.mkdir(parents=True, exist_ok=True)
    docs_figure_dir.mkdir(parents=True, exist_ok=True)
    png_path = figure_dir / f"{stem}.png"
    pdf_path = figure_dir / f"{stem}.pdf"
    fig.savefig(png_path, dpi=300, bbox_inches="tight")
    fig.savefig(pdf_path, bbox_inches="tight")
    shutil.copyfile(png_path, docs_figure_dir / png_path.name)
    plt.close(fig)


def as_float(row: dict[str, str], key: str) -> float:
    return float(row[key].strip())


def make_battery_assets(
    data_root: Path,
    metrics_dir: Path,
    figure_dir: Path,
    docs_figure_dir: Path,
) -> list[dict[str, object]]:
    rows = read_rows(data_root / BATTERY_SOURCE)
    selected = [
        row
        for row in rows
        if row.get("是否采集", "").strip() == "是"
        and row.get("每区间使用分钟数", "").strip()
        and row.get("到关机累计剩余分钟(从该区间放电至3.20V)", "").strip()
    ]
    selected.sort(key=lambda row: as_float(row, "电压中点(V)"), reverse=True)
    total_runtime_min = max(
        as_float(row, "到关机累计剩余分钟(从该区间放电至3.20V)") for row in selected
    )
    derived = []
    for row in selected:
        remaining = as_float(row, "到关机累计剩余分钟(从该区间放电至3.20V)")
        derived.append(
            {
                "voltage_mid_v": f"{as_float(row, '电压中点(V)'):.2f}",
                "charge_percent": row["对应电量百分比(%)"].strip(),
                "bin_duration_min": f"{as_float(row, '每区间使用分钟数'):.2f}",
                "remaining_runtime_min": f"{remaining:.2f}",
                "elapsed_runtime_min": f"{total_runtime_min - remaining:.2f}",
            }
        )
    write_rows(
        metrics_dir / "battery_discharge_profile.csv",
        [
            "voltage_mid_v",
            "charge_percent",
            "bin_duration_min",
            "remaining_runtime_min",
            "elapsed_runtime_min",
        ],
        derived,
    )

    elapsed_hours = np.array([float(row["elapsed_runtime_min"]) / 60.0 for row in derived])
    voltage_v = np.array([float(row["voltage_mid_v"]) for row in derived])
    fig, axis = plt.subplots(figsize=(6.4, 3.6))
    axis.plot(elapsed_hours, voltage_v, color=BLUE, linewidth=2.0)
    axis.scatter(elapsed_hours, voltage_v, color=BLUE, s=12, zorder=3)
    axis.set_xlabel("Elapsed runtime (h)")
    axis.set_ylabel("Cell voltage (V)")
    axis.set_title("Archived Battery Discharge Profile")
    axis.set_xlim(left=0)
    save_figure(fig, "battery_discharge_profile", figure_dir, docs_figure_dir)

    return [
        {
            "evaluation_id": "battery_discharge",
            "metric_name": "usable_runtime",
            "statistic": "maximum_recorded",
            "value": f"{total_runtime_min:.2f}",
            "unit": "min",
            "sample_count": len(selected),
            "source_relative_path": BATTERY_SOURCE.as_posix(),
            "scope": "one archived node battery discharge profile",
            "limitation": "load and configuration metadata are incomplete",
        },
        {
            "evaluation_id": "battery_discharge",
            "metric_name": "measured_voltage_range",
            "statistic": "min_to_max",
            "value": f"{voltage_v.min():.2f} to {voltage_v.max():.2f}",
            "unit": "V",
            "sample_count": len(selected),
            "source_relative_path": BATTERY_SOURCE.as_posix(),
            "scope": "valid voltage bins in one archived profile",
            "limitation": "not a guaranteed runtime specification",
        },
    ]


def make_drift_assets(
    data_root: Path,
    metrics_dir: Path,
    figure_dir: Path,
    docs_figure_dir: Path,
) -> list[dict[str, object]]:
    rows = [
        row
        for row in read_rows(data_root / DRIFT_SOURCE)
        if row.get("rest", "").strip() == "1"
    ]
    times = np.array([as_float(row, "t") for row in rows])
    axes = {
        "roll": np.array([as_float(row, "roll") for row in rows]),
        "pitch": np.array([as_float(row, "pitch") for row in rows]),
        "yaw": np.array([as_float(row, "yaw") for row in rows]),
    }
    start_time = float(times[0])
    baseline_mask = times <= start_time + 5.0
    final_mask = times >= float(times[-1]) - 5.0
    baselines = {name: float(np.median(values[baseline_mask])) for name, values in axes.items()}
    end_values = {name: float(np.median(values[final_mask])) for name, values in axes.items()}
    delta_values = {name: values - baselines[name] for name, values in axes.items()}

    second_index = np.floor(times - start_time).astype(int)
    trace_rows: list[dict[str, object]] = []
    for second in range(int(second_index.max()) + 1):
        mask = second_index == second
        if not np.any(mask):
            continue
        trace_rows.append(
            {
                "time_s": f"{float(np.mean(times[mask] - start_time)):.3f}",
                "roll_delta_deg": f"{float(np.mean(delta_values['roll'][mask])):.6f}",
                "pitch_delta_deg": f"{float(np.mean(delta_values['pitch'][mask])):.6f}",
                "yaw_delta_deg": f"{float(np.mean(delta_values['yaw'][mask])):.6f}",
            }
        )
    write_rows(
        metrics_dir / "static_attitude_change_trace.csv",
        ["time_s", "roll_delta_deg", "pitch_delta_deg", "yaw_delta_deg"],
        trace_rows,
    )

    fig, axis = plt.subplots(figsize=(6.4, 3.6))
    trace_time = np.array([float(row["time_s"]) for row in trace_rows])
    axis.plot(trace_time, [float(row["roll_delta_deg"]) for row in trace_rows], label="Roll", color=BLUE)
    axis.plot(trace_time, [float(row["pitch_delta_deg"]) for row in trace_rows], label="Pitch", color=ORANGE)
    axis.plot(trace_time, [float(row["yaw_delta_deg"]) for row in trace_rows], label="Yaw", color=GREEN)
    axis.axhline(0.0, color=GRAY, linewidth=0.8)
    axis.set_xlabel("Time after baseline (s)")
    axis.set_ylabel("Attitude change (deg)")
    axis.set_title("Archived Static Attitude Change")
    axis.legend(frameon=False, ncol=3, loc="upper left")
    save_figure(fig, "static_attitude_change", figure_dir, docs_figure_dir)

    metric_rows: list[dict[str, object]] = [
        {
            "evaluation_id": "static_attitude_change",
            "metric_name": "recording_duration",
            "statistic": "observed",
            "value": f"{times[-1] - start_time:.3f}",
            "unit": "s",
            "sample_count": len(rows),
            "source_relative_path": DRIFT_SOURCE.as_posix(),
            "scope": "rest-marked samples after initialization",
            "limitation": "single archived recording without current configuration snapshot",
        }
    ]
    for name in ("roll", "pitch", "yaw"):
        metric_rows.append(
            {
                "evaluation_id": "static_attitude_change",
                "metric_name": f"{name}_end_minus_baseline",
                "statistic": "median_last_5s_minus_median_first_5s",
                "value": f"{end_values[name] - baselines[name]:.6f}",
                "unit": "deg",
                "sample_count": len(rows),
                "source_relative_path": DRIFT_SOURCE.as_posix(),
                "scope": "rest-marked samples after initialization",
                "limitation": "Euler-angle change is not an external accuracy measurement",
            }
        )
        metric_rows.append(
            {
                "evaluation_id": "static_attitude_change",
                "metric_name": f"{name}_standard_deviation",
                "statistic": "sample_standard_deviation",
                "value": f"{float(np.std(delta_values[name], ddof=1)):.6f}",
                "unit": "deg",
                "sample_count": len(rows),
                "source_relative_path": DRIFT_SOURCE.as_posix(),
                "scope": "rest-marked samples after initialization",
                "limitation": "not a multi-IMU synchronization measurement",
            }
        )
    return metric_rows


def short_joint_name(name: str) -> str:
    if name.startswith("r_"):
        name = name[2:]
    return name.replace("_", " ").title()


def make_joint_tracking_assets(
    data_root: Path,
    metrics_dir: Path,
    figure_dir: Path,
    docs_figure_dir: Path,
) -> list[dict[str, object]]:
    rows = read_rows(data_root / JOINT_SOURCE)
    derived = []
    for row in rows:
        derived.append(
            {
                "condition": row["condition"].strip(),
                "date": row["date"].strip(),
                "trajectory": row["trajectory"].strip(),
                "joint": row["joint"].strip(),
                "rmse_deg": f"{as_float(row, 'rmse_deg'):.6f}",
                "mae_deg": f"{as_float(row, 'mae_deg'):.6f}",
                "max_abs_error_deg": f"{as_float(row, 'max_abs_error_deg'):.6f}",
                "tracking_rate": f"{as_float(row, 'tracking_rate'):.6f}",
                "n_samples": row["n_samples"].strip(),
                "duration_s": f"{as_float(row, 'duration_s'):.6f}",
            }
        )
    write_rows(
        metrics_dir / "historical_robot_joint_tracking_metrics.csv",
        [
            "condition",
            "date",
            "trajectory",
            "joint",
            "rmse_deg",
            "mae_deg",
            "max_abs_error_deg",
            "tracking_rate",
            "n_samples",
            "duration_s",
        ],
        derived,
    )

    labels = [short_joint_name(row["joint"]) for row in derived]
    positions = np.arange(len(derived))
    rmse = np.array([float(row["rmse_deg"]) for row in derived])
    mae = np.array([float(row["mae_deg"]) for row in derived])
    width = 0.36
    fig, axis = plt.subplots(figsize=(6.4, 3.6))
    axis.bar(positions - width / 2, rmse, width, label="RMSE", color=BLUE)
    axis.bar(positions + width / 2, mae, width, label="MAE", color=ORANGE)
    axis.set_xticks(positions)
    axis.set_xticklabels(labels, rotation=0)
    axis.set_ylabel("Tracking error (deg)")
    axis.set_title("Historical Robot Joint Tracking")
    axis.legend(frameon=False)
    save_figure(fig, "historical_joint_tracking_metrics", figure_dir, docs_figure_dir)

    metric_rows: list[dict[str, object]] = []
    for row in derived:
        for metric_name in ("rmse_deg", "mae_deg", "max_abs_error_deg"):
            metric_label = metric_name[:-4]
            metric_rows.append(
                {
                    "evaluation_id": "historical_robot_joint_tracking",
                    "metric_name": f"{row['joint']}_{metric_label}",
                    "statistic": metric_label,
                    "value": row[metric_name],
                    "unit": "deg",
                    "sample_count": row["n_samples"],
                    "source_relative_path": JOINT_SOURCE.as_posix(),
                    "scope": f"{row['condition']} {row['trajectory']} on one historical robot path",
                    "limitation": "not a body-tracker accuracy result or safety guarantee",
                }
            )
    return metric_rows


def quantile(values: list[float], q: float) -> float:
    return float(np.percentile(np.array(values), q))


def make_simulation_assets(
    data_root: Path,
    metrics_dir: Path,
    figure_dir: Path,
    docs_figure_dir: Path,
) -> list[dict[str, object]]:
    rows = read_rows(data_root / SIMULATION_SOURCE)
    grouped: dict[tuple[str, str], list[dict[str, str]]] = defaultdict(list)
    for row in rows:
        grouped[(row["joint_name"].strip(), row["trajectory"].strip())].append(row)

    progress: dict[tuple[str, int], list[float]] = defaultdict(list)
    for (_, trajectory), trials in grouped.items():
        trials.sort(key=lambda row: int(row["iteration"]))
        first_cost = as_float(trials[0], "best_so_far")
        if math.isclose(first_cost, 0.0):
            continue
        for row in trials:
            progress[(trajectory, int(row["iteration"]))].append(
                as_float(row, "best_so_far") / first_cost
            )

    derived = []
    for (trajectory, iteration), values in sorted(progress.items()):
        derived.append(
            {
                "trajectory": trajectory,
                "iteration": iteration,
                "median_relative_best_cost": f"{quantile(values, 50):.6f}",
                "p25_relative_best_cost": f"{quantile(values, 25):.6f}",
                "p75_relative_best_cost": f"{quantile(values, 75):.6f}",
                "group_count": len(values),
            }
        )
    write_rows(
        metrics_dir / "simulation_pd_optimization_progress.csv",
        [
            "trajectory",
            "iteration",
            "median_relative_best_cost",
            "p25_relative_best_cost",
            "p75_relative_best_cost",
            "group_count",
        ],
        derived,
    )

    fig, axis = plt.subplots(figsize=(6.4, 3.6))
    color_by_trajectory = {"sine": BLUE, "step": ORANGE}
    for trajectory in sorted({row["trajectory"] for row in derived}):
        subset = [row for row in derived if row["trajectory"] == trajectory]
        iterations = np.array([int(row["iteration"]) for row in subset])
        median = np.array([float(row["median_relative_best_cost"]) for row in subset])
        p25 = np.array([float(row["p25_relative_best_cost"]) for row in subset])
        p75 = np.array([float(row["p75_relative_best_cost"]) for row in subset])
        color = color_by_trajectory.get(trajectory, PURPLE)
        axis.plot(iterations, median, label=trajectory.title(), color=color, linewidth=2.0)
        axis.fill_between(iterations, p25, p75, color=color, alpha=0.16, linewidth=0)
    axis.set_xlabel("Optimization iteration")
    axis.set_ylabel("Relative best cost")
    axis.set_title("Simulation PD Optimization Progress")
    axis.legend(frameon=False, title="Trajectory")
    save_figure(fig, "simulation_pd_optimization_progress", figure_dir, docs_figure_dir)

    metric_rows = []
    for trajectory in sorted({row["trajectory"] for row in derived}):
        final_row = [row for row in derived if row["trajectory"] == trajectory][-1]
        metric_rows.append(
            {
                "evaluation_id": "simulation_pd_optimization",
                "metric_name": f"{trajectory}_final_relative_best_cost",
                "statistic": "median_across_joint_groups",
                "value": final_row["median_relative_best_cost"],
                "unit": "ratio",
                "sample_count": final_row["group_count"],
                "source_relative_path": SIMULATION_SOURCE.as_posix(),
                "scope": f"historical {trajectory} simulation tuning groups",
                "limitation": "controller-simulation result, not body-tracker accuracy",
            }
        )
    return metric_rows


def main() -> None:
    args = parse_args()
    data_root = args.data_root.expanduser().resolve()
    output_root = args.output_root.expanduser().resolve()
    metrics_dir = output_root / "analysis" / "metrics"
    figure_dir = output_root / "analysis" / "figures"
    docs_figure_dir = output_root / "docs" / "_static" / "figures"
    configure_plotting()

    summary_rows: list[dict[str, object]] = []
    summary_rows.extend(make_battery_assets(data_root, metrics_dir, figure_dir, docs_figure_dir))
    summary_rows.extend(make_drift_assets(data_root, metrics_dir, figure_dir, docs_figure_dir))
    summary_rows.extend(make_joint_tracking_assets(data_root, metrics_dir, figure_dir, docs_figure_dir))
    summary_rows.extend(make_simulation_assets(data_root, metrics_dir, figure_dir, docs_figure_dir))
    write_rows(
        metrics_dir / "historical_metric_summary.csv",
        [
            "evaluation_id",
            "metric_name",
            "statistic",
            "value",
            "unit",
            "sample_count",
            "source_relative_path",
            "scope",
            "limitation",
        ],
        summary_rows,
    )
    print(f"Wrote derived metrics to {metrics_dir}")
    print(f"Wrote figures to {figure_dir} and {docs_figure_dir}")


if __name__ == "__main__":
    main()