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

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

from src.models.student import LIPEV2StudentGold

def debug_gaze360(model_path, h5_path):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = LIPEV2StudentGold().to(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict, strict=False)
    model.eval()

    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

        target_indices = list(range(10)) + [1004]
        
        print(f"{'Sample':<6} | {'P_Pred':<8} | {'P_GT':<8} | {'Y_Pred':<8} | {'Y_GT':<8} | {'Error':<8}")
        print("-" * 65)

        total_err = 0
        with torch.no_grad():
            for i in target_indices:
                if i >= left_patches.shape[0]: continue
                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 = gaze_gt[i] * (180.0 / np.pi)

                out_l = model(lp, lm)
                out_r = model(rp, lm)
                
                def logits_to_deg(out):
                    idx = torch.arange(90).float().to(device)
                    p_deg = (torch.sum(torch.softmax(out[0], dim=1) * idx, dim=1) * 2 - 90)
                    y_deg = (torch.sum(torch.softmax(out[1], dim=1) * idx, dim=1) * 2 - 90)
                    return p_deg.item(), y_deg.item()

                p_l, y_l = logits_to_deg(out_l)
                p_r, y_r = logits_to_deg(out_r)
                
                p_pred = (p_l + p_r) / 2
                y_pred = (y_l + y_r) / 2
                
                p_gt, y_gt = gt[0], gt[1]
                err = (abs(p_pred - p_gt) + abs(y_pred - y_gt)) / 2
                total_err += err
                
                print(f"{i:<6} | {p_pred:8.2f} | {p_gt:8.2f} | {y_pred:8.2f} | {y_gt:8.2f} | {err:8.2f}")

        print("-" * 65)
        print(f"Average of target samples: {total_err/len(target_indices):.4f}")

if __name__ == "__main__":
    debug_gaze360('checkpoints/gold_swa/best_gold_p00.pt', 'data/processed/gaze360_robust_v16.h5')