#!/usr/bin/env python3
"""
visualize_mprim.py
-------------------
Vẽ nhanh tất cả primitive tại 1 góc bắt đầu (startangle) để kiểm tra bằng mắt
trước khi tích hợp vào planner thật.

Cách dùng:
    python3 visualize_mprim.py diff.mprim --angle 0
    python3 visualize_mprim.py omni.mprim --angle 4
"""

import argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def parse_mprim(path):
    with open(path) as f:
        lines = [l.strip() for l in f.readlines()]

    i = 0
    resolution = float(lines[i].split(":")[1]); i += 1
    numberofangles = int(lines[i].split(":")[1]); i += 1
    total = int(lines[i].split(":")[1]); i += 1

    prims = []
    while i < len(lines) and lines[i]:
        primid = int(lines[i].split(":")[1]); i += 1
        startangle = int(lines[i].split(":")[1]); i += 1
        endpose = lines[i].split(":")[1].split(); i += 1
        costmult = float(lines[i].split(":")[1]); i += 1
        n = int(lines[i].split(":")[1]); i += 1
        poses = []
        for _ in range(n):
            x, y, th = map(float, lines[i].split())
            poses.append((x, y, th))
            i += 1
        prims.append({
            "primid": primid, "startangle": startangle,
            "endpose_c": [int(v) for v in endpose],
            "costmult": costmult, "poses": poses
        })
    return resolution, numberofangles, prims


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("mprim_file")
    ap.add_argument("--angle", type=int, default=0, help="startangle_c cần vẽ")
    ap.add_argument("-o", "--output", default=None, help="file ảnh output (mặc định: <mprim>_angleN.png)")
    args = ap.parse_args()

    resolution, numberofangles, prims = parse_mprim(args.mprim_file)
    subset = [p for p in prims if p["startangle"] == args.angle]

    if not subset:
        print(f"Không tìm thấy primitive nào với startangle_c={args.angle}")
        return

    import math

    fig, ax = plt.subplots(figsize=(6, 6))
    tick = resolution * 0.4  # độ dài vạch heading tại mỗi pose
    for p in subset:
        xs = [pt[0] for pt in p["poses"]]
        ys = [pt[1] for pt in p["poses"]]
        line, = ax.plot(xs, ys, marker="o", markersize=2,
                        label=f'id{p["primid"]} (cost {p["costmult"]:g})')
        # vạch heading: kiểm tra bằng mắt heading có bám tangent quỹ đạo không
        # (bug kiểu "heading cuối lệch 10 độ" hiện rõ ở đây)
        for (x, y, th) in p["poses"]:
            ax.plot([x, x + tick * math.cos(th)], [y, y + tick * math.sin(th)],
                    color=line.get_color(), linewidth=0.6, alpha=0.6)
        ax.annotate(str(p["primid"]), (xs[-1], ys[-1]), fontsize=8)

    ax.set_aspect("equal")
    ax.grid(True)
    ax.set_title(f"Primitives tại startangle_c={args.angle} / {numberofangles} (resolution={resolution}m)")
    ax.set_xlabel("x (m)")
    ax.set_ylabel("y (m)")
    ax.legend(fontsize=7, loc="upper left", bbox_to_anchor=(1.02, 1))
    plt.tight_layout()

    out = args.output or args.mprim_file.replace(".mprim", f"_angle{args.angle}.png")
    plt.savefig(out, dpi=130)
    print(f"[OK] Đã lưu hình: {out}")


if __name__ == "__main__":
    main()