import math
import cv2
import numpy as np
import torch
from diffusers import AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import CLIPVisionModelWithProjection
from preservation import PreservedPipeline, preservation_mask

BASE_MODEL = "OnomaAIResearch/Illustrious-XL-v2.0"
BASE_FILE = "Illustrious-XL-v2.0.safetensors"
BASE_CONFIG = "stabilityai/stable-diffusion-xl-base-1.0"
CONTROL_MODEL = "diffusers/controlnet-canny-sdxl-1.0"
ADAPTER_MODEL = "h94/IP-Adapter"
VAE_MODEL = "madebyollin/sdxl-vae-fp16-fix"
SIZE = 1024


def prepare_illustrious(input_path, boxes=(), text_mask=None):
    with Image.open(input_path) as source:
        original = source.convert("RGB")
    width, height = original.size
    scale = min(SIZE / width, SIZE / height)
    size = (round(width * scale), round(height * scale))
    x, y = (SIZE - size[0]) // 2, (SIZE - size[1]) // 2
    image = Image.new("RGB", (SIZE, SIZE))
    image.paste(original.resize(size, Image.Resampling.LANCZOS), (x, y))

    edges = cv2.Canny(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY), 100, 150)
    control = Image.fromarray(np.repeat(edges[:, :, None], 3, axis=2))
    sx, sy = size[0] / width, size[1] / height
    padded_mask = Image.new("L", (SIZE, SIZE))
    if text_mask is not None:
        padded_mask.paste(text_mask.resize(size, Image.Resampling.NEAREST), (x, y))
    mask = preservation_mask((SIZE, SIZE), [
        (x + math.floor(x1 * sx), y + math.floor(y1 * sy),
         x + math.ceil(x2 * sx), y + math.ceil(y2 * sy))
        for x1, y1, x2, y2 in boxes
    ], padded_mask)
    return image, control, mask, (x, y, x + size[0], y + size[1]), (width, height)


def load_illustrious(device):
    encoder = CLIPVisionModelWithProjection.from_pretrained(
        ADAPTER_MODEL, subfolder="models/image_encoder", torch_dtype=torch.float16,
    )
    pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_single_file(
        hf_hub_download(repo_id=BASE_MODEL, filename=BASE_FILE),
        config=BASE_CONFIG,
        controlnet=ControlNetModel.from_pretrained(CONTROL_MODEL, torch_dtype=torch.float16),
        image_encoder=encoder,
        vae=AutoencoderKL.from_pretrained(VAE_MODEL, torch_dtype=torch.float16),
        torch_dtype=torch.float16,
    )
    pipe.load_ip_adapter(
        ADAPTER_MODEL, subfolder="sdxl_models",
        weight_name="ip-adapter-plus_sdxl_vit-h.safetensors",
    )
    pipe.set_ip_adapter_scale(0.5)
    return pipe.to(device)


def infer_illustrious(pipe, image, control, mask, strength, steps):
    return PreservedPipeline(pipe, [mask])(
        prompt="masterpiece, best quality, anime illustration",
        negative_prompt="photorealistic, low quality, blurry",
        image=image,
        control_image=control,
        ip_adapter_image=image,
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=5.0,
    ).images[0]


def save_illustrious(result, crop, size, output_path):
    result.crop(crop).resize(size, Image.Resampling.LANCZOS).save(output_path)