"""Stage 1 training: Generator A and B."""

import torch
import torch.nn as nn
from typing import Dict

from saltts.training.trainer_base import BaseTrainer
from saltts.training.losses import PerceptualLoss, FlowNLLLoss, AdversarialLoss
from saltts.models.generator_a.flow import FlowGenerator
from saltts.models.generator_a.diffusion import DiffusionGenerator
from saltts.models.generator_a.gan import GANGenerator
from saltts.models.generator_b.parametric import ParametricGenerator


class Stage1Trainer(BaseTrainer):
    """Trainer for Stage 1: Generator A and B."""

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

        Args:
            generator_a_type: Type of Generator A (flow/diffusion/gan)
            device: Device
            **kwargs: Additional arguments for BaseTrainer
        """
        # Initialize Generator A based on type
        if generator_a_type == "flow":
            generator_a = FlowGenerator()
        elif generator_a_type == "diffusion":
            generator_a = DiffusionGenerator()
        elif generator_a_type == "gan":
            generator_a = GANGenerator()
        else:
            raise ValueError(f"Unknown generator type: {generator_a_type}")

        # Initialize Generator B
        generator_b = ParametricGenerator()

        # Combine models
        model = nn.ModuleDict({
            "generator_a": generator_a,
            "generator_b": generator_b,
        })

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

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

        # Loss functions
        self.perceptual_loss = PerceptualLoss()

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

            # Forward pass
            self.optimizer.zero_grad()

            # Train Generator A
            pred_a = self.model["generator_a"](prosody, None)
            loss_a = self.perceptual_loss(pred_a, target)

            # Train Generator B
            audio_b = self.model["generator_b"].synthesize(prosody)
            # Convert to tensor for loss computation
            # loss_b = ...

            # Total loss
            loss = loss_a  # + loss_b
            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:
                prosody = batch["prosody"].to(self.device)
                target = batch["audio"].to(self.device)

                pred_a = self.model["generator_a"](prosody, None)
                loss = self.perceptual_loss(pred_a, target)

                total_loss += loss.item()

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