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

# Thêm src vào path
sys.path.append(str(Path(__file__).parent))
from src.models.student import LIPEV2Student

def run_demo(h5_path):
    print(f"--- ĐANG CHẠY DEMO TÍCH HỢP VỚI FILE: {os.path.basename(h5_path)} ---")
    
    if not os.path.exists(h5_path):
        print("Lỗi: Không tìm thấy file dữ liệu.")
        return

    # 1. Khởi tạo Model Student (Phase 2)
    model = LIPEV2Student()
    model.eval()
    
    # 2. Load dữ liệu thực tế (Phase 1)
    with h5py.File(h5_path, 'r') as f:
        # Lấy mẫu số 100 để minh họa
        idx = 100
        patch = torch.from_numpy(f['left_patches'][idx]).float() / 255.0
        landmark = torch.from_numpy(f['landmarks'][idx]).float().view(-1)
        gaze_gt = f['left_gaze'][idx]
        
        # Thêm batch dimension (N=1)
        patch = patch.unsqueeze(0) # (1, 4, 8, 8)
        landmark = landmark.unsqueeze(0) # (1, 956)
        
    print(f"Dữ liệu đầu vào: Patch shape {patch.shape}, Landmark shape {landmark.shape}")
    print(f"Nhãn Gaze thực tế (Ground Truth): Pitch={gaze_gt[0]:.4f}, Yaw={gaze_gt[1]:.4f}")

    # 3. Chạy Inference qua Model
    with torch.no_grad():
        # Thử nghiệm State A (Full)
        prediction = model(patch, landmark, state='A')
        
    print(f"Kết quả dự đoán từ Model: Pitch={prediction[0,0]:.4f}, Yaw={prediction[0,1]:.4f}")
    print("(Lưu ý: Kết quả dự đoán hiện tại là ngẫu nhiên vì chưa qua bước Trưng cất Knowledge Distillation)")

    # 4. Lưu ảnh minh họa patch thực tế
    lp = patch[0].numpy() * 255
    l_combined = np.zeros((16, 16), dtype='uint8')
    l_combined[0:8, 0:8] = lp[0]
    l_combined[0:8, 8:16] = lp[1]
    l_combined[8:16, 0:8] = lp[2]
    l_combined[8:16, 8:16] = lp[3]
    l_view = cv2.resize(l_combined, (128, 128), interpolation=cv2.INTER_NEAREST)
    
    output_path = "data/verification/demo_patch_p03.png"
    cv2.imwrite(output_path, l_view)
    print(f"--- Đã lưu ảnh patch thực tế dùng trong demo tại: {output_path} ---")

if __name__ == '__main__':
    # Sử dụng p03.h5 vì nó đã hoàn thành và có kích thước lớn
    run_demo('data/processed/p03.h5')