import torch
import torch.optim as optim
import lightning.pytorch as pl
from lightning.pytorch.tuner import Tuner
from tqdm import tqdm
from torch.optim.lr_scheduler import OneCycleLR
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import albumentations as A
from pytorch_grad_cam.utils.image import show_cam_on_image
from albumentations.pytorch import ToTensorV2
import config
from yolo_lightning import YOLOv3Lightning
import torch
import cv2
import numpy as np
import gradio as gr
import os
from utils_app import *

model = YOLOv3Lightning (config)

model.load_state_dict(torch.load("custom_yolo_model.pth",  map_location=torch.device('cpu')), strict=False)
model.setup(stage="test")

IMAGE_SIZE = 416

ANCHORS = [
    [(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
    [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
    [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
]  # Note these have been rescaled to be between [0, 1]
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]

scaled_anchors = (
    torch.tensor(config.ANCHORS)
    * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
)

def process_image_and_plot(image, model, scaled_anchors):

    transformed_image = config.transforms(image=image)["image"].unsqueeze(0)
    output = model(transformed_image)
    bboxes = [[] for _ in range(1)]

    for i in range(3):
        batch_size, A, S, _, _ = output[i].shape
        anchor = scaled_anchors[i]
        boxes_scale_i = cells_to_bboxes(output[i], anchor, S=S, is_preds=True)
        for idx, box in enumerate(boxes_scale_i):
            bboxes[idx] += box

    nms_boxes = non_max_suppression(
        bboxes[0], iou_threshold=0.5, threshold=0.4, box_format="midpoint",
    )
    fig = plot_image(transformed_image[0].permute(1, 2, 0), nms_boxes)

    cam = YoloCAM(model=model, target_layers=[model.model.layers[-2]], use_cuda=False)
    grayscale_cam = cam(transformed_image, scaled_anchors)[0, :, :]
    img = cv2.resize(image, (416, 416))
    img = np.float32(img) / 255
    cam_image = show_cam_on_image(img, grayscale_cam, use_rgb=True)

    return fig,cam_image


examples = [
    ["images/2012_004288.jpg"],
    ["images/2012_004314.jpg"],
    ["images/car.jpg"],
]
def processed_image(image):
    figure,gradcam = process_image_and_plot(image, model, scaled_anchors)
    return figure,gradcam

title = "YoloV3 on Pascal VOC Dataset (GradCAM)"
description = f"Pytorch Implemetation of YoloV3 trained from scratch on Pascal VOC dataset with GradCAM \n Class in pascol voc: {', '.join(config.PASCAL_CLASSES)}"
demo = gr.Interface(processed_image,
    inputs=[
        gr.Image(label="Input Image"),
    ],
    outputs=[gr.Plot(),gr.Image(shape=(32, 32), label="GradCAM Prediction")],
    title=title,
    description=description,
   examples=examples,
)
demo.launch()