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

# Define the Baseline Architecture (Addition Fusion)
class MiniConvEmbedder(nn.Module):
    def __init__(self):
        super(MiniConvEmbedder, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=0)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=0)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=0)
        self.relu = nn.ReLU(inplace=True)
        self.gap = nn.AdaptiveAvgPool2d(1)
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.gap(x)
        return torch.flatten(x, 1)

class LIPEV2StudentBaseline(nn.Module):
    def __init__(self):
        super(LIPEV2StudentBaseline, self).__init__()
        self.appearance_net = MiniConvEmbedder()
        self.geo_mlp = nn.Sequential(
            nn.Linear(956, 256),
            nn.LayerNorm(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.05),
            nn.Linear(256, 256),
            nn.ReLU(inplace=True)
        )
        self.pitch_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(inplace=True), nn.Linear(64, 90))
        self.yaw_head = nn.Sequential(nn.Linear(256, 64), nn.ReLU(inplace=True), nn.Linear(64, 90))

    def forward(self, patches=None, landmarks=None, state='A'):
        geo_feat = self.geo_mlp(landmarks)
        if state == 'A' and patches is not None:
            batch_size = patches.shape[0]
            patches = patches.view(-1, 1, patches.shape[2], patches.shape[3])
            app_tokens = self.appearance_net(patches)
            app_feat = app_tokens.view(batch_size, -1)
            combined = app_feat + geo_feat # Baseline used Addition
        else:
            combined = geo_feat
        return self.pitch_head(combined), self.yaw_head(combined)

def evaluate_baseline_on_gaze360(model_path, h5_path):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Evaluating BASELINE model: {model_path}")
    
    model = LIPEV2StudentBaseline().to(device)
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.eval()

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

    with h5py.File(h5_path, 'r') as f:
        lp, rp, lm, g_gt = f['left_patches'][:], f['right_patches'][:], f['landmarks'][:], f['gaze'][:]
        with torch.no_grad():
            for i in tqdm(range(len(lp)), desc="Testing Baseline"):
                p_l, y_l = model(torch.from_numpy(lp[i]).float().unsqueeze(0).to(device), torch.from_numpy(lm[i]).float().view(1,-1).to(device), state='A')
                p_r, y_r = model(torch.from_numpy(rp[i]).float().unsqueeze(0).to(device), torch.from_numpy(lm[i]).float().view(1,-1).to(device), state='A')
                
                def l2d(p, y):
                    idx = torch.arange(90).float().to(device)
                    pp, yp = torch.softmax(p, 1), torch.softmax(y, 1)
                    return (torch.sum(pp*idx,1)*2-90), (torch.sum(yp*idx,1)*2-90)
                
                pl, yl = l2d(p_l, y_l)
                pr, yr = l2d(p_r, y_r)
                pf, yf = (pl+pr)/2, (yl+yr)/2
                
                gt_d = torch.from_numpy(g_gt[i]).to(device) * (180.0/np.pi)
                yaw_gt_deg = gt_d[1].item()
                
                error = (torch.abs(pf-gt_d[0]) + torch.abs(yf-gt_d[1])).item()
                
                results['all']['error'] += error
                results['all']['count'] += 1
                if abs(yaw_gt_deg) <= 45.0:
                    results['frontal_45']['error'] += error
                    results['frontal_45']['count'] += 1

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

if __name__ == "__main__":
    # Đánh giá model p08 từ baseline_v16 (Addition Fusion) tốt nhất
    evaluate_baseline_on_gaze360('checkpoints/baseline_v16/best_student_p08.pt', 'data/processed/gaze360_robust_v16_test_B.h5')