# sr_diffusion.py
# -----------------------------------------------------------
# Super-Resolution Diffusion for CIFAR-10 (32×32 → 128×128)
# Stage-2 after your text-to-image model
# -----------------------------------------------------------

import os, math, argparse, random, glob
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, utils as vutils
from tqdm import tqdm

# ------------------ utils ------------------
def set_seed(seed: int = 0):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def _raw_module(m):
    return getattr(m, "_orig_mod", m)

# ------------------ time embedding ------------------
class SinusoidalTimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        emb = torch.log(torch.tensor(10000.0)) / (half - 1)
        emb = torch.exp(torch.arange(half, device=t.device) * -emb)
        emb = t[:, None].float() * emb[None, :]
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
        if self.dim % 2 == 1:
            emb = F.pad(emb, (0,1))
        return emb

# ------------------ building blocks ------------------
class ResBlock(nn.Module):
    def __init__(self, dim_in, dim_out, cond_dim):
        super().__init__()
        self.norm1 = nn.GroupNorm(32, dim_in)
        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, padding=1)
        self.norm2 = nn.GroupNorm(32, dim_out)
        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, padding=1)
        self.emb = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, dim_out * 2)
        )
        self.act = nn.SiLU()
        self.skip = nn.Conv2d(dim_in, dim_out, 1) if dim_in != dim_out else nn.Identity()

    def forward(self, x, cond):
        h = self.conv1(self.act(self.norm1(x)))
        scale, shift = self.emb(cond).chunk(2, dim=1)
        scale, shift = scale[:, :, None, None], shift[:, :, None, None]
        h = h * (1 + scale) + shift
        h = self.conv2(self.act(self.norm2(h)))
        return h + self.skip(x)

class Down(nn.Module):
    def __init__(self, dim_in, dim_out, cond_dim):
        super().__init__()
        self.block1 = ResBlock(dim_in, dim_out, cond_dim)
        self.block2 = ResBlock(dim_out, dim_out, cond_dim)
        self.down = nn.Conv2d(dim_out, dim_out, 4, stride=2, padding=1)

    def forward(self, x, cond):
        x = self.block1(x, cond)
        x = self.block2(x, cond)
        x = self.down(x)
        return x

class Up(nn.Module):
    def __init__(self, dim_in, dim_out, cond_dim):
        super().__init__()
        self.block1 = ResBlock(dim_in, dim_out, cond_dim)
        self.block2 = ResBlock(dim_out, dim_out, cond_dim)
        self.up = nn.ConvTranspose2d(dim_out, dim_out, 4, stride=2, padding=1)

    def forward(self, x, cond):
        x = self.block1(x, cond)
        x = self.block2(x, cond)
        x = self.up(x)
        return x

class UNetSR(nn.Module):
    def __init__(self, in_ch=3, base=64, cond_dim=256):
        super().__init__()
        self.in_conv = nn.Conv2d(in_ch, base, 3, padding=1)
        self.down1 = Down(base, base*2, cond_dim)
        self.down2 = Down(base*2, base*4, cond_dim)
        self.mid1 = ResBlock(base*4, base*4, cond_dim)
        self.mid2 = ResBlock(base*4, base*4, cond_dim)
        self.up2 = Up(base*4, base*2, cond_dim)
        self.up1 = Up(base*2, base, cond_dim)
        self.out_norm = nn.GroupNorm(8, base)
        self.out_conv = nn.Conv2d(base, in_ch, 3, padding=1)
        self.act = nn.SiLU()

    def forward(self, x, cond):
        x0 = self.in_conv(x)
        d1 = self.down1(x0, cond)
        d2 = self.down2(d1, cond)
        m = self.mid1(d2, cond)
        m = self.mid2(m, cond)
        u2 = self.up2(m + d2, cond)
        u1 = self.up1(u2 + d1, cond)
        h = self.act(self.out_norm(u1))
        return self.out_conv(h)

# ------------------ diffusion ------------------
class Diffusion(nn.Module):
    def __init__(self, model: nn.Module, T=1000, schedule="cosine"):
        super().__init__()
        self.model = model
        self.T = T
        self.register_buffer("betas", self.make_betas(T, schedule))
        alphas = 1.0 - self.betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        self.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))
        from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
        self.scheduler = DDPMScheduler(num_train_timesteps=self.T, beta_schedule="squaredcos_cap_v2")

    @staticmethod
    def make_betas(T, schedule):
        if schedule == "linear":
            return torch.linspace(1e-4, 0.02, T)
        s = 0.008
        steps = T + 1
        x = torch.linspace(0, T, steps)
        alphas_cum = torch.cos(((x / T) + s) / (1 + s) * math.pi * 0.5) ** 2
        alphas_cum = alphas_cum / alphas_cum[0]
        betas = 1 - (alphas_cum[1:] / alphas_cum[:-1])
        return betas.clamp(max=0.999)

    @torch.no_grad()
    def sample(self, shape, steps=1000, cond=None):
        device = next(self.parameters()).device
        x = torch.randn(shape, device=device)
        self.scheduler.set_timesteps(steps)
        for t in self.scheduler.timesteps:
            eps_pred = self.model(x, torch.full((shape[0],), t, device=device, dtype=torch.long), cond)
            x = self.scheduler.step(eps_pred, t, x).prev_sample
        return x

