import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, ConcatDataset
import os
import time
import numpy as np
from tqdm import tqdm
import sys
from pathlib import Path

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

from src.models.student import LIPEV2Student
from src.models.loss import GazeDistillationLoss
from src.data.dataset import GazeDataset

def get_dann_alpha(epoch, max_epochs):
    p = float(epoch) / max_epochs
    alpha = 2. / (1. + np.exp(-10 * p)) - 1
    return alpha

def train_dann_epoch(model, loader, criterion, domain_criterion, optimizer, device, epoch, max_epochs, w_aw, w_kd, w_domain=0.1):
    model.train()
    running_loss = 0.0
    running_aw = 0.0
    running_domain = 0.0
    
    alpha = get_dann_alpha(epoch, max_epochs)
    
    pbar = tqdm(loader, desc=f"Epoch {epoch} [DANN + AdaLN, alpha={alpha:.2f}]")
    for batch in pbar:
        patches, landmarks, gaze, t_p_logits, t_yaw_logits, domains = [b.to(device) for b in batch]
        teacher_logits = (t_p_logits, t_yaw_logits)
            
        optimizer.zero_grad()
        
        # Forward pass (DANN + AdaLN)
        s_p_logits, s_y_logits, domain_logits = model(patches, landmarks, state='A', alpha=alpha, domain_id=domains)
        
        # 1. Gaze Loss
        criterion.w_aw = w_aw
        criterion.w_kd = w_kd
        gaze_loss, l_aw, l_kd = criterion((s_p_logits, s_y_logits), gaze, teacher_logits)
        
        # 2. Domain Loss
        d_loss = domain_criterion(domain_logits, domains)
        
        total_loss = gaze_loss + w_domain * d_loss
        
        total_loss.backward()
        optimizer.step()
        
        running_loss += total_loss.item()
        running_aw += l_aw.item() if isinstance(l_aw, torch.Tensor) else l_aw
        running_domain += d_loss.item()
        
        pbar.set_postfix({'loss': f"{total_loss.item():.4f}", 'aw': f"{l_aw:.4f}", 'dom': f"{d_loss.item():.4f}"})
        
    return running_loss / len(loader), running_aw / len(loader), running_domain / len(loader)

def train_dann(h5_dir, source_suffix='_v16_new.h5', target_file='gaze360_robust_v16_new.h5', num_epochs=10):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"EXPERIMENT 2: DANN + AdaLN (Clean Data) on {device}")
    
    all_files = os.listdir(h5_dir)
    source_files = [os.path.join(h5_dir, f) for f in all_files if f.endswith(source_suffix) and f.startswith('p')]
    target_path = os.path.join(h5_dir, target_file)
    
    if not source_files:
        print(f"Waiting for source files with suffix {source_suffix}...")
        return

    source_ds = GazeDataset(source_files, transform=True, domain_id=0)
    target_ds = GazeDataset([target_path], transform=True, domain_id=1)
    
    train_ds = ConcatDataset([source_ds, target_ds])
    train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4)
    
    model = LIPEV2Student().to(device)
    criterion = GazeDistillationLoss()
    domain_criterion = torch.nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)
    
    os.makedirs('checkpoints/dann_adaln', exist_ok=True)
    
    for epoch in range(1, num_epochs + 1):
        w_aw, w_kd = 1.0, 0.5 
        t_loss, l_aw, d_loss = train_dann_epoch(model, train_loader, criterion, domain_criterion, optimizer, device, epoch, num_epochs, w_aw, w_kd)
        print(f"Epoch {epoch}: Total {t_loss:.4f}, AW {l_aw:.4f}, Dom {d_loss:.4f}")
        
        if epoch == num_epochs:
            torch.save(model.state_dict(), f'checkpoints/dann_adaln/student_adaln_final.pt')

if __name__ == '__main__':
    train_dann('data/processed', num_epochs=5) # 5 epochs for quick check