import torch
import numpy as np
import h5py
import os
import sys
from pathlib import Path
from tqdm import tqdm

# Add project root to path
sys.path.append(str(Path(__file__).parent.parent))

from src.models.student import LIPEV2Student, LIPEV2StudentGold, LIPEV2StudentBaseline

def evaluate_on_gaze360(model_path, h5_path, invert_yaw=False, invert_pitch=False):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Evaluating model: {model_path}")
    print(f"On dataset: {h5_path}")
    print(f"Correction: Invert Yaw={invert_yaw}, Invert Pitch={invert_pitch}")
    print(f"Device: {device}")

    # Detect Architecture
    is_gold = 'gold' in model_path.lower() or 'id_5' in model_path.lower() or 'id_2' in model_path.lower() or 'id_6' in model_path.lower() or 'id_8' in model_path.lower()
    is_baseline = 'baseline' in model_path.lower()
    
    if is_gold:
        print("Architecture: V5-GOLD (DualPool / ID 5,6,8)")
        model = LIPEV2StudentGold().to(device)
    elif is_baseline:
        print("Architecture: Baseline (Addition)")
        model = LIPEV2StudentBaseline().to(device)
    else:
        print("Architecture: Standard (Concatenation / ID 1,3,4,7)")
        model = LIPEV2Student().to(device)

    state_dict = torch.load(model_path, map_location=device)
    
    # Handle SWA weights if necessary
    if 'n_averaged' in state_dict:
        # It's an AveragedModel from SWA
        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith('module.'):
                new_state_dict[k[7:]] = v
            else:
                new_state_dict[k] = v
        state_dict = new_state_dict

    model.load_state_dict(state_dict, strict=False)
    model.eval()

    results = {
        'all': {'error': 0.0, 'count': 0},
        'frontal_45': {'error': 0.0, 'count': 0}
    }

    with h5py.File(h5_path, 'r') as f:
        left_patches = f['left_patches'][:]
        right_patches = f['right_patches'][:]
        landmarks = f['landmarks'][:]
        gaze_gt = f['gaze'][:] # (pitch, yaw) in radians

        num_samples = left_patches.shape[0]
        
        with torch.no_grad():
            for i in tqdm(range(num_samples), desc="Testing"):
                # Prepare inputs
                lp = torch.from_numpy(left_patches[i]).float().unsqueeze(0).to(device) / 255.0
                rp = torch.from_numpy(right_patches[i]).float().unsqueeze(0).to(device) / 255.0
                lm = torch.from_numpy(landmarks[i]).float().view(1, -1).to(device)
                gt = torch.from_numpy(gaze_gt[i]).float().to(device)

                # Predict
                if is_gold:
                    out = model(lp, lm)
                else:
                    out = model(lp, lm, state='A')
                
                p_l, y_l = out[0], out[1]
                
                if is_gold:
                    out_r = model(rp, lm)
                else:
                    out_r = model(rp, lm, state='A')
                p_r, y_r = out_r[0], out_r[1]
                
                # Convert Logits to Degrees
                def logits_to_deg(p_logits, y_logits):
                    idx = torch.arange(90).float().to(device)
                    p_prob = torch.softmax(p_logits, dim=1)
                    y_prob = torch.softmax(y_logits, dim=1)
                    p_deg = (torch.sum(p_prob * idx, dim=1) * 2 - 90)
                    y_deg = (torch.sum(y_prob * idx, dim=1) * 2 - 90)
                    return p_deg, y_deg

                p_deg_l, y_deg_l = logits_to_deg(p_l, y_l)
                p_deg_r, y_deg_r = logits_to_deg(p_r, y_r)
                
                p_final = (p_deg_l + p_deg_r) / 2
                y_final = (y_deg_l + y_deg_r) / 2
                
                # Apply Coordinate Correction (if needed)
                if invert_pitch: p_final = -p_final
                if invert_yaw: y_final = -y_final
                
                if args.swap_axes:
                    p_final, y_final = y_final, p_final
                
                # Ground Truth to Degrees (robust_v16: 0=Pitch, 1=Yaw)
                gt_deg = gt * (180.0 / np.pi)
                pitch_gt = gt_deg[0]
                yaw_gt = gt_deg[1]
                
                # --- NEW: Standard 3D Angular Error Calculation ---
                def angles_to_unit_vector(pitch_deg, yaw_deg):
                    p = np.radians(pitch_deg)
                    y = np.radians(yaw_deg)
                    # Standard mapping: x=cos(p)sin(y), y=sin(p), z=cos(p)cos(y)
                    # Note: Coordinate system depends on dataset conventions, 
                    # but for angular distance, consistency is key.
                    vx = np.cos(p) * np.sin(y)
                    vy = np.sin(p)
                    vz = np.cos(p) * np.cos(y)
                    return np.array([vx, vy, vz])

                v_pred = angles_to_unit_vector(p_final.item(), y_final.item())
                v_gt = angles_to_unit_vector(pitch_gt.item(), yaw_gt.item())
                
                # Dot product for cosine similarity
                cos_sim = np.clip(np.dot(v_pred, v_gt), -1.0, 1.0)
                angular_error = np.degrees(np.arccos(cos_sim))
                
                # Update "All Cases"
                results['all']['error'] += angular_error
                results['all']['count'] += 1
                
                # Update "Frontal 45"
                if abs(yaw_gt.item()) <= 45.0:
                    results['frontal_45']['error'] += angular_error
                    results['frontal_45']['count'] += 1

    print(f"\n" + "="*45)
    print(f"{'SUBSET':<20} | {'SAMPLES':<10} | {'Ang Error (deg)':<10}")
    print(f"-"*45)
    
    for key, data in results.items():
        if data['count'] > 0:
            mae = data['error'] / data['count']
            name = "All Cases" if key == 'all' else "Frontal +/- 45"
            print(f"{name:<20} | {data['count']:<10} | {mae:.4f}")
    
    print(f"="*45)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='checkpoints/baseline_v16/best_student_p11.pt')
    parser.add_argument('--h5', type=str, default='data/processed/gaze360_robust_v16_test_B.h5')
    parser.add_argument('--invert_yaw', action='store_true', default=False)
    parser.add_argument('--invert_pitch', action='store_true', default=False)
    parser.add_argument('--swap_axes', action='store_true', default=False)
    args = parser.parse_args()
    
    evaluate_on_gaze360(
        model_path=args.model,
        h5_path=args.h5,
        invert_yaw=args.invert_yaw,
        invert_pitch=args.invert_pitch
    )