"""projtex PBR: albedo, roughness and metallic maps from one lit photo of an object.

Qwen-Image-Edit-2511 + the Lightning 4-step LoRA + a LoRA trained (DiffSynth-Studio) on synthetic Blender renders
of MatSynth materials, one instruction per map. API: /decompose (image, steps) -> (albedo, roughness, metallic)."""
import math
import os

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from safetensors.torch import load_file

try:
    import spaces
except ImportError:  # outside a ZeroGPU Space (an HF Job testing the same code)
    class spaces:  # noqa: N801
        @staticmethod
        def GPU(*a, **k):
            return (lambda f: f) if not a or not callable(a[0]) else a[0]

from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline

LORA_REPO = os.environ.get("PBR_LORA_REPO", "Mike0021/qwen-image-edit-2511-pbr-lora")
LORA_FILE = os.environ.get("PBR_LORA_FILE", "pbr2/step-4215.safetensors")
PROMPTS = {
    "albedo": "Turn this photo into its albedo map: the base colour of every surface, with all lighting, shadows "
              "and reflections removed.",
    "roughness": "Turn this photo into its roughness map: grayscale, black where the surface is glossy, white where "
                 "it is rough.",
    "metallic": "Turn this photo into its metallic map: white where the surface is bare metal, black everywhere else.",
}
SCHEDULER = {  # Qwen-Image-Lightning's sampling (exponential time shift, shift 3)
    "base_image_seq_len": 256, "base_shift": math.log(3), "invert_sigmas": False, "max_image_seq_len": 8192,
    "max_shift": math.log(3), "num_train_timesteps": 1000, "shift": 1.0, "shift_terminal": None,
    "stochastic_sampling": False, "time_shift_type": "exponential", "use_beta_sigmas": False,
    "use_dynamic_shifting": True, "use_exponential_sigmas": False, "use_karras_sigmas": False}


def pbr_lora_state_dict(path: str) -> dict:
    """DiffSynth saves PEFT keys of the bare transformer ('...to_q.lora_A.default.weight'); diffusers wants them
    under 'transformer.' without the adapter name."""
    sd = load_file(path)
    return {"transformer." + k.replace(".default", ""): v for k, v in sd.items()}


pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511", scheduler=FlowMatchEulerDiscreteScheduler.from_config(SCHEDULER),
    torch_dtype=torch.bfloat16)
pipe.load_lora_weights("lightx2v/Qwen-Image-Edit-2511-Lightning",
                       weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
                       adapter_name="lightning")
pipe.load_lora_weights(pbr_lora_state_dict(hf_hub_download(LORA_REPO, LORA_FILE, token=os.environ.get("HF_TOKEN"))),
                       adapter_name="pbr")
names, weights = ["lightning", "pbr"], [1.0, 1.0]
if os.environ.get("PBR_LORA_FILE2"):  # a blend of two PBR LoRAs, e.g. PBR_LORA_WEIGHTS="0.5,0.5"
    pipe.load_lora_weights(pbr_lora_state_dict(hf_hub_download(LORA_REPO, os.environ["PBR_LORA_FILE2"],
                                                               token=os.environ.get("HF_TOKEN"))), adapter_name="pbr2")
    w1, w2 = (float(x) for x in os.environ.get("PBR_LORA_WEIGHTS", "0.5,0.5").split(","))
    names, weights = ["lightning", "pbr", "pbr2"], [1.0, w1, w2]
pipe.set_adapters(names, adapter_weights=weights)
pipe.fuse_lora(adapter_names=names, lora_scale=1.0)
pipe.unload_lora_weights()
pipe.to("cuda")


def _square(img: Image.Image, side: int = 768):
    w, h = img.size
    s = max(w, h)
    canvas = Image.new("RGB", (s, s), img.getpixel((0, 0)))
    canvas.paste(img, ((s - w) // 2, (s - h) // 2))
    return canvas.resize((side, side), Image.LANCZOS), ((s - w) // 2, (s - h) // 2, w, h, s)


def _unsquare(img: Image.Image, box) -> Image.Image:
    x, y, w, h, s = box
    return img.resize((s, s), Image.LANCZOS).crop((x, y, x + w, y + h))


@spaces.GPU(duration=45, size="xlarge")
def decompose(image, steps: int = 4):
    """Albedo (sRGB), roughness and metallic (grey: 0 black .. 1 white) of the object in ``image``."""
    img = Image.open(image).convert("RGB") if isinstance(image, str) else image.convert("RGB")
    sq, box = _square(img)
    outs = []
    for key in ("albedo", "roughness", "metallic"):
        out = pipe(image=[sq], prompt=PROMPTS[key], negative_prompt=" ", num_inference_steps=int(steps),
                   true_cfg_scale=1.0, height=768, width=768,
                   generator=torch.Generator("cuda").manual_seed(0)).images[0]
        out = _unsquare(out.convert("RGB"), box)
        outs.append(out if key == "albedo" else out.convert("L"))
    return tuple(outs)


if __name__ == "__main__" or os.environ.get("SYSTEM") == "spaces":
    import gradio as gr

    with gr.Blocks(title="projtex PBR") as demo:
        gr.Markdown("# PBR maps from a photo of an object\nAlbedo, roughness and metallic in one call: "
                    "Qwen-Image-Edit-2511 + Lightning + a LoRA trained on renders of MatSynth materials "
                    "(used by the projection-texturing Space's `pbr` option).")
        with gr.Row():
            inp = gr.Image(type="filepath", label="Photo")
            steps = gr.Slider(2, 8, value=4, step=1, label="Steps")
        btn = gr.Button("Decompose")
        with gr.Row():
            a, r, m = gr.Image(label="Albedo"), gr.Image(label="Roughness"), gr.Image(label="Metallic")
        btn.click(decompose, [inp, steps], [a, r, m], api_name="decompose")
    demo.queue().launch()