"""StyleController: Continuous Image Stylization with Smooth Transitions.

A Gradio demo for controllable image stylization using StyleController,
which combines a LoRA adapter with strength-aware projectors on the
Qwen-Image-Edit-2509 backbone.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / diffusers / transformers
import copy
import bisect
import logging
import math

import numpy as np
import torch
from PIL import Image
from safetensors.torch import load_file
from scipy.interpolate import BSpline, make_interp_spline

import gradio as gr
from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage
from diffusers.image_processor import VaeImageProcessor
from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from transformers.models.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
from transformers.models.qwen2 import Qwen2Tokenizer

from models.pipeline_qwenimage_edit_plus import (
    QwenImageEditPlusPipeline,
    calculate_dimensions,
    VAE_IMAGE_SIZE,
)
from models.styctrl import register_styctrl, add_styctrl, activate
from models.transformer_qwenimage import QwenImageTransformer2DModel

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

BASE_MODEL = "Qwen/Qwen-Image-Edit-2509"
STYLECONTROLLER_REPO = "ReyChiaro/StyleController"
VARIANT = "qwenimage_edit_2509"

LORA_RANK = 128
LORA_LAYER_START = 0
LORA_LAYER_END = 59  # inclusive
TARGET_MODULES = [
    "attn.to_q",
    "attn.to_k",
    "attn.to_v",
    "attn.to_out.0",
    "attn.add_q_proj",
    "attn.add_k_proj",
    "attn.add_v_proj",
    "attn.to_add_out",
    "img_mlp.net.2",
    "txt_mlp.net.2",
]
ANCHOR_STRENGTHS = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
PROJECTOR_TYPE = "low_rank_linear"
BSPLINE_ORDER = 3
BSPLINE_MODE = "control"
DEFAULT_PROMPT = (
    "Apply the visual style of the second reference image to the first content image "
    "while preserving the content structure and layout."
)
DEFAULT_STEPS = 16
DEFAULT_SEED = 42


def _resize_to_content_aspect(content_image: Image.Image):
    """Size the generation to the CONTENT image's aspect ratio.

    The output aspect ratio must be dictated by the input (content) image, not
    the style reference. We snap the content aspect ratio to the closest model-
    supported resolution around the ideal inference area (~1 megapixel, on a
    32-pixel grid), center-crop the content image to exactly that aspect ratio,
    and resize it to those dimensions. Returns (cropped_content, width, height)
    so the caller can pass explicit width/height to the pipeline. Works for any
    arbitrary user-uploaded input resolution.
    """
    src_w, src_h = content_image.size
    target_w, target_h = calculate_dimensions(VAE_IMAGE_SIZE, src_w / src_h)

    # Center-crop the content image to the target aspect ratio, then resize.
    target_ratio = target_w / target_h
    src_ratio = src_w / src_h
    if src_ratio > target_ratio:
        # Source is wider than target: crop width.
        new_w = round(src_h * target_ratio)
        left = (src_w - new_w) // 2
        box = (left, 0, left + new_w, src_h)
    else:
        # Source is taller than target: crop height.
        new_h = round(src_w / target_ratio)
        top = (src_h - new_h) // 2
        box = (0, top, src_w, top + new_h)
    cropped = content_image.crop(box).resize((target_w, target_h), Image.LANCZOS)
    return cropped, target_w, target_h

# ---------------------------------------------------------------------------
# Download helper
# ---------------------------------------------------------------------------

from huggingface_hub import hf_hub_download


def _download_weights():
    """Download LoRA + projector checkpoints from the HF repo."""
    weights_dir = os.path.join(os.path.dirname(__file__), "weights")
    os.makedirs(weights_dir, exist_ok=True)

    lora_path = hf_hub_download(
        STYLECONTROLLER_REPO,
        f"{VARIANT}/styctrl_{VARIANT}.safetensors",
        repo_type="model",
    )
    projector_paths = []
    for i in range(1, 10):
        p = hf_hub_download(
            STYLECONTROLLER_REPO,
            f"{VARIANT}/projectors/s{i:02d}.safetensors",
            repo_type="model",
        )
        projector_paths.append(p)
    return lora_path, projector_paths


# ---------------------------------------------------------------------------
# Projector interpolation (ported from infer_styctrl.py)
# ---------------------------------------------------------------------------


def _projector_only(state, path):
    ps = {k: v for k, v in state.items() if ".projector." in k}
    if not ps:
        raise ValueError(f"No projector parameters found in {path}.")
    return ps


def _load_projector_checkpoints(paths):
    states = []
    expected_keys = None
    expected_shapes = {}
    for raw_path in paths:
        state = _projector_only(load_file(raw_path, device="cpu"), raw_path)
        keys = set(state)
        if expected_keys is None:
            expected_keys = keys
            expected_shapes = {k: v.shape for k, v in state.items()}
        elif keys != expected_keys:
            missing = sorted(expected_keys - keys)
            extra = sorted(keys - expected_keys)
            raise ValueError(
                f"Projector keys do not match for {raw_path}; missing={missing}, extra={extra}"
            )
        else:
            mismatched = [k for k, v in state.items() if v.shape != expected_shapes[k]]
            if mismatched:
                raise ValueError(f"Projector tensor shapes do not match for {raw_path}: {mismatched}")
        states.append(state)
    return states


def _zero_state_like(state):
    return {k: torch.zeros_like(v) for k, v in state.items()}


def _to_device_dtype(state, device, dtype):
    return {k: v.to(device=device, dtype=dtype) for k, v in state.items()}


def _open_clamped_knot_vector(strengths, degree):
    parameters = np.asarray(strengths, dtype=np.float64)
    interior = [
        parameters[i : i + degree].mean()
        for i in range(1, len(parameters) - degree)
    ]
    return np.concatenate(
        (
            np.repeat(parameters[0], degree + 1),
            np.asarray(interior, dtype=np.float64),
            np.repeat(parameters[-1], degree + 1),
        )
    )


def interpolate_projector_states(
    paths,
    anchor_strengths,
    query_strength,
    method="bspline",
    order=3,
    spline_mode="control",
    device="cuda",
    dtype=torch.bfloat16,
    endpoint_state=None,
):
    """Interpolate projector tensors at an explicit strength coordinate."""
    states = _load_projector_checkpoints(paths)
    if len(states) == 1:
        if np.isclose(query_strength, 0.0):
            return _to_device_dtype(_zero_state_like(states[0]), device, dtype)
        return _to_device_dtype(states[0], device, dtype)
    if len(anchor_strengths) != len(states):
        raise ValueError("anchor_strengths must contain exactly one value per projector checkpoint.")

    pairs = sorted(zip(anchor_strengths, states), key=lambda item: item[0])
    strengths = [float(item[0]) for item in pairs]
    states = [item[1] for item in pairs]
    if len(set(strengths)) != len(strengths):
        raise ValueError("Anchor strengths must be unique.")
    if strengths[0] < 0 or strengths[-1] > 1:
        raise ValueError("Anchor strengths must lie within [0, 1].")

    # A zero projector makes the StyCtrl LoRA contribution zero.
    if strengths[0] > 0:
        strengths.insert(0, 0.0)
        states.insert(0, _zero_state_like(states[0]))

    if endpoint_state is not None and strengths[-1] < 1.0:
        if endpoint_state.keys() != states[0].keys():
            raise ValueError("Endpoint projector keys do not match anchor checkpoints.")
        mismatched = [
            k for k, v in endpoint_state.items() if v.shape != states[0][k].shape
        ]
        if mismatched:
            raise ValueError(f"Endpoint projector tensor shapes do not match anchors: {mismatched}")
        strengths.append(1.0)
        states.append(endpoint_state)

    if query_strength < strengths[0] or query_strength > strengths[-1]:
        raise ValueError(
            f"Requested strength {query_strength} is outside the available projector range "
            f"[{strengths[0]}, {strengths[-1]}]."
        )

    if method != "bspline" or spline_mode == "interpolating":
        for anchor_strength, state in zip(strengths, states):
            if np.isclose(query_strength, anchor_strength):
                return _to_device_dtype(state, device, dtype)

    if method == "linear":
        upper = bisect.bisect_right(strengths, query_strength)
        lower = upper - 1
        x0, x1 = strengths[lower], strengths[upper]
        alpha = (query_strength - x0) / (x1 - x0)
        result = {
            k: (1.0 - alpha) * states[lower][k].float() + alpha * states[upper][k].float()
            for k in states[0]
        }
        return _to_device_dtype(result, device, dtype)

    if method != "bspline":
        raise ValueError(f"Unsupported interpolation method: {method}")
    if order < 1 or order >= len(strengths):
        raise ValueError(
            f"B-spline order must satisfy 1 <= order < number of anchors ({len(strengths)})."
        )

    x = np.asarray(strengths, dtype=np.float64)
    result = {}
    for key in states[0]:
        y = torch.stack([s[key].float() for s in states], dim=0).numpy()
        if spline_mode == "interpolating":
            value = make_interp_spline(x, y, k=order)(np.asarray(query_strength))
        elif spline_mode == "control":
            knots = _open_clamped_knot_vector(x, order)
            value = BSpline(knots, y, k=order, axis=0)(np.asarray(query_strength))
        else:
            raise ValueError(f"Unsupported spline mode: {spline_mode}")
        result[key] = torch.from_numpy(np.asarray(value))
    return _to_device_dtype(result, device, dtype)


# ---------------------------------------------------------------------------
# Model loading (module scope)
# ---------------------------------------------------------------------------

logger.info("Downloading StyleController weights ...")
_lora_path, _projector_paths = _download_weights()
logger.info("Weights downloaded.")

logger.info("Loading base pipeline components from %s ...", BASE_MODEL)
_dtype = torch.bfloat16

_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    BASE_MODEL, subfolder="scheduler", torch_dtype=_dtype
)
_scheduler_val = copy.deepcopy(_scheduler)

_vae = AutoencoderKLQwenImage.from_pretrained(
    BASE_MODEL, subfolder="vae", torch_dtype=_dtype
).to("cuda").requires_grad_(False)

_transformer = QwenImageTransformer2DModel.from_pretrained(
    BASE_MODEL, subfolder="transformer", torch_dtype=_dtype
).to("cuda").requires_grad_(False)

_tokenizer = Qwen2Tokenizer.from_pretrained(BASE_MODEL, subfolder="tokenizer")
_text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    BASE_MODEL, subfolder="text_encoder", torch_dtype=_dtype
).to("cuda").requires_grad_(False)
_processor = Qwen2_5_VLProcessor.from_pretrained(BASE_MODEL, subfolder="processor")

logger.info("Base pipeline components loaded.")

# Register StyCtrl LoRA
register_styctrl(
    model=_transformer,
    target_modules=TARGET_MODULES,
    device="cuda",
    dtype=_dtype,
)

_lora_layer_indices = list(range(LORA_LAYER_START, LORA_LAYER_END + 1))

add_styctrl(
    model=_transformer,
    rank=LORA_RANK,
    lora_layer_indices=_lora_layer_indices,
    adapter_name="styctrl",
    proj_type=PROJECTOR_TYPE,
    bias=False,
    device="cuda",
    dtype=_dtype,
)

# Load main LoRA weights (lora_A, lora_B) from checkpoint
_lora_state = load_file(_lora_path, device="cpu")
_lora_load_result = _transformer.load_state_dict(_lora_state, strict=False)
if _lora_load_result.unexpected_keys:
    logger.warning("Unexpected keys in LoRA checkpoint: %s", _lora_load_result.unexpected_keys[:5])
logger.info("LoRA loaded: %d tensors loaded, %d missing keys", len(_lora_state), len(_lora_load_result.missing_keys))
activate(_transformer, "styctrl")

# Pre-load projector states into CPU memory for fast interpolation
_projector_anchor_states = _load_projector_checkpoints(_projector_paths)

# Pre-compute the endpoint projector state (identity = from main LoRA)
_endpoint_projector_state = {
    k: v.detach().cpu()
    for k, v in _transformer.state_dict().items()
    if ".projector." in k
}

logger.info("StyCtrl LoRA registered and activated.")

# Build pipeline
_pipe = QwenImageEditPlusPipeline(
    transformer=_transformer,
    vae=_vae,
    scheduler=_scheduler_val,
    text_encoder=_text_encoder,
    tokenizer=_tokenizer,
    processor=_processor,
)

logger.info("Pipeline ready.")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


@spaces.GPU(duration=120, size="xlarge")
def stylize(
    content_image: Image.Image,
    style_image: Image.Image,
    strength: float = 0.55,
    seed: int = DEFAULT_SEED,
    steps: int = DEFAULT_STEPS,
    progress=gr.Progress(track_tqdm=True),
) -> Image.Image:
    """Stylize a content image using a style reference at a given strength.

    Args:
        content_image: The content image to stylize.
        style_image: The style reference image.
        strength: Style strength in [0, 1]. 0 = no style, 1 = full style.
        seed: Random seed for reproducibility.
        steps: Number of inference steps.
    """
    if content_image is None or style_image is None:
        return None

    logger.info("stylize called: strength=%.2f, seed=%d, steps=%d", strength, seed, steps)
    content_image = content_image.convert("RGB")
    style_image = style_image.convert("RGB")

    # The OUTPUT aspect ratio is dictated by the CONTENT image (cropped to the
    # closest supported aspect ratio at the model's ideal inference resolution),
    # NOT by the style reference. The style image keeps its own aspect ratio and
    # is resized adaptively inside the pipeline for conditioning only.
    content_image, out_w, out_h = _resize_to_content_aspect(content_image)
    logger.info("Content-driven output size: %dx%d (WxH)", out_w, out_h)

    # Interpolate projector states at the requested strength
    logger.info("Interpolating projector states ...")
    interp_states = interpolate_projector_states(
        paths=_projector_paths,
        anchor_strengths=ANCHOR_STRENGTHS,
        query_strength=strength,
        method="bspline",
        order=BSPLINE_ORDER,
        spline_mode=BSPLINE_MODE,
        device="cuda",
        dtype=_dtype,
        endpoint_state=_endpoint_projector_state,
    )
    logger.info("Loading interpolated states into transformer ...")
    _transformer.load_state_dict(interp_states, strict=False)
    for _, c in _pipe.components.items():
        if hasattr(c, "parameters"):
            for p in c.parameters():
                p.requires_grad_(False)

    w = torch.tensor([strength], device="cuda", dtype=_dtype)
    attention_kwargs = {
        "lora_layer_indices": _lora_layer_indices,
        "enable_lora": True,
        "w": w,
    }

    generator = torch.Generator(device="cuda").manual_seed(seed)

    logger.info("Running pipeline inference ...")
    output = _pipe(
        image=[content_image, style_image],
        prompt=DEFAULT_PROMPT,
        height=out_h,
        width=out_w,
        attention_kwargs=attention_kwargs,
        num_inference_steps=steps,
        generator=generator,
    )
    logger.info("Pipeline done, returning image.")
    return output.images[0]


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # 🎨 StyleController
        **Continuous Image Stylization with Smooth Transitions**

        Upload a content image and a style reference, then control the style
        strength with a slider. Built on Qwen-Image-Edit with a LoRA adapter
        and strength-aware projectors.

        [Paper](https://arxiv.org/abs/2608.08125) ·
        [GitHub](https://github.com/ReyChiaro/StyleController) ·
        [Model](https://huggingface.co/ReyChiaro/StyleController)
        """
    )

    with gr.Row():
        with gr.Column():
            content_input = gr.Image(
                label="Content Image",
                type="pil",
                height=400,
            )
            style_input = gr.Image(
                label="Style Reference Image",
                type="pil",
                height=400,
            )
            strength_slider = gr.Slider(
                label="Style Strength",
                minimum=0.0,
                maximum=1.0,
                step=0.05,
                value=0.55,
                info="0 = original content, 1 = full stylization",
            )
            run_btn = gr.Button("Stylize", variant="primary", scale=1)

        with gr.Column():
            output_image = gr.Image(
                label="Stylized Output",
                type="pil",
                height=400,
            )

    with gr.Accordion("Advanced Settings", open=False):
        seed_input = gr.Number(
            label="Seed",
            value=DEFAULT_SEED,
            precision=0,
            info="Random seed for reproducibility",
        )
        steps_input = gr.Slider(
            label="Inference Steps",
            minimum=4,
            maximum=50,
            step=1,
            value=DEFAULT_STEPS,
            info="More steps = higher quality but slower",
        )

    run_btn.click(
        fn=stylize,
        inputs=[content_input, style_input, strength_slider, seed_input, steps_input],
        outputs=output_image,
        api_name="stylize",
    )

    gr.Examples(
        examples=[
            ["examples/content_fox.jpg", "examples/style_oil_painting.jpg", 0.65],
            ["examples/content_astronaut.jpg", "examples/style_3d_render.jpg", 0.7],
            ["examples/content_owl.jpg", "examples/style_watercolor.jpg", 0.6],
            ["examples/content_cat.jpg", "examples/style_flat_vector.jpg", 0.7],
            ["examples/content_lighthouse.jpg", "examples/style_anime_illustration.jpg", 0.65],
        ],
        inputs=[content_input, style_input, strength_slider],
        outputs=output_image,
        fn=stylize,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True)