# =================================================================================
# UKÁZKA PSEUDOKÓDU V KNIHOVNĚ PYTORCH:
# =================================================================================

import torch
import torch.nn as nn

# ----------------------------------------------------------------------
# INICIALIZACE (provádí se jednou před zahájením trénovacího cyklu)
# ----------------------------------------------------------------------
# Předpokládá se existence modelů: generator a discriminator
loss_function = nn.BCELoss()

# Dva oddělené optimizéry pro každou síť zvlášť
d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=0.0002)
g_optimizer = torch.optim.Adam(generator.parameters(), lr=0.0002)


# ----------------------------------------------------------------------
# TĚLO TRÉNOVACÍ SMYČKY (jeden krok pro jednu dávku dat)
# ----------------------------------------------------------------------

# --------------------------------------------------
# FÁZE 1: Trénování diskriminátoru
# --------------------------------------------------
d_optimizer.zero_grad()

# 1. Trénování na reálných datech (cíl = 1)
real_labels = torch.ones(batch_size, 1)
output_real = discriminator(real_images)
d_loss_real = loss_function(output_real, real_labels)

# 2. Trénování na falešných datech (cíl = 0)
noise = torch.randn(batch_size, latent_dim)
fake_images = generator(noise)
fake_labels = torch.zeros(batch_size, 1)

# Použije se .detach(), aby se gradient v této fázi nešířil do generátoru.
# Tenzor je odpojen od dynamického výpočetního grafu (requires_grad=False).
# Důvody:
# A) Zamezení zbytečným výpočtům v generátoru (šetří se čas a VRAM).
# B) Ochrana gradientů generátoru: v této fázi jsou optimalizovány výhradně váhy diskriminátoru.
# Tenzor fake_images zůstává v paměti a ve Fázi 2 je odeslán znovu bez .detach().
output_fake = discriminator(fake_images.detach())
d_loss_fake = loss_function(output_fake, fake_labels)

# Sečtení chyb a aktualizace POUZE diskriminátoru
d_loss = d_loss_real + d_loss_fake
d_loss.backward()
d_optimizer.step()


# --------------------------------------------------
# FÁZE 2: Trénování generátoru
# --------------------------------------------------
g_optimizer.zero_grad()

# Falešná data jsou znovu ohodnocena diskriminátorem
output_fake_for_g = discriminator(fake_images)

# Non-saturating trik: Cíl je nastaven na 1 (snaha oklamat diskriminátor)
trick_labels = torch.ones(batch_size, 1)
g_loss = loss_function(output_fake_for_g, trick_labels)

# Výpočet gradientů a aktualizace POUZE generátoru
g_loss.backward()
g_optimizer.step()