"""FlashReFocus for the Hugging Face Space (ZeroGPU), on the flashrefocus package (a copy of the repository's).

ZeroGPU gives a GPU only inside @spaces.GPU calls: every model is built on the CPU, moved to `cuda` at import time as an
ordinary module (Stage I with two separate U-Nets, no weights swapped in place), and the per-session image and depth
stay on the CPU between calls. The weights are downloaded on start; the FlashReFocus ones come from a private model
repository read with the HF_TOKEN secret."""
import os
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from PIL import Image

HERE = Path(__file__).resolve().parent
WEIGHT_REPO = os.environ.get("FLASHREFOCUS_WEIGHT_REPO", "ming0531/FlashReFocus-weights")


def fetch_weights():
    """Weights into ./weights (skipped when FLASHREFOCUS_WEIGHTS already points to a local copy)."""
    if "FLASHREFOCUS_WEIGHTS" in os.environ:
        return
    w = HERE / "weights"
    sd_files = ["model_index.json", "scheduler/scheduler_config.json", "tokenizer/*", "text_encoder/config.json",
                "text_encoder/model.safetensors", "unet/config.json", "unet/diffusion_pytorch_model.safetensors",
                "vae/config.json", "vae/diffusion_pytorch_model.safetensors"]
    snapshot_download("sd2-community/stable-diffusion-2-1-base", allow_patterns=sd_files,
                      local_dir=w / "stable-diffusion-2-1-base")
    snapshot_download("stable-diffusion-v1-5/stable-diffusion-v1-5",
                      allow_patterns=["vae/config.json", "vae/diffusion_pytorch_model.safetensors"],
                      local_dir=w / "stable-diffusion-v1-5")
    hf_hub_download("depth-anything/Depth-Anything-V2-Large", "depth_anything_v2_vitl.pth", local_dir=w)
    for f in ("stage1.pkl", "refiner.pt", "renderer.pt"):
        hf_hub_download(WEIGHT_REPO, f, local_dir=w, token=os.environ.get("HF_TOKEN"))
    os.environ["FLASHREFOCUS_WEIGHTS"] = str(w)


fetch_weights()
os.environ.setdefault("FLASHREFOCUS_BUILD_DEVICE", "cpu")      # third_party/PiSA-SR here builds on this device
from flashrefocus import Refocuser, Scene, paths  # noqa: E402
from flashrefocus.deblurring import Deblurrer  # noqa: E402
from flashrefocus.image_io import image_to_tensor, tensor_to_uint8  # noqa: E402
from flashrefocus.rendering import DepthEstimator, depth_edits, load_renderer  # noqa: E402


def build_refocuser():
    deblurrer = Deblurrer.from_pretrained(separate_unets=True, device="cpu")
    del deblurrer.model.text_encoder               # the empty-prompt embedding is computed once, at construction
    return Refocuser(deblurrer.to("cuda"), DepthEstimator.from_pretrained(device="cpu").to("cuda"),
                     load_renderer(paths.RENDERER, "cpu").to("cuda"))


refocuser = build_refocuser()


@torch.no_grad()
def restore(rgb, omega):
    """HxWx3 uint8 (sides multiples of 16) -> restored image (HxWx3 uint8) and its depth (HxW float32, 1 = far)."""
    refocuser.deblurrer.detail_weight = float(omega)
    scene = refocuser.prepare(image_to_tensor(rgb).cuda())
    return tensor_to_uint8(scene.image), scene.depth[0].cpu().numpy()


@torch.no_grad()
def render(rgb, depth, focal_depth, strength):
    """Bokeh (HxWx3 uint8) of the restored image with blur strength K and focal depth D_f; depth may be edited."""
    d = torch.from_numpy(np.ascontiguousarray(depth, dtype=np.float32))[None].cuda()
    return tensor_to_uint8(refocuser.render(Scene(image_to_tensor(rgb).cuda(), d), strength, focal_depth))


def two_focal_planes(depth, focal_depth_a, focal_depth_b):
    return depth_edits.two_focal_planes(torch.from_numpy(depth)[None], focal_depth_a, focal_depth_b)[0].numpy()


def tilt_shift(height, width, row, band=0.06):
    return depth_edits.tilt_shift(height, width, row, band=band)[0].numpy()


def depth_from_image(image, height, width):
    """A user depth map (near is bright, any size) -> D in [0, 1], 1 = far."""
    grey = Image.fromarray(np.asarray(image)).convert("L").resize((width, height), Image.BILINEAR)
    return (1.0 - np.asarray(grey, dtype=np.float32) / 255.0).astype(np.float32)