"""GAN implementation for financial data generation."""

import torch
import torch.nn as nn
import torch.optim as optim
from typing import Dict, Tuple

class Generator(nn.Module):
    def __init__(self, latent_dim: int, feature_dim: int, hidden_dims: List[int]):
        super().__init__()
        
        layers = []
        prev_dim = latent_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.BatchNorm1d(hidden_dim),
                nn.LeakyReLU(0.2),
                nn.Dropout(0.3)
            ])
            prev_dim = hidden_dim
            
        layers.append(nn.Linear(prev_dim, feature_dim))
        layers.append(nn.Tanh())
        
        self.model = nn.Sequential(*layers)
        
    def forward(self, z: torch.Tensor) -> torch.Tensor:
        return self.model(z)

class Discriminator(nn.Module):
    def __init__(self, feature_dim: int, hidden_dims: List[int]):
        super().__init__()
        
        layers = []
        prev_dim = feature_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.LeakyReLU(0.2),
                nn.Dropout(0.3)
            ])
            prev_dim = hidden_dim
            
        layers.append(nn.Linear(prev_dim, 1))
        layers.append(nn.Sigmoid())
        
        self.model = nn.Sequential(*layers)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

class FinancialGAN:
    def __init__(self, config: Dict):
        """Initialize the GAN."""
        self.latent_dim = config['model']['latent_dim']
        self.feature_dim = config['model']['feature_dim']
        self.hidden_dims = config['model']['hidden_dims']
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        self.generator = Generator(
            self.latent_dim, 
            self.feature_dim, 
            self.hidden_dims
        ).to(self.device)
        
        self.discriminator = Discriminator(
            self.feature_dim, 
            self.hidden_dims[::-1]
        ).to(self.device)
        
        self.g_optimizer = optim.Adam(
            self.generator.parameters(), 
            lr=config['model']['learning_rate']
        )
        self.d_optimizer = optim.Adam(
            self.discriminator.parameters(), 
            lr=config['model']['learning_rate']
        )
        
        self.criterion = nn.BCELoss()
        
    def train_step(self, real_data: torch.Tensor) -> Tuple[float, float]:
        """Perform one training step."""
        batch_size = real_data.size(0)
        real_label = torch.ones(batch_size, 1).to(self.device)
        fake_label = torch.zeros(batch_size, 1).to(self.device)
        
        # Train Discriminator
        self.d_optimizer.zero_grad()
        d_real_output = self.discriminator(real_data)
        d_real_loss = self.criterion(d_real_output, real_label)
        
        z = torch.randn(batch_size, self.latent_dim).to(self.device)
        fake_data = self.generator(z)
        d_fake_output = self.discriminator(fake_data.detach())
        d_fake_loss = self.criterion(d_fake_output, fake_label)
        
        d_loss = d_real_loss + d_fake_loss
        d_loss.backward()
        self.d_optimizer.step()
        
        # Train Generator
        self.g_optimizer.zero_grad()
        g_output = self.discriminator(fake_data)
        g_loss = self.criterion(g_output, real_label)
        g_loss.backward()
        self.g_optimizer.step()
        
        return g_loss.item(), d_loss.item()
        
    def generate_samples(self, num_samples: int) -> torch.Tensor:
        """Generate synthetic financial data."""
        self.generator.eval()
        with torch.no_grad():
            z = torch.randn(num_samples, self.latent_dim).to(self.device)
            samples = self.generator(z)
        self.generator.train()
        return samples