import os, sys, time, math, argparse
import numpy as np, torch, torch.nn as nn
import PIL.Image as Image

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

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)
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)

def grid(imgs, path):
    imgs=(imgs*0.5+0.5).clip(0,1).transpose(0,2,3,1).astype(np.uint8) if imgs.min()<-0.5 else imgs.astype(np.uint8)
    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 main():
    a=argparse.Namespace(data="/work/logos/logos_big.npy", out="/work/models/logo-gan-v2",
                         steps=8000, bs=16, zdim=128, lr=2e-4, betas=(0.5,0.999), dev="cuda",
                         ckpt_every=1000, sample_every=1000)
    os.makedirs(a.out, exist_ok=True)
    dev = a.dev if (a.dev=="cuda" and torch.cuda.is_available()) else "cpu"
    imgs = np.load(a.data).astype(np.float32)          # (N,64,64,3) 0..1
    imgs = (imgs - 0.5)/0.5                            # -> [-1,1]
    X = torch.from_numpy(imgs).permute(0,3,1,2).to(dev)
    N = X.shape[0]
    g = G(a.zdim).to(dev); d = D().to(dev)
    go = torch.optim.Adam(g.parameters(), a.lr, betas=a.betas)
    do = torch.optim.Adam(d.parameters(), a.lr, betas=a.betas)
    logf = open(os.path.join(a.out,"train.log"),"a")
    def lg(s): print(s, flush=True); logf.write(s+"\n"); logf.flush()
    lg(f"START v2 N={N} dev={dev} steps={a.steps} bs={a.bs} zdim={a.zdim} lr={a.lr}")
    g_params=sum(p.numel() for p in g.parameters()); d_params=sum(p.numel() for p in d.parameters())
    lg(f"params generator={g_params} discriminator={d_params} total={g_params+d_params}")
    torch.manual_seed(0); np.random.seed(0)
    for step in range(1, a.steps+1):
        idx = torch.randint(0, N, (a.bs,))
        real = X[idx]
        z = torch.randn(a.bs, a.zdim, device=dev)
        # D
        do.zero_grad()
        fake = g(z).detach()
        dloss = -(torch.log(torch.sigmoid(d(real)).clamp(min=1e-7)).mean() + torch.log(1-torch.sigmoid(d(fake)).clamp(min=1e-7)).mean())
        dloss.backward(); do.step()
        # G
        go.zero_grad()
        fake = g(z)
        gloss = -torch.log(torch.sigmoid(d(fake)).clamp(min=1e-7)).mean()
        gloss.backward(); go.step()
        if step % 200 == 0:
            lg(f"step {step}/{a.steps} dloss={float(dloss):.4f} gloss={float(gloss):.4f}")
        if step % a.ckpt_every == 0:
            torch.save({"g":g.state_dict(),"d":d.state_dict(),"step":step,"zdim":a.zdim}, os.path.join(a.out,f"ckpt_{step:05d}.pt"))
        if step % a.sample_every == 0:
            g.eval()
            with torch.no_grad():
                fake = g(torch.randn(64, a.zdim, device=dev)).cpu().numpy()
            grid(fake, os.path.join(a.out, f"grid_{step:05d}.png"))
            g.train()
    # final
    g.eval()
    with torch.no_grad():
        fake = g(torch.randn(64, a.zdim, device=dev)).cpu().numpy()
    grid(fake, os.path.join(a.out, "grid_final.png"))
    torch.save({"g":g.state_dict(),"d":d.state_dict(),"step":a.steps,"zdim":a.zdim}, os.path.join(a.out,"final.pt"))
    lg(f"DONE final step {a.steps}")
    lg("files: " + ", ".join(sorted(os.listdir(a.out))))
    logf.close()

if __name__=="__main__":
    main()