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

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

from src.models.student import LIPEV2StudentGold, LIPEV2StudentBaseline
from src.utils.hardening import apply_phase1_hardening, preprocess_for_model

def logits_to_deg(p_logits, y_logits, device):
    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

def run_benchmark(subjects, gold_dir, v16_dir, data_dir, use_hardening=True):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    print(f"Hardening: {use_hardening}")
    
    results = {}

    for sub in subjects:
        print(f"\n>>> Benchmarking Subject: {sub}")
        
        # Load Models
        gold_path = os.path.join(gold_dir, f"best_gold_{sub}.pt")
        v16_path = os.path.join(v16_dir, f"best_student_{sub}.pt")
        
        if not os.path.exists(gold_path) or not os.path.exists(v16_path):
            print(f"Missing checkpoint for {sub}, skipping.")
            continue

        model_gold = LIPEV2StudentGold().to(device)
        state_gold = torch.load(gold_path, map_location=device)
        if 'n_averaged' in state_gold: # Handle SWA
            new_state = {}
            for k, v in state_gold.items():
                if k.startswith('module.'): new_state[k[7:]] = v
                else: new_state[k] = v
            state_gold = new_state
        model_gold.load_state_dict(state_gold, strict=False)
        model_gold.eval()

        model_v16 = LIPEV2StudentBaseline().to(device)
        state_v16 = torch.load(v16_path, map_location=device)
        model_v16.load_state_dict(state_v16)
        model_v16.eval()

        # Load Data
        h5_path = os.path.join(data_dir, f"{sub}_v16.h5") # Using v16 processed data
        if not os.path.exists(h5_path):
            print(f"Data {h5_path} not found.")
            continue

        sub_res = {
            'gold': {'err': 0.0, 'cnt': 0},
            'v16': {'err': 0.0, 'cnt': 0}
        }

        with h5py.File(h5_path, 'r') as f:
            lp = f['left_patches'][:]
            rp = f['right_patches'][:]
            lm = f['landmarks'][:]
            
            # Fallback for gaze labels
            if 'gaze' in f:
                gt = f['gaze'][:]
            elif 'left_gaze' in f:
                gt = f['left_gaze'][:]
            else:
                print(f"Error: No gaze labels in {h5_path}")
                continue
            
            num_samples = len(lp)
            for i in tqdm(range(num_samples), desc=f"Testing {sub}"):
                # 1. Prepare Inputs
                l_patch = lp[i]
                r_patch = rp[i]
                lmark = lm[i]
                target = torch.from_numpy(gt[i]).to(device) * (180.0 / np.pi) # To Degrees
                
                # 2. Apply Hardening if requested
                if use_hardening:
                    l_patch, lmark = apply_phase1_hardening(l_patch, lmark)
                    r_patch, _ = apply_phase1_hardening(r_patch, lmark) # Keep same jitter
                
                # 3. Final Preprocessing
                lp_t, lm_t = preprocess_for_model(l_patch, lmark)
                rp_t, _ = preprocess_for_model(r_patch, lmark)
                
                lp_t, rp_t, lm_t = lp_t.unsqueeze(0).to(device), rp_t.unsqueeze(0).to(device), lm_t.unsqueeze(0).to(device)

                # 4. Inference Gold
                with torch.no_grad():
                    p_gl, y_gl = model_gold(lp_t, lm_t)
                    p_gr, y_gr = model_gold(rp_t, lm_t)
                    pgl_d, ygl_d = logits_to_deg(p_gl, y_gl, device)
                    pgr_d, ygr_d = logits_to_deg(p_gr, y_gr, device)
                    pg_f = (pgl_d + pgr_d) / 2
                    yg_f = (ygl_d + ygr_d) / 2
                    
                    err_g = (torch.abs(pg_f - target[1]) + torch.abs(yg_f - target[0])).item()
                    sub_res['gold']['err'] += err_g
                    sub_res['gold']['cnt'] += 1

                # 5. Inference V16
                with torch.no_grad():
                    p_vl, y_vl = model_v16(lp_t, lm_t, state='A')
                    p_vr, y_vr = model_v16(rp_t, lm_t, state='A')
                    pvl_d, yvl_d = logits_to_deg(p_vl, y_vl, device)
                    pvr_d, yvr_d = logits_to_deg(p_vr, y_vr, device)
                    pv_f = (pvl_d + pvr_d) / 2
                    yv_f = (yvl_d + yvr_d) / 2
                    
                    err_v = (torch.abs(pv_f - target[1]) + torch.abs(yv_f - target[0])).item()
                    sub_res['v16']['err'] += err_v
                    sub_res['v16']['cnt'] += 1

        results[sub] = sub_res
        mae_g = sub_res['gold']['err'] / (sub_res['gold']['cnt'] * 2)
        mae_v = sub_res['v16']['err'] / (sub_res['v16']['cnt'] * 2)
        print(f"Subject {sub} | Gold MAE: {mae_g:.2f}° | V16 MAE: {mae_v:.2f}°")

    # Final Summary
    print("\n" + "="*50)
    print(f"{'Subject':<10} | {'Gold MAE':<15} | {'V16 MAE':<15} | {'Improv':<10}")
    print("-"*50)
    total_g_err, total_g_cnt = 0, 0
    total_v_err, total_v_cnt = 0, 0
    
    for sub in subjects:
        if sub in results:
            r = results[sub]
            mg = r['gold']['err'] / (r['gold']['cnt'] * 2)
            mv = r['v16']['err'] / (r['v16']['cnt'] * 2)
            imp = ((mv - mg) / mv) * 100 if mv > 0 else 0
            print(f"{sub:<10} | {mg:>13.2f}° | {mv:>13.2f}° | {imp:>8.1f}%")
            total_g_err += r['gold']['err']
            total_g_cnt += r['gold']['cnt']
            total_v_err += r['v16']['err']
            total_v_cnt += r['v16']['cnt']
            
    avg_g = total_g_err / (total_g_cnt * 2)
    avg_v = total_v_err / (total_v_cnt * 2)
    avg_imp = ((avg_v - avg_g) / avg_v) * 100
    print("-"*50)
    print(f"{'AVERAGE':<10} | {avg_g:>13.2f}° | {avg_v:>13.2f}° | {avg_imp:>8.1f}%")
    print("="*50)

if __name__ == '__main__':
    subjects = ['p00', 'p01', 'p02', 'p08', 'p11']
    gold_dir = 'checkpoints/gold_swa'
    v16_dir = 'checkpoints/baseline_v16'
    data_dir = 'data/processed'
    
    run_benchmark(subjects, gold_dir, v16_dir, data_dir, use_hardening=True)