# GENAI Experiment 7: CycleGAN unpaired style transfer
# ===========================================================================

# # Part 1 — Create the unpaired dataset (two domains)

import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import matplotlib.pyplot as plt
torch.manual_seed(0); np.random.seed(0)
S = 16
yy, xx = np.mgrid[0:S, 0:S]

def render(shape, cx, cy, r, style):
    def mask(rr):
        if shape == 0: return (xx - cx)**2 + (yy - cy)**2 <= rr*rr
        return (abs(xx - cx) <= 0.85*rr) & (abs(yy - cy) <= 0.85*rr)
    img = mask(r).astype(np.float32)
    return img if style == "light" else 1.0 - img                       # "dark" style = colours inverted

def make_domain(style, n, seed):
    rng = np.random.default_rng(seed); params = []
    for _ in range(n):
        r = rng.uniform(4, 6.5); params.append((rng.integers(2), *rng.uniform(r, S - r, 2), r))
    return np.stack([render(*p, style) for p in params]), params

XA, _ = make_domain("light", 2000, 1)      # domain A
XB, _ = make_domain("dark", 2000, 2)       # domain B (different random shapes -> unpaired)
to_t = lambda X: torch.tensor(X).unsqueeze(1) * 1.6 - 0.8      # 1x16x16, range [-0.8, 0.8] (keeps tanh away from saturation)
A, B = to_t(XA), to_t(XB)

fig, ax = plt.subplots(2, 8, figsize=(10, 2.6))
for i in range(8):
    ax[0, i].imshow(XA[i], cmap="gray"); ax[0, i].axis("off"); ax[1, i].imshow(XB[i], cmap="gray"); ax[1, i].axis("off")
ax[0, 0].set_title("Domain A: light shapes", loc="left", fontsize=9); ax[1, 0].set_title("Domain B: dark shapes (unpaired)", loc="left", fontsize=9); plt.show()

# # Part 2 — Two generators and two discriminators

def Generator():                     # small CNN: image in -> image out (same size)
    return nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(),
                         nn.Conv2d(16, 16, 3, padding=1), nn.ReLU(),
                         nn.Conv2d(16, 1, 3, padding=1), nn.Tanh())

def Discriminator():                 # PatchGAN-style: outputs a 4x4 grid of real/fake scores
    return nn.Sequential(nn.Conv2d(1, 16, 4, 2, 1), nn.LeakyReLU(0.2),
                         nn.Conv2d(16, 32, 4, 2, 1), nn.LeakyReLU(0.2),
                         nn.Conv2d(32, 1, 3, 1, 1))

G_AB, G_BA, D_A, D_B = Generator(), Generator(), Discriminator(), Discriminator()
opt_G = torch.optim.Adam(list(G_AB.parameters()) + list(G_BA.parameters()), lr=2e-4, betas=(0.5, 0.999))
opt_D = torch.optim.Adam(list(D_A.parameters()) + list(D_B.parameters()), lr=2e-4, betas=(0.5, 0.999))
mse, l1 = nn.MSELoss(), nn.L1Loss()
LAMBDA_CYC, LAMBDA_ID = 1, 0      # lab uses 10 for real photos; see the note below. Identity loss off (wrong target for an inversion)

# # Part 3 — Training: adversarial + cycle-consistency (+ identity) losses
# Note on λ_cycle (tested): with these tiny networks and simple images, λ_cycle = 10 lets the cycle
# loss dominate. The generators then find a shortcut: they hide enough information to reconstruct the
# input, but do not really translate the style (A→B error stayed ≈ 0.63). With λ_cycle = 1 the
# adversarial loss gets enough weight and both directions translate correctly (error ≈ 0.01). If the
# exam fixes λ = 10, train longer or give the discriminators a lower learning rate.