# ------------------ Wrapper ------------------
class CondUNetSR(nn.Module):
    def __init__(self, base=64, time_dim=128, cond_dim=256):
        super().__init__()
        self.time_emb = SinusoidalTimeEmbedding(time_dim)
        self.time_mlp = nn.Sequential(
            nn.Linear(time_dim, cond_dim), nn.SiLU(), nn.Linear(cond_dim, cond_dim)
        )
        # learnable projection for conditioning image
        self.cond_proj = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.SiLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.SiLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(32, cond_dim)
        )
        self.net = UNetSR(in_ch=3, base=base, cond_dim=cond_dim)

    def forward(self, x, t, cond_img=None):
        t_emb = self.time_mlp(self.time_emb(t))  # [B, 256]
        if cond_img is not None:
            cond_emb = self.cond_proj(cond_img)  # [B, 256]
            cond = t_emb + cond_emb
        else:
            cond = t_emb
        return self.net(x, cond)

# ------------------ Trainer ------------------
class SRTrainer:
    def __init__(self, args):
        self.args = args
        set_seed(args.seed)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        tfm_hr = transforms.Compose([
            transforms.Resize(args.hr_size),
            transforms.ToTensor(),
            transforms.Lambda(lambda x: x * 2 - 1),
        ])
        self.ds = datasets.CIFAR10(args.data_root, train=True, transform=tfm_hr, download=True)
        self.loader = DataLoader(self.ds, batch_size=args.bs, shuffle=True,
                                 num_workers=args.workers, drop_last=True)

        self.model = CondUNetSR(base=args.base, time_dim=128, cond_dim=args.cond_dim).to(self.device)
        self.diff = Diffusion(self.model, T=args.T, schedule=args.schedule).to(self.device)
        self.opt = torch.optim.AdamW(self.model.parameters(), lr=args.lr)

        os.makedirs(args.outdir, exist_ok=True)

        # preview set
        test_ds = datasets.CIFAR10(args.data_root, train=False,
                                   transform=transforms.ToTensor(), download=True)
        self.sample_imgs = torch.stack([test_ds[i][0] for i in range(16)]).to(self.device) * 2 - 1

        # auto resume
        self.start_epoch = 0
        if args.resume:
            self.start_epoch = self.try_resume()

    # ---- checkpoints ----
    def save_ckpt(self, epoch):
        ckpt = {
            "model": _raw_module(self.model).state_dict(),
            "opt": self.opt.state_dict(),
            "epoch": epoch
        }
        path = os.path.join(self.args.outdir, f"ckpt_epoch{epoch}.pt")
        torch.save(ckpt, path)
        print(f"[ckpt] saved {path}")

    def load_ckpt(self, path):
        ckpt = torch.load(path, map_location=self.device)
        self.model.load_state_dict(ckpt["model"])
        self.opt.load_state_dict(ckpt["opt"])
        print(f"[ckpt] loaded {path} (epoch {ckpt['epoch']})")
        return ckpt["epoch"]

    def try_resume(self):
        ckpts = sorted(glob.glob(os.path.join(self.args.outdir, "ckpt_epoch*.pt")))
        if len(ckpts) == 0:
            print("[resume] no checkpoints found")
            return 0
        last_ckpt = ckpts[-1]
        epoch = self.load_ckpt(last_ckpt)
        return epoch

    # ---- training ----
    def train(self):
        self.model.train()
        for epoch in range(self.start_epoch, self.args.epochs):
            losses = []
            pbar = tqdm(self.loader, desc=f"Epoch {epoch+1}/{self.args.epochs}")
            for hr, _ in pbar:
                hr = hr.to(self.device)
                lr = F.interpolate(hr, size=32, mode="bilinear")   # low-res input
                t = torch.randint(0, self.args.T, (hr.size(0),), device=self.device)
                noise = torch.randn_like(hr)
                a_t = self.diff.sqrt_alphas_cumprod[t][:, None, None, None]
                am1_t = self.diff.sqrt_one_minus_alphas_cumprod[t][:, None, None, None]
                hr_noisy = a_t * hr + am1_t * noise
                pred = self.model(hr_noisy, t, cond_img=lr)
                loss = F.mse_loss(pred, noise)
                self.opt.zero_grad(set_to_none=True)
                loss.backward()
                nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.opt.step()
                losses.append(loss.item())
                pbar.set_postfix(loss=loss.item())
            avg_loss = sum(losses) / len(losses)
            print(f"Epoch {epoch+1}/{self.args.epochs} | Avg Loss: {avg_loss:.4f}")

            # ---- checkpoint + sample ----
            if (epoch + 1) % self.args.sample_every == 0:
                self.save_ckpt(epoch + 1)
                self.model.eval()
                with torch.no_grad():
                    sr = self.super_resolve(self.sample_imgs, steps=self.args.steps)
                    grid = vutils.make_grid(sr, nrow=4)
                    save_path = os.path.join(self.args.outdir, f"sr_epoch{epoch+1}.png")
                    vutils.save_image(grid, save_path)
                    print(f"[sample] saved {save_path}")
                self.model.train()

    @torch.no_grad()
    def super_resolve(self, lr_imgs, steps=1000):
        self.model.eval()
        sr_imgs = self.diff.sample((lr_imgs.size(0), 3, self.args.hr_size, self.args.hr_size),
                                   steps=steps, cond=lr_imgs)
        return (sr_imgs.clamp(-1,1)+1)*0.5

