"""
Material maps for the textures this Space makes, and for any picture of a
surface (/material_maps): one picture in, the maps a PBR renderer needs out,
all tiling when the picture tiles. `compute` runs on the GPU and is called
from inside app.py's GPU functions, so a texture and its maps are one run.

How it goes:
- Two small ESRGAN nets trained on texture sets (Joey Ballentine's Material
  Map Generator, Apache-2.0): one gives a tangent-space normal map, the other
  displacement and roughness. They are what reads a painted brick as raised
  and its white mortar as sunk, which brightness alone gets backwards.
- Lotus-G normal (Apache-2.0, Stable Diffusion 2 fine-tuned for surface
  normals, one step) gives the broad shape: the rounded top of a cobble, the
  bevel of a brick. Its low frequencies and the ESRGAN detail are merged in
  slope space with the slopes of the displacement map, so normal and height
  agree.
- CLIP (MIT) names the material class, which sets the roughness level and
  whether anything is metal; the ESRGAN roughness adds the variation.

Every convolution pads circularly (the ESRGAN input is wrapped, the Lotus UNet
and VAE have their padding mode switched), and every filter wraps, so a
seamless picture gives seamless maps.
"""

import os
import tempfile
import time

os.environ["GRADIO_TEMP_DIR"] = os.path.join(tempfile.gettempdir(), "gradio")
os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True)

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from diffusers import AutoencoderKL, UNet2DConditionModel
from huggingface_hub import hf_hub_download
from PIL import Image
from spandrel import ModelLoader
from transformers import CLIPModel, CLIPProcessor

import embeddings
from embeddings import CLASSES, CLIP_REPO, LOTUS_REPO, features

DEVICE = "cuda"
MAX_SIDE = 2048
# The side the models read the picture at; the maps go out at the picture's own
# size. Measured 2026-10-07 on six Poly Haven scans at 2K: read at 2048 the nets
# gave normals twice too steep, 14 to 31 degrees off the scan against 8 to 25
# when read at 1024 and enlarged, with nothing above 1024 that matched the scan.
MODEL_SIDE = 1024
# Slopes scaled to half: measured against Poly Haven scans (scripts/compare-normal-maps.mjs),
# full strength was about twice too steep, and invented relief on painted, flat surfaces.
NORMAL_STRENGTH = 0.5
LOTUS_SIDE = 768  # Lotus-G is Stable Diffusion 2 base: 512 to 768 is home.
MAPS_REPO = "InvokeAI/pbr-material-maps"
MAPS_REVISION = "b7ca9ebc6e14688a69d41872d2b9c80ea453e8f0"


def _log(*parts):
    print("[maps]", *parts, flush=True)


def _circular(module: torch.nn.Module) -> None:
    for m in module.modules():
        if isinstance(m, torch.nn.Conv2d) and m.padding not in (0, (0, 0)):
            m.padding_mode = "circular"


def _esrgan(name: str) -> torch.nn.Module:
    path = hf_hub_download(MAPS_REPO, name, revision=MAPS_REVISION)
    return ModelLoader().load_from_file(path).model.eval().half().to(DEVICE)


normal_net = _esrgan("normal_map_generator.safetensors")
franken_net = _esrgan("franken_map_generator.safetensors")

lotus_unet = UNet2DConditionModel.from_pretrained(LOTUS_REPO, subfolder="unet", torch_dtype=torch.float16).to(DEVICE)
lotus_vae = AutoencoderKL.from_pretrained(LOTUS_REPO, subfolder="vae", torch_dtype=torch.float16).to(DEVICE)
_circular(lotus_unet)
_circular(lotus_vae)
# The task embedding that selects the normal head (see Lotus's infer.py).
_task = torch.tensor([[1.0, 0.0]])
TASK_EMB = torch.cat([torch.sin(_task), torch.cos(_task)], dim=-1).half().to(DEVICE)

clip_model = CLIPModel.from_pretrained(CLIP_REPO).eval().half().to(DEVICE)
clip_processor = CLIPProcessor.from_pretrained(CLIP_REPO)

# The empty prompt Lotus runs with and the class words for CLIP, embedded once
# in a process of their own (embeddings.py says why).
_embedded = embeddings.load()
EMPTY_PROMPT = _embedded["empty_prompt"].half().to(DEVICE)
CLASS_EMB = _embedded["class_emb"].to(DEVICE)
_log("models ready")


# --- plain-array helpers, all wrapping at the edges -------------------------

def _blur(a: torch.Tensor, sigma: float) -> torch.Tensor:
    """Separable Gaussian on an (H, W) tensor, wrapping around the edges."""
    if sigma <= 0:
        return a
    radius = max(1, int(3 * sigma))
    x = torch.arange(-radius, radius + 1, device=a.device, dtype=a.dtype)
    k = torch.exp(-(x**2) / (2 * sigma**2))
    k = k / k.sum()
    out = F.pad(a[None, None], (radius, radius, 0, 0), mode="circular")
    out = F.conv2d(out, k.view(1, 1, 1, -1))
    out = F.pad(out, (0, 0, radius, radius), mode="circular")
    return F.conv2d(out, k.view(1, 1, -1, 1))[0, 0]


