#!/opt/homebrew/bin/python3.11
"""
Face Trace Visualizer - Visualize face tracking paths

Output:
1. Trace path visualization (matplotlib)
2. Trace statistics CSV
"""

import sys
import json
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from collections import defaultdict
from typing import Dict


def visualize_traces(face_data: Dict, output_path: str = None) -> None:
    """
    Visualize face trace paths
    """
    frames = face_data.get("frames", {})
    traces = face_data.get("traces", {})
    metadata = face_data.get("metadata", {})
    
    if not frames or not traces:
        print("No frames or traces found")
        return
    
    video_width = metadata.get("width", 640)
    video_height = metadata.get("height", 360)
    video_duration = metadata.get("total_duration", 15)
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    
    ax1 = axes[0, 0]
    ax2 = axes[0, 1]
    ax3 = axes[1, 0]
    ax4 = axes[1, 1]
    
    colors = plt.cm.tab10(np.linspace(0, 1, len(traces)))
    
    trace_data = {}
    for trace_id_str, trace in traces.items():
        trace_id = int(trace_id_str)
        path = trace.get("path", [])
        
        trace_data[trace_id] = {
            "frames": [p["frame"] for p in path],
            "x": [p["bbox"]["x"] + p["bbox"]["width"] / 2 for p in path],
            "y": [p["bbox"]["y"] + p["bbox"]["height"] / 2 for p in path],
            "confidence": [p["confidence"] for p in path],
            "pose": [p["pose_angle"] for p in path],
        }
    
    for trace_id, color in zip(sorted(trace_data.keys()), colors):
        data = trace_data[trace_id]
        
        ax1.plot(data["frames"], data["x"], color=color, label=f"Trace {trace_id}", linewidth=2)
        ax1.scatter(data["frames"], data["x"], color=color, s=30)
        
        ax2.plot(data["frames"], data["y"], color=color, label=f"Trace {trace_id}", linewidth=2)
        ax2.scatter(data["frames"], data["y"], color=color, s=30)
        
        ax3.plot(data["frames"], data["confidence"], color=color, label=f"Trace {trace_id}", linewidth=2)
        ax3.scatter(data["frames"], data["confidence"], color=color, s=30)
    
    ax1.set_xlabel("Frame Number")
    ax1.set_ylabel("X Position (center)")
    ax1.set_title("Face X Position Over Time")
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    ax2.set_xlabel("Frame Number")
    ax2.set_ylabel("Y Position (center)")
    ax2.set_title("Face Y Position Over Time")
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    ax3.set_xlabel("Frame Number")
    ax3.set_ylabel("Detection Confidence")
    ax3.set_title("Face Detection Confidence Over Time")
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    
    pose_colors = {
        "frontal": "green",
        "three_quarter": "blue",
        "profile_left": "orange",
        "profile_right": "red",
        "unknown": "gray",
    }
    
    for trace_id, color in zip(sorted(trace_data.keys()), colors):
        data = trace_data[trace_id]
        poses = data["pose"]
        frames = data["frames"]
        
        pose_counts = defaultdict(int)
        for pose in poses:
            pose_counts[pose] += 1
        
        ax4.bar(
            [f"Trace {trace_id}\n{pose}" for pose in pose_counts.keys()],
            pose_counts.values(),
            color=[pose_colors.get(pose, "gray") for pose in pose_counts.keys()],
            alpha=0.7,
            label=f"Trace {trace_id}",
        )
    
    ax4.set_xlabel("Trace / Pose")
    ax4.set_ylabel("Count")
    ax4.set_title("Pose Distribution by Trace")
    ax4.tick_params(axis='x', rotation=45)
    
    plt.tight_layout()
    
    if output_path:
        plt.savefig(output_path, dpi=150, bbox_inches="tight")
        print(f"\n✅ Visualization saved to: {output_path}")
    else:
        plt.show()


def export_trace_csv(face_data: Dict, output_path: str) -> None:
    """
    Export trace statistics to CSV
    """
    traces = face_data.get("traces", {})
    
    import csv
    
    with open(output_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "trace_id",
            "start_frame",
            "end_frame",
            "duration_frames",
            "duration_seconds",
            "total_appearances",
            "avg_confidence",
            "pose_three_quarter",
            "pose_profile_right",
            "pose_profile_left",
            "pose_frontal",
        ])
        
        for trace_id_str, trace in sorted(traces.items(), key=lambda x: int(x[0])):
            poses = trace.get("pose_angles", [])
            pose_counts = defaultdict(int)
            for pose in poses:
                pose_counts[pose] += 1
            
            writer.writerow([
                trace["trace_id"],
                trace["start_frame"],
                trace["end_frame"],
                trace["duration_frames"],
                trace["duration_seconds"],
                trace["total_appearances"],
                trace["avg_confidence"],
                pose_counts.get("three_quarter", 0),
                pose_counts.get("profile_right", 0),
                pose_counts.get("profile_left", 0),
                pose_counts.get("frontal", 0),
            ])
    
    print(f"\n✅ CSV exported to: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="Visualize face traces")
    parser.add_argument("--face-json", required=True, help="Path to face_traced.json")
    parser.add_argument("--output-plot", help="Output plot path (PNG)")
    parser.add_argument("--output-csv", help="Output CSV path")
    args = parser.parse_args()
    
    with open(args.face_json) as f:
        face_data = json.load(f)
    
    print("=" * 60)
    print("Face Trace Visualizer")
    print("=" * 60)
    print(f"\nInput: {args.face_json}")
    print(f"Traces: {len(face_data.get('traces', {}))}")
    
    if args.output_plot:
        visualize_traces(face_data, args.output_plot)
    
    if args.output_csv:
        export_trace_csv(face_data, args.output_csv)
    
    if not args.output_plot and not args.output_csv:
        visualize_traces(face_data)


if __name__ == "__main__":
    main()