# GENAI Experiment 5: Autoencoder DeepFake: shared encoder + two decoders, reconstruction analysis
# ===========================================================================

# # Part 1 — Create the two "people" (synthetic faces)
# * Person A: round face, small eyes, no eyebrows.
# * Person B: square face, big eyes, eyebrows and a beard.
# * Expression (smile ↔ frown) and pose (head turned left ↔ right) are random for every image.
#
# *If real face images are allowed:* load any two folders of face crops with PIL, convert to
# grayscale, resize to 24×24 and use them in place of XA, XB. Everything else stays the same.

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 = 24
yy, xx = np.mgrid[0:S, 0:S]

def face(person, expr, pose):
    """person 'A' or 'B'; expr in [-1, 1] (frown..smile); pose in [-3, 3] (head turn)."""
    img = np.zeros((S, S), np.float32)
    cx, cy = 12, 12
    if person == "A": skin = (xx - cx)**2 + (yy - cy)**2 <= 10**2                    # round face
    else:             skin = (abs(xx - cx) <= 9) & (abs(yy - cy) <= 10)              # square face
    img[skin] = 0.8
    fx = cx + pose                                                                    # features shift with pose
    er = 1.2 if person == "A" else 2.0                                               # eye size
    for ex in (fx - 4, fx + 4):
        img[(xx - ex)**2 + (yy - 9)**2 <= er**2] = 0.05
    if person == "B":
        img[(abs(yy - 6) <= 0.6) & (abs(xx - fx) <= 7) & (abs(xx - fx) >= 1.5)] = 0.05     # eyebrows
        img[skin & (yy >= 18)] = 0.35                                                    # beard
    mouth_y = 16 - expr * 0.12 * (xx - fx)**2 + expr * 1.5                               # smile / frown curve
    img[(abs(yy - mouth_y) <= 0.7) & (abs(xx - fx) <= 4)] = 0.05
    return img

def make_person(person, n, seed):
    rng = np.random.default_rng(seed)
    params = np.c_[rng.uniform(-1, 1, n), rng.uniform(-3, 3, n)]
    return np.stack([face(person, e, p) for e, p in params]), params

XA, PA = make_person("A", 1500, 1)
XB, PB = make_person("B", 1500, 2)
fig, ax = plt.subplots(2, 6, figsize=(9, 3.2))
for i in range(6):
    ax[0, i].imshow(XA[i], cmap="gray", vmin=0, vmax=1); ax[0, i].axis("off")
    ax[1, i].imshow(XB[i], cmap="gray", vmin=0, vmax=1); ax[1, i].axis("off")
ax[0, 0].set_title("Person A", loc="left", fontsize=9); ax[1, 0].set_title("Person B", loc="left", fontsize=9); plt.show()

# # Part 2 — Shared encoder + two decoders, trained only to reconstruct

D, L = S * S, 16
A = torch.tensor(XA).view(-1, D); B = torch.tensor(XB).view(-1, D)
encoder   = nn.Sequential(nn.Linear(D, 256), nn.ReLU(), nn.Linear(256, L))
decoder_A = nn.Sequential(nn.Linear(L, 256), nn.ReLU(), nn.Linear(256, D), nn.Sigmoid())
decoder_B = nn.Sequential(nn.Linear(L, 256), nn.ReLU(), nn.Linear(256, D), nn.Sigmoid())
opt = torch.optim.Adam(list(encoder.parameters()) + list(decoder_A.parameters()) + list(decoder_B.parameters()), lr=1e-3)

