# examples/gan.py
"""Offline float64 CPU GAN mechanics, grounded in Goodfellow et al. (2014).

Section 3, Equation 1 and Algorithm 1: https://arxiv.org/html/1406.2661v1#S3
The generator uses Section 3's non-saturating alternative, not minimax descent.
This authored two-dimensional mixture and tanh MLP pair are not original-paper
architectures, pretrained models, a convergence run or a benchmark reproduction.
"""

import argparse
from contextlib import contextmanager
import hashlib
import json
import math
from pathlib import Path

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


def positive_integer(value, name):
    """Reject empty dimensions, booleans and nonintegral dimensions."""
    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
        raise ValueError(f"{name} must be a positive integer")


def matrix(value, width, name):
    """Validate a finite nonempty float64 CPU minibatch without broadcasting."""
    if (not isinstance(value, torch.Tensor)
            or value.ndim != 2
            or value.shape[0] < 1
            or value.shape[1] != width
            or value.dtype != torch.float64
            or value.device.type != "cpu"
            or not torch.isfinite(value).all()):
        raise ValueError(f"{name} must be finite CPU float64 [B,{width}]")


# region: synthetic_data
def synthetic_data(batch=4, latent=2, seed=7):
    """Draw an equal mixture at (-1,-1)/(1,1), with isotropic stddev 0.15.

    Real samples and the two noise minibatches use independent local RNGs.
    The generator update receives fresh noise, as in Algorithm 1.
    """
    positive_integer(batch, "batch")
    positive_integer(latent, "latent")
    real_rng = torch.Generator().manual_seed(seed)
    d_rng = torch.Generator().manual_seed(seed + 1)
    g_rng = torch.Generator().manual_seed(seed + 2)
    mode = 2 * torch.randint(0, 2, (batch, 1), generator=real_rng) - 1
    real = mode + 0.15 * torch.randn(
        batch, 2, generator=real_rng, dtype=torch.float64
    )
    z_d = torch.randn(batch, latent, generator=d_rng, dtype=torch.float64)
    z_g = torch.randn(batch, latent, generator=g_rng, dtype=torch.float64)
    return real, z_d, z_g
# endregion: synthetic_data


# region: generator
def generator(latent=2, hidden=4, seed=7):
    """Construct G: [B,Z] -> [B,2], preserving global RNG state."""
    positive_integer(latent, "latent")
    positive_integer(hidden, "hidden")
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(seed)
        return nn.Sequential(
            nn.Linear(latent, hidden), nn.Tanh(), nn.Linear(hidden, 2)
        ).double()
# endregion: generator


# region: discriminator
def discriminator(hidden=4, seed=8):
    """Construct logits a: [B,2] -> [B,1]; D(x) = sigmoid(a(x))."""
    positive_integer(hidden, "hidden")
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(seed)
        return nn.Sequential(
            nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, 1)
        ).double()


def logistic_loss(logits, target):
    """Return mean BCE-with-logits for a scalar binary target."""
    matrix(logits, 1, "logits")
    if target not in (0, 1):
        raise ValueError("target must be 0 or 1")
    return F.binary_cross_entropy_with_logits(
        logits, torch.full_like(logits, target), reduction="mean"
    )
# endregion: discriminator


# region: discriminator_objective
def discriminator_objective(g, d, real, z_d):
    """Minimize negative Equation 1 in D, with G detached."""
    matrix(real, 2, "real")
    matrix(z_d, g[0].in_features, "z_d")
    if real.shape[0] != z_d.shape[0]:
        raise ValueError("real and noise batch sizes must agree")
    fake = g(z_d).detach()
    real_logits, fake_logits = d(real), d(fake)
    loss = logistic_loss(real_logits, 1) + logistic_loss(fake_logits, 0)
    return loss, real_logits, fake_logits
# endregion: discriminator_objective


def parameters(module):
    """Snapshot every parameter in a flat detached CPU tensor."""
    values = [p.detach().flatten() for p in module.parameters()]
    return torch.cat(values).clone()


def gradients(module):
    """Snapshot gradients, requiring every trainable parameter to have one."""
    values = [p.grad.flatten() for p in module.parameters() if p.requires_grad]
    result = torch.cat(values).detach().clone()
    if not torch.isfinite(result).all():
        raise FloatingPointError("nonfinite gradient")
    return result


