"""Stage 2 training: Stream Transformer."""

import torch
from typing import Dict

from saltts.training.trainer_base import BaseTrainer
from saltts.training.losses import ProsodyLoss
from saltts.core.stream_transformer import StreamTransformer


class Stage2Trainer(BaseTrainer):
    """Trainer for Stage 2: Stream Transformer."""

    def __init__(self, device: str = "cuda", **kwargs):
        """Initialize Stage 2 trainer.

        Args:
            device: Device
            **kwargs: Additional arguments for BaseTrainer
        """
        # Initialize Stream Transformer
        model = StreamTransformer()

        # Optimizer
        optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001, weight_decay=0.01)

        super().__init__(model, optimizer, device, **kwargs)

        # Loss function
        self.prosody_loss = ProsodyLoss()

    def train_epoch(self, dataloader) -> Dict[str, float]:
        """Train for one epoch."""
        self.model.train()
        total_loss = 0.0

        for batch in dataloader:
            # Move to device
            text1 = batch["text"].to(self.device)
            target_prosody = batch["prosody"].to(self.device)

            # Forward pass
            self.optimizer.zero_grad()
            pred_prosody = self.model(text1)

            # Compute loss (simplified)
            loss = torch.nn.functional.mse_loss(pred_prosody, target_prosody)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()

            total_loss += loss.item()
            self.global_step += 1

        return {"loss": total_loss / len(dataloader)}

    def validate(self, dataloader) -> Dict[str, float]:
        """Validate model."""
        self.model.eval()
        total_loss = 0.0

        with torch.no_grad():
            for batch in dataloader:
                text1 = batch["text"].to(self.device)
                target_prosody = batch["prosody"].to(self.device)

                pred_prosody = self.model(text1)
                loss = torch.nn.functional.mse_loss(pred_prosody, target_prosody)

                total_loss += loss.item()

        return {"val_loss": total_loss / len(dataloader)}
