"""可视化停车机器人四舵轮的模式切换与转向响应。

默认读取 MyParking/logs/wheel-speed 中最新的 *_snapshot.csv，输出四张 PNG：
1. 四个舵轮的目标角、实际角和模式切换时刻；
2. 四个舵轮的转角误差；
3. 四个转向 PID 输出；
4. 八个电机的最终命令速度与 CAN 反馈速度。

示例：
    python plot_steering_response.py
    python plot_steering_response.py --input "D:\\logs\\xxx_snapshot.csv"
"""

from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


SCRIPT_DIRECTORY = Path(__file__).resolve().parent
PROJECT_DIRECTORY = SCRIPT_DIRECTORY.parent.parent
DEFAULT_LOG_DIRECTORY = PROJECT_DIRECTORY / "logs" / "wheel-speed"

WHEELS = (
    ("LeftFront", "左前", "tab:blue"),
    ("LeftRear", "左后", "tab:orange"),
    ("RightFront", "右前", "tab:green"),
    ("RightRear", "右后", "tab:red"),
)

MOTORS = (
    ("LFL", "左前左"), ("LFR", "左前右"),
    ("LRL", "左后左"), ("LRR", "左后右"),
    ("RFL", "右前左"), ("RFR", "右前右"),
    ("RRL", "右后左"), ("RRR", "右后右"),
)

MODE_NAMES = {0: "正常", 1: "蟹行", 2: "自转"}


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="绘制四舵轮模式切换响应图")
    parser.add_argument("--input", type=Path, help="指定 *_snapshot.csv；缺省时取最新文件")
    parser.add_argument("--output", type=Path, help="图片输出目录；缺省时写入本次日志同级 plots")
    parser.add_argument("--from-seconds", type=float, default=0.0, help="从第几秒开始显示")
    parser.add_argument("--to-seconds", type=float, help="显示到第几秒结束")
    return parser.parse_args()


def find_snapshot(path: Path | None) -> Path:
    if path is not None:
        if not path.is_file():
            raise FileNotFoundError(f"找不到快照文件：{path}")
        return path

    candidates = sorted(
        DEFAULT_LOG_DIRECTORY.glob("*_snapshot.csv"),
        key=lambda item: item.stat().st_mtime,
        reverse=True,
    )
    if not candidates:
        raise FileNotFoundError(
            f"{DEFAULT_LOG_DIRECTORY} 中没有 *_snapshot.csv。\n"
            "请先在 M 层点击 StartWheelSpeedDiagnostic，完成模式切换后点击 StopWheelSpeedDiagnostic。"
        )
    return candidates[0]


def load_snapshot(path: Path) -> pd.DataFrame:
    frame = pd.read_csv(path, comment="#")
    required = {"ElapsedMs", "ManualControlMode", "SendThresSpeed"}
    missing = required.difference(frame.columns)
    if missing:
        raise ValueError(f"CSV 缺少字段：{', '.join(sorted(missing))}。请部署最新 MedullaAdapter.dll 后重新记录。")

    frame = frame.apply(pd.to_numeric, errors="coerce")
    frame = frame.dropna(subset=["ElapsedMs"]).sort_values("ElapsedMs")
    if frame.empty:
        raise ValueError("CSV 中没有有效数据行。")
    frame["ElapsedSeconds"] = (frame["ElapsedMs"] - frame["ElapsedMs"].iloc[0]) / 1000.0
    return frame


def crop(frame: pd.DataFrame, start: float, end: float | None) -> pd.DataFrame:
    result = frame[frame["ElapsedSeconds"] >= start]
    if end is not None:
        result = result[result["ElapsedSeconds"] <= end]
    if result.empty:
        raise ValueError("所选时间范围内没有数据。")
    return result


def require_columns(frame: pd.DataFrame, names: list[str]) -> None:
    missing = [name for name in names if name not in frame.columns]
    if missing:
        raise ValueError("CSV 缺少字段：" + ", ".join(missing))


def add_mode_markers(axis: plt.Axes, frame: pd.DataFrame) -> None:
    modes = frame["ManualControlMode"].round().astype("Int64")
    changes = modes.ne(modes.shift())
    for _, row in frame.loc[changes].iterrows():
        mode = int(row["ManualControlMode"])
        axis.axvline(row["ElapsedSeconds"], color="0.55", linestyle="--", linewidth=0.8, alpha=0.75)
        axis.text(
            row["ElapsedSeconds"], 0.99, MODE_NAMES.get(mode, f"模式{mode}"),
            transform=axis.get_xaxis_transform(), rotation=90,
            va="top", ha="right", fontsize=8, color="0.35",
        )


def save_steering_angle_plot(frame: pd.DataFrame, output: Path, prefix: str) -> None:
    required = []
    for key, _, _ in WHEELS:
        required.extend([f"TargetTh{key}", f"ActualTh{key}"])
    require_columns(frame, required)

    figure, axes = plt.subplots(2, 2, figsize=(14, 8), sharex=True)
    for axis, (key, label, color) in zip(axes.flat, WHEELS):
        time = frame["ElapsedSeconds"]
        axis.plot(time, frame[f"TargetTh{key}"], label="目标角", color=color, linewidth=1.8)
        axis.plot(time, frame[f"ActualTh{key}"], label="实际角", color="0.15", linewidth=1.1)
        axis.axhline(120, color="tab:red", linestyle=":", linewidth=0.8, label="机械限位 ±120°")
        axis.axhline(-120, color="tab:red", linestyle=":", linewidth=0.8)
        add_mode_markers(axis, frame)
        axis.set_title(f"{label}舵轮")
        axis.set_ylabel("转角 (deg)")
        axis.grid(alpha=0.25)
        axis.legend(loc="best", fontsize=8)
    for axis in axes[1]:
        axis.set_xlabel("时间 (s)")
    figure.suptitle("四舵轮目标转角与实际转角")
    figure.tight_layout()
    figure.savefig(output / f"{prefix}_steering_angles.png", dpi=180)
    plt.close(figure)


