"""Stage 3 training: Integration (Mixer + Post-processing)."""

import torch
from typing import Dict

from saltts.training.trainer_base import BaseTrainer
from saltts.core.mixer import Mixer
from saltts.postprocessing.pipeline import PostProcessingPipeline


class Stage3Trainer(BaseTrainer):
    """Trainer for Stage 3: Integration."""

    def __init__(
        self,
        freeze_generators: bool = True,
        device: str = "cuda",
        **kwargs,
    ):
        """Initialize Stage 3 trainer.

        Args:
            freeze_generators: Whether to freeze generator weights
            device: Device
            **kwargs: Additional arguments for BaseTrainer
        """
        # Initialize Mixer
        model = Mixer()

        # Optimizer
        optimizer = torch.optim.Adam(model.parameters(), lr=0.00005)

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

        # Post-processing pipeline
        self.postproc = PostProcessingPipeline()

    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
            cfg = batch["cfg"].to(self.device)
            side_feat = batch["side_feat"].to(self.device)
            aux_params = batch["aux_params"].to(self.device)
            target = batch["target"].to(self.device)

            # Forward pass
            self.optimizer.zero_grad()
            mixed = self.model(cfg, side_feat, aux_params)

            # Compute loss
            loss = torch.nn.functional.mse_loss(mixed, target)

            loss.backward()
            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:
                cfg = batch["cfg"].to(self.device)
                side_feat = batch["side_feat"].to(self.device)
                aux_params = batch["aux_params"].to(self.device)
                target = batch["target"].to(self.device)

                mixed = self.model(cfg, side_feat, aux_params)
                loss = torch.nn.functional.mse_loss(mixed, target)

                total_loss += loss.item()

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