import os, sys, time, math, argparse
import numpy as np, torch, torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import PIL.Image as Image

torch.manual_seed(0); np.random.seed(0)

A = argparse.ArgumentParser()
A.add_argument("--data", default="/work/logos/logos64.npy")
A.add_argument("--out", default="/work/models/logo-gan")
A.add_argument("--steps", type=int, default=4000)
A.add_argument("--batch", type=int, default=64)
A.add_argument("--zdim", type=int, default=128)
A.add_argument("--lr", type=float, default=2e-4)
A.add_argument("--ckpt-every", type=int, default=500)
A.add_argument("--sample-every", type=int, default=500)
A.add_argument("--device", default="cuda")
a = A.parse_args()

os.makedirs(a.out, exist_ok=True)
log = open(os.path.join(a.out, "train.log"), "w")
def L(*x):
    s = " ".join(str(v) for v in x)
    print(s, flush=True); log.write(s + "\n"); log.flush()

dev = a.device if (a.device == "cuda" and torch.cuda.is_available()) else "cpu"
L("device", dev)

data = np.load(a.data)  # (N,3,64,64) 0..1
L("data", data.shape, "range", float(data.min()), float(data.max()))
X = torch.from_numpy(data).permute(0,3,1,2).to(dev)  # (N,3,64,64)
ds = TensorDataset(X); dl = DataLoader(ds, batch_size=a.batch, shuffle=True, drop_last=True)
n = len(data)

def conv_block(in_c, out_c, ks=4, s=2, p=1):
    return nn.Sequential(
        nn.Conv2d(in_c, out_c, ks, s, p, bias=False),
        nn.BatchNorm2d(out_c), nn.LeakyReLU(0.2, inplace=True))

class G(nn.Module):
    def __init__(self, zdim=128):
        super().__init__()
        self.fc = nn.Sequential(nn.Linear(zdim, 256*8*8), nn.BatchNorm1d(256*8*8), nn.ReLU())
        self.up = nn.Sequential(
            nn.ConvTranspose2d(256,128,4,2,1,bias=False), nn.BatchNorm2d(128), nn.ReLU(),
            nn.ConvTranspose2d(128,64,4,2,1,bias=False), nn.BatchNorm2d(64), nn.ReLU(),
            nn.ConvTranspose2d(64,3,4,2,1,bias=True), nn.Tanh())
    def forward(self, z):
        h = self.fc(z).view(-1,256,8,8)
        return self.up(h)  # (B,3,64,64) -1..1

class D(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3,64,4,2,1,bias=False), nn.BatchNorm2d(64), nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(64,128,4,2,1,bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(128,256,4,2,1,bias=False), nn.BatchNorm2d(256), nn.LeakyReLU(0.2,inplace=True),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(256,1))
    def forward(self, x): return self.net(x).squeeze(-1)  # (B,)

g = G(a.zdim).to(dev); d = D().to(dev)
opt_g = torch.optim.Adam(g.parameters(), a.lr, betas=(0.5,0.999))
opt_d = torch.optim.Adam(d.parameters(), a.lr, betas=(0.5,0.999))
BCE = nn.BCEWithLogitsLoss()
zdim = a.zdim

def sample_grid(zs, path, step):
    with torch.no_grad():
        imgs = g(zs).cpu().numpy()  # (B,3,64,64) -1..1
    imgs = (imgs*0.5+0.5).clip(0,1).transpose(0,2,3,1)  # (B,64,64,3)
    imgs = (imgs*255).astype(np.uint8)
    # 8x8 grid
    cols = rows = 8
    cell = 64; gap = 2
    W = cols*cell + (cols+1)*gap; H = rows*cell + (rows+1)*gap
    canvas = np.full((H, W, 3), 255, dtype=np.uint8)
    for i in range(min(len(imgs), cols*rows)):
        r = i // cols; c = i % cols
        y = gap + r*(cell+gap); x = gap + c*(cell+gap)
        canvas[y:y+cell, x:x+cell] = imgs[i]
    Image.fromarray(canvas).save(path)

def gen_samples(step):
    zs = torch.randn(a.batch, zdim, device=dev)
    name = f"grid_{step:05d}.png" if isinstance(step,int) else f"grid_{step}.png"
    sample_grid(zs, os.path.join(a.out, name), step)

it = iter(dl)
best_g = None
for step in range(1, a.steps+1):
    try:
        real = next(it)
    except StopIteration:
        it = iter(dl); real = next(it)
    real = real if isinstance(real, torch.Tensor) else real[0]
    b = real.size(0)
    # ---- D step ----
    opt_d.zero_grad()
    noise = torch.randn(b, zdim, device=dev)
    fake = g(noise).detach()
    loss_d = BCE(d(real), torch.ones(b, device=dev)) + BCE(d(fake), torch.zeros(b, device=dev))
    loss_d.backward(); opt_d.step()
    # ---- G step ----
    opt_g.zero_grad()
    noise = torch.randn(b, zdim, device=dev)
    fake = g(noise)
    loss_g = BCE(d(fake), torch.ones(b, device=dev))
    loss_g.backward(); opt_g.step()
    if step % 100 == 0:
        L(f"step {step}/{a.steps} d_loss {loss_d.item():.4f} g_loss {loss_g.item():.4f}")
    if step % a.ckpt_every == 0:
        torch.save({"g": g.state_dict(), "d": d.state_dict(), "step": step, "zdim": zdim},
                   os.path.join(a.out, "ckpt.pt"))
        gen_samples(step)
    if step % a.sample_every == 0:
        gen_samples(step)

# final
torch.save({"g": g.state_dict(), "d": d.state_dict(), "step": a.steps, "zdim": zdim},
           os.path.join(a.out, "final.pt"))
gen_samples(a.steps)
# also save a fixed-seed grid for the card (reproducible)
torch.manual_seed(1234)
gen_samples("final")
L(f"DONE final step {a.steps}")
L("ckpt files:", sorted(os.listdir(a.out)))
log.close()