def save_error_and_pid_plot(frame: pd.DataFrame, output: Path, prefix: str) -> None:
    error_columns = [f"ErrorTh{key}" for key, _, _ in WHEELS]
    pid_columns = [f"PidOut{key}" for key, _, _ in WHEELS]
    require_columns(frame, error_columns + pid_columns)

    figure, axes = plt.subplots(2, 1, figsize=(14, 9), sharex=True)
    time = frame["ElapsedSeconds"]
    for key, label, color in WHEELS:
        axes[0].plot(time, frame[f"ErrorTh{key}"], label=label, color=color, linewidth=1.2)
        axes[1].plot(time, frame[f"PidOut{key}"], label=label, color=color, linewidth=1.2)
    axes[0].axhline(2, color="0.4", linestyle=":", linewidth=0.9, label="到位阈值 ±2°")
    axes[0].axhline(-2, color="0.4", linestyle=":", linewidth=0.9)
    for axis in axes:
        add_mode_markers(axis, frame)
        axis.grid(alpha=0.25)
        axis.legend(loc="best", ncol=3, fontsize=9)
    axes[0].set_ylabel("目标角 - 实际角 (deg)")
    axes[1].set_ylabel("转向 PID 输出 (m/s)")
    axes[1].set_xlabel("时间 (s)")
    figure.suptitle("转角误差与转向 PID 输出")
    figure.tight_layout()
    figure.savefig(output / f"{prefix}_steering_error_pid.png", dpi=180)
    plt.close(figure)


def save_motor_speed_plot(frame: pd.DataFrame, output: Path, prefix: str) -> None:
    command_columns = [f"Pid{name}" for name, _ in MOTORS]
    feedback_columns = [f"Actual{name}" for name, _ in MOTORS]
    require_columns(frame, command_columns + feedback_columns)

    figure, axes = plt.subplots(4, 2, figsize=(15, 12), sharex=True)
    time = frame["ElapsedSeconds"]
    for axis, (name, label) in zip(axes.flat, MOTORS):
        axis.plot(time, frame[f"Pid{name}"], label="最终命令", color="tab:blue", linewidth=1.2)
        axis.plot(time, frame[f"Actual{name}"], label="CAN反馈", color="tab:orange", linewidth=1.0)
        add_mode_markers(axis, frame)
        axis.set_title(f"{label}电机 ({name})")
        axis.set_ylabel("速度 (m/s)")
        axis.grid(alpha=0.25)
        axis.legend(loc="best", fontsize=8)
    for axis in axes[-1]:
        axis.set_xlabel("时间 (s)")
    figure.suptitle("八个电机最终速度命令与 CAN 实际速度反馈")
    figure.tight_layout()
    figure.savefig(output / f"{prefix}_motor_command_feedback.png", dpi=180)
    plt.close(figure)


def save_summary_plot(frame: pd.DataFrame, output: Path, prefix: str) -> None:
    require_columns(frame, ["SendThresSpeed"])
    figure, axes = plt.subplots(2, 1, figsize=(14, 7), sharex=True)
    time = frame["ElapsedSeconds"]
    axes[0].step(time, frame["ManualControlMode"], where="post", color="tab:purple", linewidth=1.5)
    axes[0].set_yticks([0, 1, 2], ["正常", "蟹行", "自转"])
    axes[0].set_ylabel("控制模式")
    axes[0].grid(alpha=0.25)
    axes[1].plot(time, frame["SendThresSpeed"], color="tab:brown", linewidth=1.4, label="SendThresSpeed")
    axes[1].set_ylabel("速度限幅 (m/s)")
    axes[1].set_xlabel("时间 (s)")
    axes[1].grid(alpha=0.25)
    axes[1].legend(loc="best")
    figure.suptitle("模式切换与整车下发速度限幅")
    figure.tight_layout()
    figure.savefig(output / f"{prefix}_mode_speed_limit.png", dpi=180)
    plt.close(figure)


def print_parameter_summary(frame: pd.DataFrame) -> None:
    parameter_names = [
        "DiffSteerKp", "DiffSteerKi", "DiffSteerKd", "DiffSteerMaxI",
        "DiffSteerDeadZone", "DiffSteerThresh", "DiffSteerSpeedAcc",
    ]
    if not set(parameter_names).issubset(frame.columns):
        return
    print("本次记录的转向 PID 参数：")
    print("  " + ", ".join(f"{name}={frame[name].iloc[0]:.6g}" for name in parameter_names))


def main() -> None:
    arguments = parse_arguments()
    snapshot_path = find_snapshot(arguments.input)
    frame = crop(load_snapshot(snapshot_path), arguments.from_seconds, arguments.to_seconds)
    output_directory = arguments.output or snapshot_path.parent / "plots"
    output_directory.mkdir(parents=True, exist_ok=True)
    prefix = snapshot_path.name.removesuffix("_snapshot.csv")

    save_steering_angle_plot(frame, output_directory, prefix)
    save_error_and_pid_plot(frame, output_directory, prefix)
    save_motor_speed_plot(frame, output_directory, prefix)
    save_summary_plot(frame, output_directory, prefix)
    print_parameter_summary(frame)
    print(f"已读取：{snapshot_path}")
    print(f"已生成四张图：{output_directory}")


if __name__ == "__main__":
    main()