# ------------------ CLI ------------------
def build_parser():
    p = argparse.ArgumentParser(description="Super-Resolution Diffusion (CIFAR 32→HR)")
    p.add_argument("--train", action="store_true")
    p.add_argument("--sample", action="store_true")
    p.add_argument("--data_root", type=str, default="./data")
    p.add_argument("--epochs", type=int, default=20)
    p.add_argument("--bs", type=int, default=128)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--T", type=int, default=1000)
    p.add_argument("--schedule", type=str, default="cosine", choices=["cosine", "linear"])
    p.add_argument("--base", type=int, default=64)
    p.add_argument("--cond_dim", type=int, default=256)
    p.add_argument("--workers", type=int, default=4)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--outdir", type=str, default="outputs_sr")
    p.add_argument("--hr_size", type=int, default=128)
    p.add_argument("--steps", type=int, default=50)
    p.add_argument("--sample_every", type=int, default=10,
                   help="Save checkpoint + preview samples every N epochs")
    p.add_argument("--resume", action="store_true",
                   help="Auto-resume from the latest checkpoint in outdir")
    p.add_argument("--prompts", nargs="+", type=str, default=None,
                   help="Text prompts for text→image→super-res pipeline")
    return p

def main():
    args = build_parser().parse_args()
    trainer = SRTrainer(args)

    if args.train:
        trainer.train()

    if args.sample:
        # ---- load checkpoint before sampling ----
        ckpts = sorted(glob.glob(os.path.join(args.outdir, "ckpt_epoch*.pt")))
        if len(ckpts) == 0:
            print("[sample] No checkpoints found in outdir, cannot sample")
            return
        last_ckpt = ckpts[-1]
        trainer.load_ckpt(last_ckpt)

        if args.prompts:
            # -------- Stage-1: text-to-image (32x32) --------
            from t2i_text import Trainer as T2ITrainer, build_parser as t2i_parser

            # load stage-1 checkpoint
            t2i_args = t2i_parser().parse_args([])
            t2i_trainer = T2ITrainer(t2i_args)
            t2i_trainer.load_ckpt("outputs/ckpt_latest.pt")  # adjust path if needed

            print(f"[stage-1] generating low-res 32x32 from prompts: {args.prompts}")
            lr_imgs = t2i_trainer.sample(prompts=args.prompts, steps=50, n_per_prompt=1)
            lr_imgs = lr_imgs.to(trainer.device) * 2 - 1  # normalize to [-1,1]

        else:
            # -------- fallback: use CIFAR test images --------
            ds = datasets.CIFAR10(args.data_root, train=False,
                                  transform=transforms.ToTensor(), download=True)
            lr_imgs = torch.stack([ds[i][0] for i in range(16)])
            lr_imgs = (lr_imgs*2-1).to(trainer.device)

        # -------- Stage-2: super-resolution (32→128) --------
        sr = trainer.super_resolve(lr_imgs, steps=args.steps)
        grid = vutils.make_grid(sr, nrow=4)
        os.makedirs(args.outdir, exist_ok=True)

        if args.prompts:
            save_path = os.path.join(args.outdir, "sr_from_prompt.png")
        else:
            save_path = os.path.join(args.outdir, "sr_samples.png")

        vutils.save_image(grid, save_path)
        print(f"[sample] saved {save_path}")


if __name__ == "__main__":
    main()