#!/usr/bin/env python3
import argparse
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

def load_unique_faces(wide_csv: str) -> pd.DataFrame:
    df = pd.read_csv(wide_csv)
    # Keep the first occurrence per face_id (faces may be listed from multiple tets)
    df = df.sort_values(["face_id", "tet_id"]).drop_duplicates(subset=["face_id"], keep="first")
    return df

def plot_faces(df: pd.DataFrame, out_png: str, view="xy", linewidth=1.5, annotate=True):
    # Choose which plane to show
    axes = {"xy": ("_x","_y"), "xz": ("_x","_z"), "yz": ("_y","_z")}
    if view not in axes:
        raise ValueError("view must be one of: xy, xz, yz")
    axa, axb = axes[view]

    fig, ax = plt.subplots(figsize=(10, 10))  # larger square figure

    all_x, all_y = [], []

    # Plot each face as a closed triangle
    for _, r in df.iterrows():
        x = [r["face0"+axa], r["face1"+axa], r["face2"+axa], r["face0"+axa]]
        y = [r["face0"+axb], r["face1"+axb], r["face2"+axb], r["face0"+axb]]
        ax.plot(x, y, "-k", lw=linewidth)

        all_x.extend(x)
        all_y.extend(y)

        if annotate:
            cx = (r["face0"+axa] + r["face1"+axa] + r["face2"+axa]) / 3.0
            cy = (r["face0"+axb] + r["face1"+axb] + r["face2"+axb]) / 3.0
            ax.annotate(f'{int(r["face_id"])}', (cx, cy),
                        textcoords="offset points", xytext=(0, 0),
                        ha="center", va="center", fontsize=8,
                        bbox=dict(boxstyle="round,pad=0.2", fc="yellow", ec="black", lw=0.5))

    # Square bounding box
    if all_x and all_y:
        min_x, max_x = min(all_x), max(all_x)
        min_y, max_y = min(all_y), max(all_y)
        span = max(max_x - min_x, max_y - min_y)
        cx = 0.5 * (max_x + min_x)
        cy = 0.5 * (max_y + min_y)
        ax.set_xlim(cx - span/2, cx + span/2)
        ax.set_ylim(cy - span/2, cy + span/2)

    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel(view[0])
    ax.set_ylabel(view[1])
    ax.grid(True, ls="--", alpha=0.4)
    ax.set_title(f"Faces with face_id annotations ({view}-view)")

    # Expand to fill entire figure
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)

    out_png = Path(out_png)
    out_png.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_png, dpi=220, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved {out_png}")

def main():
    ap = argparse.ArgumentParser(description="Plot faces from parsed CUDA debug CSV and annotate face_id.")
    ap.add_argument("--wide-csv", required=True, help="Path to *_wide.csv produced by the parser.")
    ap.add_argument("--out", default="faces_annotated.png", help="Output PNG path.")
    ap.add_argument("--view", choices=["xy","xz","yz"], default="xy", help="Projection plane to plot.")
    ap.add_argument("--linewidth", type=float, default=1.5, help="Triangle edge width.")
    ap.add_argument("--no-annotate", action="store_true", help="Disable face_id annotations.")
    args = ap.parse_args()

    df = load_unique_faces(args.wide_csv)
    if df.empty:
        print("No faces found in CSV.")
        return

    plot_faces(df, args.out, view=args.view, linewidth=args.linewidth, annotate=not args.no_annotate)

if __name__ == "__main__":
    main()