import cv2
import numpy as np
import math
import torch
import random

from torch.utils.data import DataLoader

torch.manual_seed(12345)
random.seed(12345)
np.random.seed(12345)


def get_dataset_x(blank_image, filter_size=50, filter_stride=2):
    full_image_tensor = torch.tensor(blank_image).type(torch.FloatTensor).permute(2, 0, 1).unsqueeze(0)
    num_windows_h = math.floor((full_image_tensor.shape[2] - filter_size) / filter_stride) + 1
    num_windows_w = math.floor((full_image_tensor.shape[3] - filter_size) / filter_stride) + 1
    windows = torch.nn.functional.unfold(full_image_tensor, (filter_size, filter_size), stride=filter_stride).reshape(
        [1, 3, 50, 50, num_windows_h * num_windows_w]).permute([0, 4, 1, 2, 3]).squeeze()

    dataset_images = [windows[idx] for idx in range(len(windows))]
    dataset = list(dataset_images)
    return dataset


from torchvision.models.resnet import resnet50
from torchvision.models.resnet import ResNet50_Weights

print("Loading resnet...")
model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
hidden_state_size = model.fc.in_features
model.fc = torch.nn.Linear(in_features=hidden_state_size, out_features=2, bias=True)
model.to("cuda")
model.load_state_dict(torch.load("model_best_epoch_4_59.62.pth", map_location=torch.device("cuda")))
model.to("cuda")

import gradio as gr


def count_barnacles(input_img, progress=gr.Progress()):
    progress(0, desc="Loading Image")
    test_dataset = get_dataset_x(input_img)
    test_dataloader = DataLoader(test_dataset, batch_size=1024, shuffle=False)
    model.eval()
    predicted_labels_list = []
    for data in progress.tqdm(test_dataloader):
        with torch.no_grad():
            data = data.to("cuda")
            predicted_labels_list += [model(data)]
    predicted_labels = torch.cat(predicted_labels_list)
    x = int(math.sqrt(predicted_labels.shape[0]))
    predicted_labels = predicted_labels.reshape([x, x, 2]).detach()
    label_img = predicted_labels[:, :, :1].cuda().numpy()
    label_img -= label_img.min()
    label_img /= label_img.max()
    label_img = (label_img * 255).astype(np.uint8)
    mask = np.array(label_img > 180, np.uint8)
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    def extract_contour_center(cnt):
        M = cv2.moments(cnt)
        cx = int(M['m10'] / M['m00'])
        cy = int(M['m01'] / M['m00'])
        return cx, cy

    filter_width = 50
    filter_stride = 2

    def rev_window_transform(point):
        wx, wy = point
        x = int(filter_width / 2) + wx * filter_stride
        y = int(filter_width / 2) + wy * filter_stride
        return x, y

    nonempty_contours = filter(lambda cnt: cv2.contourArea(cnt) != 0, contours)
    windows = map(extract_contour_center, nonempty_contours)
    points = map(rev_window_transform, windows)

    blank_img_copy = input_img.copy()
    for x, y in points:
        blank_img_copy = cv2.circle(blank_img_copy, (x, y), radius=4, color=(255, 0, 0), thickness=-1)
    return blank_img_copy, int(len(list(points)))


demo = gr.Interface(count_barnacles, gr.Image(shape=(500, 500), type="numpy"),
                    outputs=["image", "number"],
                    examples="examples")
demo.queue(concurrency_count=10).launch()