hist = []
for epoch in range(40):
    perm = torch.randperm(len(A)); tot = 0
    for i in range(0, len(A), 64):
        a, b = A[perm[i:i+64]], B[perm[i:i+64]]
        loss = F.mse_loss(decoder_A(encoder(a)), a) + F.mse_loss(decoder_B(encoder(b)), b)   # A->A and B->B only
        opt.zero_grad(); loss.backward(); opt.step(); tot += loss.item()
    hist.append(tot / (len(A) // 64 + 1))
    if epoch % 10 == 0 or epoch == 39: print(f"epoch {epoch:2d}  loss {hist[-1]:.5f}")
plt.plot(hist); plt.yscale("log"); plt.xlabel("epoch"); plt.ylabel("MSE"); plt.title("Training loss"); plt.show()

# # Part 3 — The swap: decoder_B(encoder(A)) and decoder_A(encoder(B))

TA, TPA = make_person("A", 6, 99)                     # NEW unseen images of A
TB, TPB = make_person("B", 6, 98)
ta, tb = torch.tensor(TA).view(-1, D), torch.tensor(TB).view(-1, D)
with torch.no_grad():
    A_rec, A_swap = decoder_A(encoder(ta)), decoder_B(encoder(ta))
    B_rec, B_swap = decoder_B(encoder(tb)), decoder_A(encoder(tb))

def rows(lst, titles):
    fig, ax = plt.subplots(len(lst), 6, figsize=(9, 1.6 * len(lst)))
    for r, (imgs, t) in enumerate(zip(lst, titles)):
        for i in range(6):
            ax[r, i].imshow(np.asarray(imgs[i]).reshape(S, S), 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()

rows([ta, A_rec, A_swap], ["A (input)", "Dec_A(Enc(A)) reconstruction", "Dec_B(Enc(A)) SWAP: A's expression, B's face"])
rows([tb, B_rec, B_swap], ["B (input)", "Dec_B(Enc(B)) reconstruction", "Dec_A(Enc(B)) SWAP: B's expression, A's face"])

# Check the swap with the true answer. The faces are synthetic, so B can be drawn with A's exact
# expression and pose and compared.

truth_B = np.stack([face("B", e, p) for e, p in TPA]).reshape(6, -1)
mse_swap = ((A_swap.numpy() - truth_B)**2).mean()
mse_input = ((ta.numpy() - truth_B)**2).mean()
print(f"MSE( swap , true B with A's expression/pose ) = {mse_swap:.4f}")
print(f"MSE( original A , same target )               = {mse_input:.4f}   -> the swap is much closer to the target")

# # Part 4 — Reconstruction-analysis system (input → preprocess → reconstruct → error → side by side)

def preprocess(img):
    """Any grayscale or RGB array -> grayscale -> 24x24 (block average / nearest) -> values 0..1."""
    img = np.asarray(img, np.float32)
    if img.ndim == 3: img = img.mean(2)
    if img.max() > 1: img = img / 255.0
    h, w = img.shape
    ys = (np.arange(S) * h / S).astype(int); xs = (np.arange(S) * w / S).astype(int)
    return img[np.ix_(ys, xs)]

def analyse(img, title, dec=decoder_A):
    x = torch.tensor(preprocess(img)).view(1, D)
    with torch.no_grad(): r = dec(encoder(x))
    err = (x - r).view(S, S).numpy(); mse = float((err**2).mean()); psnr = 10 * np.log10(1 / mse)
    fig, ax = plt.subplots(1, 3, figsize=(7, 2.5))
    for a, im, t in zip(ax, [x.view(S, S), r.view(S, S), np.abs(err)], ["original", "reconstructed", "|error|"]):
        a.imshow(im, cmap="gray" if t != "|error|" else "hot", vmin=0, vmax=1); a.set_title(t, fontsize=8); a.axis("off")
    plt.suptitle(f"{title}:  MSE = {mse:.4f}   PSNR = {psnr:.1f} dB", fontsize=9); plt.show()
    return mse

genuine = analyse(np.kron(TA[0], np.ones((4, 4))), "Genuine face of A (96x96 input, resized)")
tampered = TA[0].copy(); tampered[8:12, 4:20] = 0.8                       # eyes painted over (manipulation)
fake = analyse(tampered, "Manipulated face (eyes removed)")
noise = analyse(np.random.rand(S, S), "Non-face input (random noise)")
print("Higher reconstruction error = image does not look like the faces the model learned -> possible manipulation.")

# Exam explanation: the autoencoder only learned faces of A and B, so genuine faces get a low error.
# Manipulated or non-face inputs get a higher error, which flags them. In the swap, the shared encoder
# carries expression and pose and the decoder decides identity. That is why Dec_B(Enc(A)) shows B
# smiling or frowning exactly like A: the core idea of autoencoder DeepFakes. *Ethics:* use only with
# consent; the same tools also help to detect fakes.