def _resize_wrap(x: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
    """Resize (N, C, H, W) so the result still tiles: pad by wrapping, scale, crop."""
    h, w = x.shape[-2:]
    if (h, w) == size:
        return x
    pad = 4
    big = F.pad(x, (pad, pad, pad, pad), mode="circular")
    sy, sx = size[0] / h, size[1] / w
    out = F.interpolate(big, size=(round((h + 2 * pad) * sy), round((w + 2 * pad) * sx)), mode="bicubic", align_corners=False)
    oy, ox = round(pad * sy), round(pad * sx)
    return out[..., oy : oy + size[0], ox : ox + size[1]]


def _enlarge(a: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
    """An (H, W) map at the picture's size, still tiling."""
    return _resize_wrap(a[None, None], size)[0, 0]


def _slopes(n: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """(3, H, W) normals, x right, y up → slopes dh/dx and dh/dy_up, tilt removed.
    A normal is (-dh/dx, -dh/dy, 1) normalised."""
    nz = n[2].clamp(min=0.2)
    p, q = -n[0] / nz, -n[1] / nz
    return p - p.mean(), q - q.mean()


def _grad(h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Central differences with wrap: dh/dx and dh/dy_up (rows run down)."""
    gx = (torch.roll(h, -1, 1) - torch.roll(h, 1, 1)) / 2
    gy = (torch.roll(h, 1, 0) - torch.roll(h, -1, 0)) / 2
    return gx, gy


def _stretch(a: torch.Tensor, lo: float = 0.005, hi: float = 0.995) -> torch.Tensor:
    flat = a.flatten()
    if flat.numel() > 1_000_000:
        flat = flat[:: flat.numel() // 1_000_000 + 1]
    a_lo, a_hi = torch.quantile(flat, lo), torch.quantile(flat, hi)
    return ((a - a_lo) / (a_hi - a_lo).clamp(min=1e-6)).clamp(0, 1)


def _smoothstep(e0: float, e1: float, x: torch.Tensor) -> torch.Tensor:
    t = ((x - e0) / (e1 - e0)).clamp(0, 1)
    return t * t * (3 - 2 * t)


# --- the models --------------------------------------------------------------

def _run_esrgan(net: torch.nn.Module, rgb: torch.Tensor) -> torch.Tensor:
    """(1, 3, H, W) in [0, 1] → (3, H, W) in [0, 1]; wrapped so the edges tile."""
    pad = 32
    x = F.pad(rgb, (pad, pad, pad, pad), mode="circular").half()
    return net(x)[0, :, pad:-pad, pad:-pad].float().clamp(0, 1)


def _run_lotus(rgb: torch.Tensor) -> torch.Tensor:
    """(1, 3, H, W) in [0, 1] → (3, H, W) unit normals, x right, y up, z out."""
    h, w = rgb.shape[-2:]
    scale = min(1.0, LOTUS_SIDE / max(h, w))
    size = (max(64, round(h * scale / 64) * 64), max(64, round(w * scale / 64) * 64))
    x = _resize_wrap(rgb, size) * 2 - 1
    latents = lotus_vae.encode(x.half()).latent_dist.mode() * lotus_vae.config.scaling_factor
    noise = torch.randn(latents.shape, generator=torch.Generator(DEVICE).manual_seed(0), device=DEVICE, dtype=latents.dtype)
    x0 = lotus_unet(
        torch.cat([latents, noise], dim=1),
        torch.tensor([999], device=DEVICE),
        encoder_hidden_states=EMPTY_PROMPT,
        class_labels=TASK_EMB,
        return_dict=False,
    )[0]
    decoded = lotus_vae.decode(x0 / lotus_vae.config.scaling_factor, return_dict=False)[0].float().clamp(-1, 1)
    n = _resize_wrap(decoded, (h, w))[0]
    return n / n.norm(dim=0, keepdim=True).clamp(min=1e-6)


def _classify(image: Image.Image) -> tuple[torch.Tensor, list[tuple[str, float]]]:
    inputs = clip_processor(images=image, return_tensors="pt").to(DEVICE)
    emb = F.normalize(features(clip_model.get_image_features(pixel_values=inputs.pixel_values.half())).float(), dim=-1)
    probs = (100 * emb @ CLASS_EMB.T).softmax(dim=-1)[0]
    order = probs.argsort(descending=True)[:3].tolist()
    return probs, [(CLASSES[i][0], round(float(probs[i]), 3)) for i in order]


@torch.no_grad()
def compute(image: Image.Image):
    """On the GPU: the four maps as arrays in [0, 1] at the picture's size, and
    what the surface was taken for. The models read it at MODEL_SIDE at most."""
    t0 = time.time()
    full, picture = (image.height, image.width), image
    if max(image.size) > MODEL_SIDE:
        if image.width == image.height and image.width % MODEL_SIDE == 0:
            image = image.reduce(image.width // MODEL_SIDE)  # whole blocks: a tile keeps its wrap
        else:
            scale = MODEL_SIDE / max(image.size)
            image = image.resize((max(8, round(image.width * scale)), max(8, round(image.height * scale))), Image.LANCZOS)
    w, h = image.size
    rgb = torch.from_numpy(np.asarray(image, np.float32) / 255).permute(2, 0, 1)[None].to(DEVICE)
    s = max(w, h) / 768  # filter sizes were tuned at 768 px

    es_normal = _run_esrgan(normal_net, rgb) * 2 - 1
    franken = _run_esrgan(franken_net, rgb)
    lotus = _run_lotus(rgb)
    probs, top = _classify(picture)  # CLIP sizes it itself
    _log(f"models {time.time() - t0:.2f}s", top)

    # Height: the texture-trained displacement. Normal: the displacement's
    # slopes, plus Lotus's broad shape and the ESRGAN normal's fine detail.
    height = _stretch(franken[2])
    pd, qd = _grad(height)
    pd, qd = pd * max(w, h) / 40, qd * max(w, h) / 40
    pl, ql = _slopes(lotus)
    pe, qe = _slopes(es_normal)
    p = (_blur(pl, 2 * s) + pe - _blur(pe, 3 * s) + pd) / 2 * NORMAL_STRENGTH
    q = (_blur(ql, 2 * s) + qe - _blur(qe, 3 * s) + qd) / 2 * NORMAL_STRENGTH
    p, q = _enlarge(p, full), _enlarge(q, full)
    normal = torch.stack([-p, -q, torch.ones_like(p)])
    normal = normal / normal.norm(dim=0, keepdim=True)

    # Roughness: the class sets the level, the ESRGAN map the variation.
    level = sum(float(probs[i]) * c[2] for i, c in enumerate(CLASSES))
    rough = franken[1]
    roughness = (level + (rough - rough.median()) * 1.2).clamp(0.04, 1)

    # Metallic: bare metal is metal all over; rusted metal only where grey
    # steel shows (low saturation). Everything else is not metal.
    metal = sum(float(probs[i]) for i, c in enumerate(CLASSES) if c[3] == "metal")
    rust = sum(float(probs[i]) for i, c in enumerate(CLASSES) if c[3] == "rust")
    mx, mn = rgb[0].max(dim=0).values, rgb[0].min(dim=0).values
    saturation = (mx - mn) / mx.clamp(min=1e-3)
    bare = _smoothstep(0.35, 0.15, saturation)
    metallic = _smoothstep(0.35, 0.65, metal + rust * bare)
    roughness = roughness - metallic * 0.15
    height, roughness, metallic = (_enlarge(m, full).clamp(0, 1) for m in (height, roughness, metallic))

    info = {
        "material": top[0][0],
        "classes": [{"label": label, "p": p_} for label, p_ in top],
        "roughness_level": round(level, 3),
        "metal": round(metal, 3),
        "gpu_seconds": round(time.time() - t0, 2),
    }
    out = (
        (normal.permute(1, 2, 0) * 0.5 + 0.5).clamp(0, 1).cpu().numpy(),
        height.cpu().numpy(),
        roughness.clamp(0, 1).cpu().numpy(),
        metallic.clamp(0, 1).cpu().numpy(),
    )
    torch.cuda.empty_cache()
    return out, info


def _save_png(array: np.ndarray, stem: str, bits: int = 8) -> str:
    path = os.path.join(os.environ["GRADIO_TEMP_DIR"], f"{stem}-{time.time_ns()}.png")
    if bits == 16:
        Image.fromarray((array * 65535).round().astype(np.uint16)).save(path)
    else:
        Image.fromarray((array * 255).round().astype(np.uint8)).save(path, optimize=False, compress_level=6)
    return path


def prepare(image) -> Image.Image:
    """The picture as `compute` takes it: RGB, no larger than MAX_SIDE."""
    if image is None:
        raise gr.Error("Upload a picture of a surface.")
    if not isinstance(image, Image.Image):
        image = Image.open(image)
    image = image.convert("RGB")
    if max(image.size) > MAX_SIDE:
        scale = MAX_SIDE / max(image.size)
        image = image.resize((max(8, round(image.width * scale)), max(8, round(image.height * scale))), Image.LANCZOS)
    return image


def save(arrays, info: dict, directx: bool = False):
    """The maps as PNG files: normal (OpenGL unless `directx`), height
    (16-bit), roughness and metallic, and the JSON note."""
    normal, height, roughness, metallic = arrays
    if directx:
        normal = normal.copy()
        normal[..., 1] = 1 - normal[..., 1]
    info["convention"] = "directx" if directx else "opengl"
    info["size"] = [normal.shape[1], normal.shape[0]]
    files = (
        _save_png(normal, "normal"),
        _save_png(height, "height", bits=16),
        _save_png(roughness, "roughness"),
        _save_png(metallic, "metallic"),
    )
    _log("done", info)
    return (*files, info)