from pathlib import Path

import gradio as gr
import torch
from PIL import Image, ImageOps
from torchvision import transforms

from utils.models import Decoder, VGGEncoder
from utils.utils import adaptive_instance_normalization


BASE_DIR = Path(__file__).resolve().parent
VGG_PATH = BASE_DIR / "weights" / "vgg_normalised.pth"
DECODER_PATH = BASE_DIR / "weights" / "decoder_final.pth"
MAX_IMAGE_SIZE = 512

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_num_threads(min(4, max(1, torch.get_num_threads())))


def load_models():
    if not VGG_PATH.exists():
        raise FileNotFoundError(f"Missing VGG weights: {VGG_PATH}")
    if not DECODER_PATH.exists():
        raise FileNotFoundError(f"Missing decoder weights: {DECODER_PATH}")

    encoder = VGGEncoder(str(VGG_PATH)).to(DEVICE).eval()
    decoder = Decoder().to(DEVICE).eval()
    decoder_state = torch.load(str(DECODER_PATH), map_location=DEVICE)
    decoder.load_state_dict(decoder_state)
    return encoder, decoder


ENCODER, DECODER = load_models()


def prepare_image(image: Image.Image) -> torch.Tensor:
    image = ImageOps.exif_transpose(image).convert("RGB")
    image.thumbnail((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE), Image.Resampling.LANCZOS)
    return transforms.ToTensor()(image).unsqueeze(0).to(DEVICE)


def tensor_to_image(tensor: torch.Tensor) -> Image.Image:
    tensor = tensor.squeeze(0).detach().cpu().clamp(0, 1)
    return transforms.ToPILImage()(tensor)


def stylize(content_image, style_image, alpha):
    if content_image is None or style_image is None:
        raise gr.Error("Upload both a content image and a style image.")

    alpha = float(alpha)
    content = prepare_image(content_image)
    style = prepare_image(style_image)

    with torch.inference_mode():
        content_features = ENCODER(content, is_test=True)
        style_features = ENCODER(style, is_test=True)
        stylized_features = adaptive_instance_normalization(
            content_features,
            style_features,
        )
        blended_features = alpha * stylized_features + (1.0 - alpha) * content_features
        output = DECODER(blended_features)

    return tensor_to_image(output)


example_dir = BASE_DIR / "examples"
examples = [
    [
        str(example_dir / "content_lenna.jpg"),
        str(example_dir / "style_sketch.png"),
        1.0,
    ],
    [
        str(example_dir / "content_golden_gate.jpg"),
        str(example_dir / "style_la_muse.jpg"),
        0.8,
    ],
]


with gr.Blocks(title="AdaIN Neural Style Transfer") as demo:
    gr.Markdown("# AdaIN Neural Style Transfer")

    with gr.Row():
        content_input = gr.Image(type="pil", label="Content image")
        style_input = gr.Image(type="pil", label="Style image")

    alpha_input = gr.Slider(
        minimum=0.0,
        maximum=1.0,
        value=1.0,
        step=0.05,
        label="Style strength",
    )
    run_button = gr.Button("Transfer style", variant="primary")
    output_image = gr.Image(type="pil", label="Stylized output")

    run_button.click(
        fn=stylize,
        inputs=[content_input, style_input, alpha_input],
        outputs=output_image,
    )

    gr.Examples(
        examples=examples,
        inputs=[content_input, style_input, alpha_input],
        outputs=output_image,
        fn=stylize,
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch(ssr_mode=False)