"""Minimal PiSA-SR inference on the diffusers repack, in default or adjustable mode."""
import argparse, os
import torch
from PIL import Image
from torchvision import transforms
import torchvision.transforms.functional as TF
from diffusers import AutoencoderKL, UNet2DConditionModel
from transformers import AutoTokenizer, CLIPTextModel

BASE = "stabilityai/stable-diffusion-2-1-base"
REPO = "ndtran0101/pisa-sr-diffusers"


def prepare(img, upscale=4):
    """Upsample a low-quality image to the diffusion input size, aligned to a multiple of 8.

    The upsample happens before the UNet, so cost scales with the output size.
    """
    img = img.resize((img.width * upscale, img.height * upscale))
    return img.resize((img.width - img.width % 8, img.height - img.height % 8), Image.LANCZOS)


def adain(target, source):
    """Match the target's per-channel mean and std to the source.

    The upstream pipeline applies this colour transfer after decoding; skipping it
    shifts colour noticeably.
    """
    t = TF.to_tensor(target).unsqueeze(0)
    s = TF.to_tensor(source).unsqueeze(0)
    t_mean, t_std = t.mean([2, 3], keepdim=True), t.std([2, 3], keepdim=True)
    s_mean, s_std = s.mean([2, 3], keepdim=True), s.std([2, 3], keepdim=True)
    out = ((t - t_mean) / (t_std + 1e-5)) * s_std + s_mean
    return transforms.ToPILImage()(out.clamp(0, 1)[0])


@torch.no_grad()
def upscale(path, base=BASE, repo=REPO, lambda_pix=None, lambda_sem=None,
            device="cuda", dtype=torch.float16, color_fix=True):
    """Upscale one image x4 and return it as a PIL image.

    Runs one UNet pass at t=1 on an empty prompt and subtracts the predicted residual;
    there is no scheduler and no noise. Passing both lambdas switches to the two-pass
    adjustable mode, which additionally loads unet_pix.

    path: low-quality input image
    repo: HuggingFace id or local directory holding unet_full/ and unet_pix/
    lambda_pix / lambda_sem: strength of the pixel and semantic branches
    """
    adjustable = lambda_pix is not None and lambda_sem is not None

    tok = AutoTokenizer.from_pretrained(base, subfolder="tokenizer")
    te = CLIPTextModel.from_pretrained(base, subfolder="text_encoder").to(device, dtype).eval()
    vae = AutoencoderKL.from_pretrained(base, subfolder="vae").to(device, dtype).eval()
    unet = UNet2DConditionModel.from_pretrained(repo, subfolder="unet_full").to(device, dtype).eval()

    src = Image.open(path).convert("RGB")
    img = prepare(src)
    x = TF.to_tensor(img).unsqueeze(0).to(device, dtype) * 2 - 1

    ids = tok("", max_length=tok.model_max_length, padding="max_length",
              truncation=True, return_tensors="pt").input_ids.to(device)
    emb = te(ids)[0].to(dtype)
    t = torch.tensor([1], device=device).long()

    z = vae.encode(x).latent_dist.sample() * vae.config.scaling_factor
    if adjustable:
        unet_pix = UNet2DConditionModel.from_pretrained(
            repo, subfolder="unet_pix").to(device, dtype).eval()
        pred_sem = unet(z, t, encoder_hidden_states=emb).sample
        pred_pix = unet_pix(z, t, encoder_hidden_states=emb).sample
        pred = lambda_pix * pred_pix + lambda_sem * (pred_sem - pred_pix)
    else:
        pred = unet(z, t, encoder_hidden_states=emb).sample

    out = vae.decode((z - pred) / vae.config.scaling_factor).sample.clamp(-1, 1)
    pil = transforms.ToPILImage()((out * 0.5 + 0.5).clamp(0, 1)[0].float().cpu())
    return adain(pil, img) if color_fix else pil


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", default="sr.png")
    ap.add_argument("--repo", default=REPO)
    ap.add_argument("--base", default=BASE)
    ap.add_argument("--lambda_pix", type=float, default=None)
    ap.add_argument("--lambda_sem", type=float, default=None)
    ap.add_argument("--no_color_fix", action="store_true")
    a = ap.parse_args()

    img = upscale(a.input, base=a.base, repo=a.repo,
                  lambda_pix=a.lambda_pix, lambda_sem=a.lambda_sem,
                  color_fix=not a.no_color_fix)
    img.save(a.output)
    print(f"{a.input} -> {a.output}  {img.size[0]}x{img.size[1]}")


if __name__ == "__main__":
    main()