# region: discriminator_update
def discriminator_update(g, d, real, z_d, learning_rate=0.05):
    """Clear D gradients and take one actual D-only SGD step."""
    if not math.isfinite(learning_rate) or learning_rate <= 0:
        raise ValueError("learning_rate must be finite and positive")
    optimizer = torch.optim.SGD(d.parameters(), lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    loss, real_logits, fake_logits = discriminator_objective(g, d, real, z_d)
    loss.backward()
    grad = gradients(d)
    optimizer.step()
    return loss.detach(), grad, real_logits.detach(), fake_logits.detach()
# endregion: discriminator_update


# region: generator_objective
@contextmanager
def frozen_discriminator(d):
    """Freeze D parameters across forward/backward, retaining the path to G."""
    flags = [p.requires_grad for p in d.parameters()]
    try:
        for p in d.parameters():
            p.requires_grad_(False)
        yield
    finally:
        for p, flag in zip(d.parameters(), flags):
            p.requires_grad_(flag)


def generator_objective(g, d, z_g):
    """Re-forward fresh fakes through updated D; minimize -log D(G(z))."""
    matrix(z_g, g[0].in_features, "z_g")
    fake = g(z_g)
    logits = d(fake)  # Neither detach() nor no_grad(): G needs this derivative.
    return logistic_loss(logits, 1), fake, logits
# endregion: generator_objective


# region: generator_update_and_trace
def alternating_step(g, d, real, z_d, z_g, learning_rate=0.05):
    """Execute D first, then non-saturating G, returning full state evidence."""
    matrix(z_g, g[0].in_features, "z_g")
    g_before, d_before = parameters(g), parameters(d)
    g_grads_before = [
        None if p.grad is None else p.grad.clone() for p in g.parameters()
    ]
    d_loss, d_grad, real_logits, fake_logits = discriminator_update(
        g, d, real, z_d, learning_rate
    )
    g_after_d, d_after_d = parameters(g), parameters(d)
    g_untouched = all(
        (old is None and p.grad is None)
        or (old is not None and p.grad is not None and torch.equal(old, p.grad))
        for old, p in zip(g_grads_before, g.parameters())
    )
    d_grad_after_d = [
        None if p.grad is None else p.grad.clone() for p in d.parameters()
    ]
    optimizer = torch.optim.SGD(g.parameters(), lr=learning_rate)
    optimizer.zero_grad(set_to_none=True)
    flags_before = [p.requires_grad for p in d.parameters()]
    with frozen_discriminator(d):
        g_loss, fake, g_logits = generator_objective(g, d, z_g)
        flags_during = [p.requires_grad for p in d.parameters()]
        g_loss.backward()
        g_grad = gradients(g)
        optimizer.step()
    d_grads_untouched = all(
        (old is None and p.grad is None)
        or (old is not None and p.grad is not None and torch.equal(old, p.grad))
        for old, p in zip(d_grad_after_d, d.parameters())
    )
    return {
        "d_loss": float(d_loss),
        "g_loss": float(g_loss.detach()),
        "real_logits": real_logits.tolist(),
        "d_fake_logits": fake_logits.tolist(),
        "g_fake_logits": g_logits.detach().tolist(),
        "generated_samples_before_g_update": fake.detach().tolist(),
        "g_before": g_before.tolist(),
        "g_after_d": g_after_d.tolist(),
        "g_after_g": parameters(g).tolist(),
        "g_gradient": g_grad.tolist(),
        "d_before": d_before.tolist(),
        "d_after_d": d_after_d.tolist(),
        "d_after_g": parameters(d).tolist(),
        "d_gradient": d_grad.tolist(),
        "g_gradients_untouched_by_d": g_untouched,
        "d_gradients_untouched_by_g": d_grads_untouched,
        "d_flags_before": flags_before,
        "d_flags_during_g": flags_during,
        "d_flags_after": [p.requires_grad for p in d.parameters()],
        "finite_gradients": bool(
            torch.isfinite(g_grad).all() and torch.isfinite(d_grad).all()
        ),
        "g_gradient_norm": float(torch.linalg.vector_norm(g_grad)),
        "d_gradient_norm": float(torch.linalg.vector_norm(d_grad)),
        "g_parameter_change": float(
            torch.linalg.vector_norm(parameters(g) - g_before)
        ),
        "d_parameter_change": float(
            torch.linalg.vector_norm(d_after_d - d_before)
        ),
    }
# endregion: generator_update_and_trace


def smoke(seed=7):
    """Return measured reductions, shapes, gradients and ordered SGD states."""
    g, d = generator(seed=seed), discriminator(seed=seed + 1)
    real, z_d, z_g = synthetic_data(seed=seed)
    trace = alternating_step(g, d, real, z_d, z_g)
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "seed": seed,
        "dtype": "float64",
        "device": "cpu",
        "learning_rate": 0.05,
        "reductions": {
            "discriminator": (
                "mean(real target-1 BCE) + mean(fake target-0 BCE)"
            ),
            "generator": "mean(fake target-1 BCE; non-saturating)",
        },
        "shapes": {
            "real": list(real.shape),
            "z_d": list(z_d.shape),
            "z_g": list(z_g.shape),
            "fake": list(torch.tensor(
                trace["generated_samples_before_g_update"]
            ).shape),
            "logits": list(torch.tensor(trace["g_fake_logits"]).shape),
        },
        "update_order": ["discriminator", "generator"],
        "boundary": (
            "One synthetic CPU update; "
            "no convergence or mode-collapse diagnosis."
        ),
        "trace": trace,
    }


def main():
    """Print one finite smoke record only when explicitly invoked."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("pass --smoke to run the offline reference")
    print(json.dumps(smoke(args.seed), allow_nan=False, sort_keys=True))


if __name__ == "__main__":
    main()