"""
viz.py

读取 sim_log.csv（由 test.py 生成）并用 matplotlib 动画显示：
- 主图：路径点（`example_path.json`）和割草机当前位置/航向箭头
- 下方两个子图：forward_sig 与 turn_sig 随时间的曲线，当前帧用竖线标示

运行：
    python viz.py

如果没有 `sim_log.csv`，脚本会提示先运行 `test.py` 生成数据。
"""
import os
import json
import csv
import sys
from math import cos, sin
import numpy as np

try:
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
except Exception as e:
    print('matplotlib required. Install with: pip install matplotlib')
    raise


SIM_CSV = 'sim_log.csv'
PATH_FILE = 'example_path.json'


def load_sim(csv_path=SIM_CSV):
    if not os.path.exists(csv_path):
        print(f"Sim log '{csv_path}' not found. Run test.py first to generate it.")
        sys.exit(1)
    t, xs, ys, headings, forward_sig, turn_sig = [], [], [], [], [], []
    with open(csv_path, 'r', newline='') as f:
        reader = csv.DictReader(f)
        for row in reader:
            t.append(float(row['t']))
            xs.append(float(row['x']))
            ys.append(float(row['y']))
            headings.append(float(row['heading']))
            forward_sig.append(int(row['forward_sig']))
            turn_sig.append(int(row['turn_sig']))
    return {'t': t, 'x': xs, 'y': ys, 'heading': headings, 'forward': forward_sig, 'turn': turn_sig}


def load_path(pth=PATH_FILE):
    if not os.path.exists(pth):
        return []
    with open(pth, 'r') as f:
        return json.load(f)


def animate_sim(simdata, path):
    t = simdata['t']
    x = simdata['x']
    y = simdata['y']
    h = simdata['heading']
    fwd = simdata['forward']
    trn = simdata['turn']

    fig = plt.figure(figsize=(10, 8))
    gs = fig.add_gridspec(3, 1, height_ratios=[3, 1, 1])

    ax_map = fig.add_subplot(gs[0])
    ax_f = fig.add_subplot(gs[1])
    ax_t = fig.add_subplot(gs[2])

    # map plot
    if path:
        px = [p[0] for p in path]
        py = [p[1] for p in path]
        ax_map.plot(px, py, '-k', lw=1, label='planned path')
        ax_map.scatter(px, py, c='k', s=10)

    line_pos, = ax_map.plot([], [], '-r', lw=2, label='actual')
    point_pos, = ax_map.plot([], [], 'ro')
    # heading arrow using quiver (initialize at first pose if available)
    if len(x) > 0 and len(h) > 0:
        quiv = ax_map.quiver([x[0]], [y[0]], [cos(h[0])], [sin(h[0])], angles='xy', scale_units='xy', scale=4, color='b')
    else:
        quiv = ax_map.quiver([0], [0], [1], [0], angles='xy', scale_units='xy', scale=4, color='b')

    ax_map.set_aspect('equal', 'box')
    ax_map.set_title('Mower path and pose')
    ax_map.legend()

    # forward subplot
    ax_f.plot(t, fwd, color='gray', alpha=0.5)
    f_line = ax_f.axvline(t[0], color='r')
    ax_f.set_ylabel('forward_sig')

    # turn subplot
    ax_t.plot(t, trn, color='gray', alpha=0.5)
    t_line = ax_t.axvline(t[0], color='r')
    ax_t.set_ylabel('turn_sig')
    ax_t.set_xlabel('time (s)')

    # set map limits
    margin = 1.0
    all_x = x + ( [p[0] for p in path] if path else [] )
    all_y = y + ( [p[1] for p in path] if path else [] )
    if all_x and all_y:
        ax_map.set_xlim(min(all_x)-margin, max(all_x)+margin)
        ax_map.set_ylim(min(all_y)-margin, max(all_y)+margin)

    def init():
        line_pos.set_data([], [])
        point_pos.set_data([], [])
        return line_pos, point_pos, quiv, f_line, t_line

    def update(i):
        # i is frame index
        xi = x[:i+1]
        yi = y[:i+1]
        line_pos.set_data(xi, yi)
        point_pos.set_data([x[i]], [y[i]])

        # update quiver in-place
        hx = cos(h[i])
        hy = sin(h[i])
        # set new offset and vector
        quiv.set_offsets(np.array([[x[i], y[i]]]))
        quiv.set_UVC(np.array([hx]), np.array([hy]))

        # update vertical lines
        f_line.set_xdata(t[i])
        t_line.set_xdata(t[i])

        return line_pos, point_pos, quiv, f_line, t_line

    interval_ms = max(10, int(1000.0 / 74.0))
    ani = FuncAnimation(fig, update, frames=len(t), init_func=init, blit=False, interval=interval_ms)
    plt.tight_layout()
    plt.show()


def main():
    sim = load_sim()
    path = load_path()
    animate_sim(sim, path)


if __name__ == '__main__':
    main()