import subprocess
import sys
import os
import glob
from tqdm import tqdm

# Configuration
H5_PATH = "data/processed/gaze360_robust_v16_test_B.h5"
SWA_DIR = "checkpoints/gold_swa"
RESULTS_CSV = "logs/gaze360_swa_summary.csv"
RESULTS_LOG = "logs/gaze360_swa_summary.log"

def main():
    # 1. Identify all available SWA checkpoints
    swa_models = glob.glob(os.path.join(SWA_DIR, "swa_gold_p*.pt"))
    swa_models.sort()
    
    if not swa_models:
        print(f"No SWA models found in {SWA_DIR}")
        return

    print(f"Found {len(swa_models)} SWA models. Starting evaluation on Gaze360...")

    # 2. Prepare result storage
    results = []
    
    print(f"{'='*70}")
    print(f"{'Fold':<8} | {'MAE (All Cases)':<20} | {'MAE (Frontal 45)':<20}")
    print(f"{'-'*70}")

    with open(RESULTS_CSV, 'w') as f_csv:
        f_csv.write("Fold,MAE_All,MAE_Frontal\n")

        for model_path in tqdm(swa_models, desc="Evaluating Folds"):
            p_id = os.path.basename(model_path).split('_')[-1].replace('.pt', '')
            
            # Execute verify_gaze360.py
            cmd = [
                sys.executable, "src/verify_gaze360.py",
                "--model", model_path,
                "--h5", H5_PATH,
                "--invert_yaw",
                "--invert_pitch"
            ]
            
            try:
                # Capture output
                process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                stdout, stderr = process.communicate()
                
                # Parse output for MAE values
                mae_all = "N/A"
                mae_frontal = "N/A"
                
                for line in stdout.splitlines():
                    if "All Cases" in line and "|" in line:
                        mae_all = line.split('|')[-1].strip()
                    if "Frontal +/- 45" in line and "|" in line:
                        mae_frontal = line.split('|')[-1].strip()
                
                print(f"{p_id:<8} | {mae_all:<20} | {mae_frontal:<20}")
                f_csv.write(f"{p_id},{mae_all},{mae_frontal}\n")
                
                if mae_all != "N/A":
                    results.append({
                        'fold': p_id,
                        'all': float(mae_all),
                        'frontal': float(mae_frontal) if mae_frontal != "N/A" else 0.0
                    })
            except Exception as e:
                print(f"Error evaluating {p_id}: {e}")

    # 3. Final Summary
    if results:
        avg_all = sum(r['all'] for r in results) / len(results)
        avg_frontal = sum(r['frontal'] for r in results) / len(results)
        
        summary_str = (
            f"\n{'='*70}\n"
            f" FINAL SUMMARY (Gaze360 SWA Evaluation)\n"
            f"{'-'*70}\n"
            f" Average MAE (All Cases):    {avg_all:.4f} deg\n"
            f" Average MAE (Frontal 45):   {avg_frontal:.4f} deg\n"
            f" Number of Folds evaluated:  {len(results)}\n"
            f"{'='*70}\n"
        )
        print(summary_str)
        
        with open(RESULTS_LOG, 'w') as f_log:
            f_log.write(summary_str)
            f_log.write("\nDetailed Results:\n")
            for r in results:
                f_log.write(f"Fold {r['fold']}: All={r['all']:.4f}, Frontal={r['frontal']:.4f}\n")

if __name__ == "__main__":
    main()