hist = {"G": [], "cycle": [], "D": []}
fixed_a, fixed_b = A[:6], B[:6]
for step in range(1, 1501):
    a = A[torch.randint(0, len(A), (32,))]; b = B[torch.randint(0, len(B), (32,))]   # random, unpaired batches
    # ----- Generators -----
    fake_b, fake_a = G_AB(a), G_BA(b)
    pb, pa = D_B(fake_b), D_A(fake_a)
    loss_gan = mse(pb, torch.ones_like(pb)) + mse(pa, torch.ones_like(pa))
    loss_cyc = l1(G_BA(fake_b), a) + l1(G_AB(fake_a), b)               # A->B->A and B->A->B
    loss_id = l1(G_AB(b), b) + l1(G_BA(a), a)
    loss_G = loss_gan + LAMBDA_CYC * loss_cyc + LAMBDA_ID * loss_id
    opt_G.zero_grad(); loss_G.backward(); opt_G.step()
    # ----- Discriminators -----
    ra, fa_ = D_A(a), D_A(fake_a.detach()); rb, fb_ = D_B(b), D_B(fake_b.detach())
    loss_D = 0.5 * (mse(ra, torch.ones_like(ra)) + mse(fa_, torch.zeros_like(fa_))
                    + mse(rb, torch.ones_like(rb)) + mse(fb_, torch.zeros_like(fb_)))
    opt_D.zero_grad(); loss_D.backward(); opt_D.step()
    hist["G"].append(loss_G.item()); hist["cycle"].append(loss_cyc.item()); hist["D"].append(loss_D.item())
    if step % 300 == 0:
        print(f"step {step:4d}  G {np.mean(hist['G'][-300:]):.3f}  cycle {np.mean(hist['cycle'][-300:]):.3f}  D {np.mean(hist['D'][-300:]):.3f}")

# # Part 4 — Loss curves

sm = lambda v: np.convolve(v, np.ones(50) / 50, "valid")
fig, ax = plt.subplots(1, 3, figsize=(14, 3.3))
for a_, k in zip(ax, ["G", "cycle", "D"]): a_.plot(sm(hist[k])); a_.set_title(f"{k} loss"); a_.set_xlabel("step")
plt.tight_layout(); plt.show()

# # Part 5 — Translate UNSEEN images: A → B and B → A, and check the result

TA, PA = make_domain("light", 6, 77); TB, PB = make_domain("dark", 6, 78)
ta, tb = to_t(TA), to_t(TB)
with torch.no_grad():
    ab, aba = G_AB(ta), G_BA(G_AB(ta)); ba = G_BA(tb)
to_img = lambda t: ((t.squeeze(1) + 0.8) / 1.6).clamp(0, 1).numpy()
rows = [(TA, "A (unseen)"), (to_img(ab), "G_AB(A): A -> B"), (to_img(aba), "G_BA(G_AB(A)): cycle"), (TB, "B (unseen)"), (to_img(ba), "G_BA(B): B -> A")]
fig, ax = plt.subplots(len(rows), 6, figsize=(8, 7))
for r, (imgs, t) in enumerate(rows):
    for i in range(6): ax[r, i].imshow(imgs[i], cmap="gray", vmin=0, vmax=1); ax[r, i].axis("off")
    ax[r, 0].set_title(t, loc="left", fontsize=8)
plt.tight_layout(); plt.show()

# Because the data is synthetic we know the correct answer for each test image:
true_B = np.stack([render(*p, "dark") for p in PA]); true_A = np.stack([render(*p, "light") for p in PB])
print(f"A->B error vs true dark image  : {np.abs(to_img(ab) - true_B).mean():.3f}   (input A vs truth: {np.abs(TA - true_B).mean():.3f})")
print(f"B->A error vs true light image  : {np.abs(to_img(ba) - true_A).mean():.3f}   (input B vs truth: {np.abs(TB - true_A).mean():.3f})")
print(f"Cycle reconstruction error A->B->A: {np.abs(to_img(aba) - TA).mean():.3f}")

# Exam explanation: the discriminators teach each generator what the other domain looks like
# (adversarial loss). The cycle loss forces A → B → A to give back the input, so the *shape, size and
# position* are kept and only the *style* changes (light ↔ dark), exactly like horse ↔ zebra or summer
# ↔ winter. The cycle loss falls during training, and the translated test images are much closer to
# the true answer than the inputs were, even though no paired examples were ever shown.
