import gradio as gr

import torch
from PIL import Image
import numpy as np
import random
from einops import rearrange
import matplotlib.pyplot as plt
from torchvision.transforms import v2

from model import MAE_ViT, MAE_Encoder, MAE_Decoder, MAE_Encoder_FeatureExtractor

path_1  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
path_2  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
path_3  = [['images/cat.jpg'], ['images/dog.jpg'], ['images/horse.jpg'], ['images/airplane.jpg'], ['images/truck.jpg']]
device = torch.device("cpu")

model_name = "model/no_mode/vit-t-mae-pretrain.pt"
model_no_mode = torch.load(model_name, map_location='cpu')
model_no_mode.eval()
model_no_mode.to(device)

model_name = "model/bottom_25/vit-t-mae-pretrain.pt"
model_pca_mode_bottom = torch.load(model_name, map_location='cpu')
model_pca_mode_bottom.eval()
model_pca_mode_bottom.to(device)

model_name = "model/top_75/vit-t-mae-pretrain.pt"
model_pca_mode_top = torch.load(model_name, map_location='cpu')
model_pca_mode_top.eval()
model_pca_mode_top.to(device)

transform = v2.Compose([
        v2.Resize((96, 96)),
        v2.ToTensor(), 
        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
    ])

# Load and Preprocess the Image
def load_image(image_path, transform):
    img = Image.open(image_path).convert('RGB')
    img = transform(img).unsqueeze(0)  # Add batch dimension
    return img

def show_image(img, title):
    img = rearrange(img, "c h w -> h w c")
    img = (img.cpu().detach().numpy() + 1) / 2  # Normalize to [0, 1]

    plt.imshow(img)
    plt.axis('off')
    plt.title(title)

# Visualize a Single Image
def visualize_single_image_no_mode(image_path):
    img = load_image(image_path, transform).to(device)
    
    # Run inference
    with torch.no_grad():
        predicted_img, mask = model_no_mode(img)
    
    # Convert the tensor back to a displayable image
    # masked image
    im_masked = img * (1 - mask)

    # MAE reconstruction pasted with visible patches
    im_paste = img * (1 - mask) + predicted_img * mask

    # remove the batch dimension
    img = img[0]
    im_masked = im_masked[0]
    predicted_img = predicted_img[0]
    im_paste = im_paste[0]
    
    # make the plt figure larger
    plt.figure(figsize=(18, 8))

    plt.subplot(1, 3, 1)
    show_image(img, "original")

    plt.subplot(1, 3, 2)
    show_image(im_masked, "masked")

    # plt.subplot(1, 4, 3)
    # show_image(predicted_img, "reconstruction")

    plt.subplot(1, 3, 3)
    show_image(im_paste, "reconstruction")

    plt.tight_layout()

    # convert the plt figure to a numpy array
    plt.savefig("output.png")

    return np.array(plt.imread("output.png"))

def visualize_single_image_pca_mode_bottom(image_path):
    img = load_image(image_path, transform).to(device)
    
    # Run inference
    with torch.no_grad():
        predicted_img, mask = model_pca_mode_bottom(img)
    
    # Convert the tensor back to a displayable image
    # masked image
    im_masked = img * (1 - mask)

    # MAE reconstruction pasted with visible patches
    im_paste = img * (1 - mask) + predicted_img * mask

    # remove the batch dimension
    img = img[0]
    im_masked = im_masked[0]
    predicted_img = predicted_img[0]
    im_paste = im_paste[0]
    
    # make the plt figure larger
    plt.figure(figsize=(18, 8))

    plt.subplot(1, 3, 1)
    show_image(img, "original")

    plt.subplot(1, 3, 2)
    show_image(im_masked, "masked")

    plt.subplot(1, 3, 3)
    show_image(predicted_img, "reconstruction")

    # plt.subplot(1, 4, 4)
    # show_image(im_paste, "reconstruction + visible")

    plt.tight_layout()

    # convert the plt figure to a numpy array
    plt.savefig("output.png")

    return np.array(plt.imread("output.png"))

def visualize_single_image_pca_mode_top(image_path):
    img = load_image(image_path, transform).to(device)
    
    # Run inference
    with torch.no_grad():
        predicted_img, mask = model_pca_mode_top(img)
    
    # Convert the tensor back to a displayable image
    # masked image
    im_masked = img * (1 - mask)

    # MAE reconstruction pasted with visible patches
    im_paste = img * (1 - mask) + predicted_img * mask

    # remove the batch dimension
    img = img[0]
    im_masked = im_masked[0]
    predicted_img = predicted_img[0]
    im_paste = im_paste[0]
    
    # make the plt figure larger
    plt.figure(figsize=(18, 8))

    plt.subplot(1, 3, 1)
    show_image(img, "original")

    plt.subplot(1, 3, 2)
    show_image(im_masked, "masked")

    plt.subplot(1, 3, 3)
    show_image(predicted_img, "reconstruction")

    # plt.subplot(1, 4, 4)
    # show_image(im_paste, "reconstruction + visible")

    plt.tight_layout()

    # convert the plt figure to a numpy array
    plt.savefig("output.png")

    return np.array(plt.imread("output.png"))

inputs_image_1 = [
    gr.components.Image(type="filepath", label="Input Image"),
]

outputs_image_1 = [
    gr.components.Image(type="numpy", label="Output Image"),
]

inputs_image_2 = [
    gr.components.Image(type="filepath", label="Input Image"),
]

outputs_image_2 = [
    gr.components.Image(type="numpy", label="Output Image"),
]

inputs_image_3 = [
    gr.components.Image(type="filepath", label="Input Image"),
]

outputs_image_3 = [
    gr.components.Image(type="numpy", label="Output Image"),
]


inference_no_mode = gr.Interface(
    fn=visualize_single_image_no_mode,
    inputs=inputs_image_1,
    outputs=outputs_image_1,
    examples=path_1,
    cache_examples = False,
    title="MAE-ViT Image Reconstruction",
    description="This is a demo of the MAE-ViT model for image reconstruction. The model is trained without PCA mode. It was trained on the STL-10 dataset. Check out the huggingface model card and the github repository for more information. https://huggingface.co/turhancan97/MAE-Models and https://github.com/turhancan97/Learning-by-Reconstruction-with-MAE",
)

inference_pca_mode_bottom = gr.Interface(
    fn=visualize_single_image_pca_mode_bottom,
    inputs=inputs_image_2,
    outputs=outputs_image_2,
    examples=path_2,
    title="MAE-ViT Image Reconstruction",
    description="This is a demo of the MAE-ViT model for image reconstruction. The model is trained with PCA mode (bottom 25%). It was trained on the STL-10 dataset.",
)

inference_pca_mode_top = gr.Interface(
    fn=visualize_single_image_pca_mode_top,
    inputs=inputs_image_3,
    outputs=outputs_image_3,
    examples=path_3,
    title="MAE-ViT Image Reconstruction",
    description="This is a demo of the MAE-ViT model for image reconstruction. The model is trained with PCA mode (top 75%). It was trained on the STL-10 dataset.",
)

gr.TabbedInterface(
    [inference_no_mode, inference_pca_mode_bottom, inference_pca_mode_top],
    tab_names=['Normal Mode', 'PCA Mode (Bottom 25%)', 'PCA Mode (Top 75%)']
).queue().launch()