import torch
from tqdm import tqdm
import numpy as np

class Trainer:
    def __init__(self, model, train_loader, val_loader, criterion, optimizer, device):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.criterion = criterion # 우리가 만든 TrustWeightedLoss
        self.optimizer = optimizer
        self.device = device

    def train(self, epochs):
        print(f"🚀 Training started on {self.device}")
        
        for epoch in range(epochs):
            self.model.train()
            total_loss = 0
            
            loop = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{epochs}")
            for batch in loop:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                ratings = batch['rating'].to(self.device)
                weights = batch['weight'].to(self.device)

                self.optimizer.zero_grad()
                outputs = self.model(input_ids, attention_mask).squeeze()
                
                # 가중 손실 함수 적용
                loss = self.criterion(outputs, ratings, weights)
                
                loss.backward()
                self.optimizer.step()
                
                total_loss += loss.item()
                loop.set_postfix(loss=loss.item())
            
            print(f"Epoch {epoch+1} Avg Loss: {total_loss / len(self.train_loader):.4f}")
            self.evaluate()

    def evaluate(self):
        self.model.eval()
        val_loss = 0
        with torch.no_grad():
            for batch in self.val_loader:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                ratings = batch['rating'].to(self.device)
                weights = batch['weight'].to(self.device)
                
                outputs = self.model(input_ids, attention_mask).squeeze()
                loss = self.criterion(outputs, ratings, weights)
                val_loss += loss.item()
        
        print(f"📊 Validation Loss: {val_loss / len(self.val_loader):.4f}")