import torch
import torch.nn.functional as F
from utils.MultiResolutionSTFTLoss import MultiResolutionSTFTLoss

# Keep STFT settings as is
stft_loss_fn = MultiResolutionSTFTLoss(
    fft_sizes=[512, 1024, 2048],
    hop_sizes=[64, 128, 256],
    win_lengths=[256, 512, 1024]
)


def feature_matching_loss(fmap_r, fmap_g):
    """
    Computes L1 distance between real and fake feature maps.
    """
    loss = 0
    for dr, dg in zip(fmap_r, fmap_g):
        for rl, gl in zip(dr, dg):
            rl = rl.detach()
            loss += torch.mean(torch.abs(rl - gl))

    return loss * 2


def discriminator_loss(disc_real_outputs, disc_generated_outputs):
    """
    Least Squares GAN Loss (LSGAN) for the Discriminator.
    Objective: Real -> 1, Fake -> 0
    """
    loss = 0
    r_losses = []
    g_losses = []

    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
        r_loss = torch.mean((dr - 1) ** 2)
        g_loss = torch.mean(dg ** 2)

        loss += (r_loss + g_loss)
        r_losses.append(r_loss.item())
        g_losses.append(g_loss.item())

    return loss, r_losses, g_losses


def generator_adv_loss(disc_generated_outputs):
    """
    Least Squares GAN Loss for the Generator.
    Objective: Fake -> 1 (Fool the discriminator)
    """
    loss = 0
    for dg in zip(disc_generated_outputs):
        dg = dg[0] # Unpack tuple
        loss += torch.mean((dg - 1) ** 2)
    return loss


def discriminator_train(
    high_quality,
    discriminator,
    generator_output
):
    y_d_rs, y_d_gs, _, _ = discriminator(high_quality, generator_output.detach())

    d_loss, _, _ = discriminator_loss(y_d_rs, y_d_gs)

    return d_loss


def generator_train(
    low_quality,
    high_quality,
    generator,
    discriminator,
    generator_output
):
    y_d_rs, y_d_gs, fmap_rs, fmap_gs = discriminator(high_quality, generator_output)

    loss_gen_adv = generator_adv_loss(y_d_gs)

    loss_fm = feature_matching_loss(fmap_rs, fmap_gs)

    stft_loss = stft_loss_fn(high_quality, generator_output)["total"]

    lambda_stft = 45.0
    lambda_fm = 2.0
    lambda_adv = 1.0

    combined_loss = (lambda_stft * stft_loss) + \
                    (lambda_fm * loss_fm) + \
                    (lambda_adv * loss_gen_adv)

    return combined_loss, loss_gen_adv