import gradio as gr
import torch
import torchvision.transforms.functional as TF
import torchvision.transforms as transforms
from src.model import Model
import os

device = "cuda" if torch.cuda.is_available() else "cpu"


def denorm_img(img: torch.Tensor):
    std = torch.Tensor([0.229, 0.224, 0.225]).reshape(-1, 1, 1)
    mean = torch.Tensor([0.485, 0.456, 0.406]).reshape(-1, 1, 1)
    return torch.clip(img * std + mean, min=0, max=1)


def main(inp1, inp2, alph, out_size=256):
    # print("inp1 ", inp1)
    # print("inp2 ", inp2)

    model = Model()
    model.load_state_dict(torch.load("./models/model_puddle.pt", map_location=torch.device(device)))
    model.eval()

    model.alpha = alph
    style = TF.to_tensor(inp1["composite"])
    content = TF.to_tensor(inp2["composite"])

    norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform = transforms.Compose([
        transforms.Resize(out_size, antialias=True)
    ])

    style, content = norm(style), norm(content)
    style, content = transform(style), transform(content)

    style, content = style.unsqueeze(0).to(device), content.unsqueeze(0).to(device)
    out = model(content, style)

    return denorm_img(out[0].detach()).permute(1, 2, 0).numpy()

def update_crop_size(crop_size):
    return gr.update(crop_size=(crop_size, crop_size))

with gr.Blocks() as demo:
    gr.Markdown("# Style Transfer with AdaIN")
    with gr.Row(variant="compact", equal_height=False):
        inp1 = gr.ImageEditor(
            type="pil",
            sources=["upload", "clipboard"],
            crop_size=(256, 256),
            eraser=False,
            brush=False,
            layers=False,
            label="Style",
            image_mode="RGB",
            transforms="crop",
            canvas_size=(512, 512)
        )
        inp2 = gr.ImageEditor(
            type="pil",
            sources=["upload", "clipboard"],
            crop_size=(256, 256),
            eraser=False,
            brush=False,
            layers=False,
            label="Content",
            image_mode="RGB",
            transforms="crop",
            canvas_size=(512, 512)
        )
        out = gr.Image(type="pil", label="Output")
    
    with gr.Row():
        out_size = gr.Dropdown(
            choices=[256, 512],
            value=256,
            multiselect=False,
            interactive=True,
            allow_custom_value=True,
            label="Output size",
            info="Size of the output image"
        )
        out_size.change(fn=update_crop_size, inputs=out_size, outputs=inp1)
        out_size.change(fn=update_crop_size, inputs=out_size, outputs=inp2)

        alph = gr.Slider(0, 1, value=1, label="Alpha", info="How much to change the original image", interactive=True, scale=3)

    with gr.Row():
        with gr.Column():
            gr.Markdown("## Style Examples")
            gr.Examples(
                examples=[
                    os.path.join(os.path.dirname(__file__), "data/styles/25.jpg"),
                    os.path.join(os.path.dirname(__file__), "data/styles/2272.jpg"),
                    os.path.join(os.path.dirname(__file__), "data/styles/2314.jpg"),
                ],
                inputs=inp1,
            )
        with gr.Column():
            gr.Markdown("## Content Examples")
            gr.Examples(
                examples=[
                    # os.path.join(os.path.dirname(__file__), "data/content/bear.jpg"),
                    os.path.join(os.path.dirname(__file__), "data/content/cat.jpg"),
                    os.path.join(os.path.dirname(__file__), "data/content/cow.jpg"),
                    os.path.join(os.path.dirname(__file__), "data/content/ducks.jpg"),
                ],
                inputs=inp2,
            )
    
    btn = gr.Button("Run")
    btn.click(fn=main, inputs=[inp1, inp2, alph, out_size], outputs=out)

demo.launch()