import numpy as np
import os
from pathlib import Path
import pandas as pd

def compare_folders(base_path, iter_a, iter_b, start_frame=190, end_frame=220, model="awrl1432"):
    path_a = Path(base_path) / "iterations" / iter_a / model
    path_b = Path(base_path) / "iterations" / iter_b / model
    
    if not path_a.exists() or not path_b.exists():
        print(f"[ERROR] Paths not found. A: {path_a}, B: {path_b}")
        return

    # Filter files in range
    files_a = sorted(list(path_a.glob("*.npy")))
    filtered_a = []
    for f in files_a:
        try:
            frame_num = int(f.stem.split('_')[-1])
            if start_frame <= frame_num <= end_frame:
                filtered_a.append(f)
        except:
            continue

    comparison_data = []

    print(f"\n🚀 Comparing {iter_a} vs {iter_b} (Frames {start_frame}-{end_frame}) for {model}...")
    print(f"{'Frame':<15} | {'Pts (A)':<10} | {'Pts (B)':<10} | {'Pts Delta %':<15} | {'Mean Mag (A)':<15} | {'Mean Mag (B)':<15} | {'Mag Delta %':<15}")
    print("-" * 115)

    for f_a in filtered_a:
        f_b = path_b / f_a.name
        if not f_b.exists():
            continue
            
        data_a = np.load(f_a)
        data_b = np.load(f_b)
        
        count_a = data_a.shape[0]
        count_b = data_b.shape[0]
        pts_delta = ((count_b - count_a) / count_a) * 100 if count_a > 0 else 0
        
        mag_a = np.mean(data_a[:, 4]) if count_a > 0 else 0
        mag_b = np.mean(data_b[:, 4]) if count_b > 0 else 0
        mag_delta = ((mag_b - mag_a) / mag_a) * 100 if mag_a > 0 else 0
        
        frame_name = f_a.stem
        print(f"{frame_name:<15} | {count_a:<10} | {count_b:<10} | {pts_delta:<15.1f}% | {mag_a:<15.4f} | {mag_b:<15.4f} | {mag_delta:<15.1f}%")
        
        comparison_data.append({
            "frame": frame_name,
            "pts_a": count_a,
            "pts_b": count_b,
            "mag_a": mag_a,
            "mag_b": mag_b
        })

    # Summary Statistics
    if comparison_data:
        df = pd.DataFrame(comparison_data)
        print("\n📊 AVG SUMMARY (Iteration A -> B):")
        print(f"Avg Point Change:     {((df['pts_b'].sum() - df['pts_a'].sum()) / df['pts_a'].sum())*100:+.2f}%")
        print(f"Avg Magnitude Change: {((df['mag_b'].mean() - df['mag_a'].mean()) / df['mag_a'].mean())*100:+.2f}%")
        print(f"Total Points A:       {df['pts_a'].sum()}")
        print(f"Total Points B:       {df['pts_b'].sum()}")
    else:
        print("[ERROR] No overlapping data found for comparison.")

if __name__ == "__main__":
    project_root = Path(__file__).parent.parent.parent
    base = project_root / "Shenron_debug"
    # Comparing 14b (Old Stable) vs 16 (New Smooth Physics)
    compare_folders(base, "14b_normalization", "16_resolution_independent", start_frame=190, end_frame=225)