import torch
import torch.optim as optim
import torch.nn as nn
from model import RouterNetwork

def train_sft(model, embedder, dataset, epochs = 5, lr = 1e-3):
    """
    dataset: List of dicts [{'question': '...', 'label': 0|1|2}]
    """
    device = next(model.parameters()).device
    class_weights = torch.tensor([1.0, 1.0, 10.0], device=device, dtype=torch.float)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss(weight=class_weights)
    
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        correct_predictions = 0
        for data in dataset:
            # 1. Embed the query (state)
            # Using torch.no_grad() for embedder so we only train the router MLP
            with torch.no_grad():
                query_emb = embedder.encode(data['question'], convert_to_tensor=True)
                
            label = torch.tensor([data['label']], dtype=torch.long).to(query_emb.device)
            
            # 2. Forward pass through the router network
            optimizer.zero_grad()
            action_logits, _ = model(query_emb) # only care about action_logits (Actor head) for SFT

            # Calculate prediction and track accuracy for this epoch
            predicted_class = torch.argmax(action_logits, dim=-1)
            if predicted_class.item() == label.item():
                correct_predictions += 1
            
            # 3. Compute loss and backpropagate
            # CrossEntropyLoss expects (N, C) for logits and (N) for labels, so we unsqueeze action_logits to add batch dimension
            loss = criterion(action_logits.unsqueeze(0), label)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            
        avg_loss = total_loss / len(dataset)
        accuracy = (correct_predictions / len(dataset)) * 100
        print(f"SFT Epoch {epoch+1} | Loss: {avg_loss:.4f} | Accuracy: {accuracy:.2f}%")