import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as func
from captum.attr import IntegratedGradients
import __main__


class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()

        # размер исходной картинки 180x180

        self.conv1 = nn.Conv2d(3, 8, 3, padding=1)
        self.batchnorm1 = nn.BatchNorm2d(8)
        self.pool1 = nn.MaxPool2d((2, 2))

        self.conv2 = nn.Conv2d(8, 16, 8, padding=1)
        self.dropout2 = nn.Dropout(0.25)

        self.batchnorm2 = nn.BatchNorm2d(16)
        self.pool2 = nn.MaxPool2d((2, 2))

        self.conv3 = nn.Conv2d(16, 32, 2, padding=1)
        self.dropout3 = nn.Dropout(0.25)
        self.batchnorm3 = nn.BatchNorm2d(32)
        self.pool3 = nn.MaxPool2d((2, 2))

        self.conv4 = nn.Conv2d(32, 16, 16, padding=1)
        self.dropout4 = nn.Dropout(0.25)
        self.batchnorm4 = nn.BatchNorm2d(16)

        # flatten
        self.flatten = nn.Flatten()

        self.fc_2_1 = nn.Linear(28224, 512)
        self.fc_2_2 = nn.Linear(512, 4)

        # linear 1
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 4)

    def forward(self, x):
        x = func.relu(self.conv1(x))
        x = self.batchnorm1(x)
        x = self.pool1(x)

        x = func.relu(self.conv2(x))
        x = self.dropout2(x)
        x = self.batchnorm2(x)
        x = self.pool2(x)

        x_1 = func.relu(self.conv3(x))
        x_1 = self.dropout3(x_1)
        x_1 = self.batchnorm3(x_1)
        x_1 = self.pool3(x_1)

        x_1 = func.relu(self.conv4(x_1))
        x_1 = self.dropout4(x_1)
        x_1 = self.batchnorm4(x_1)

        x_1 = self.flatten(x_1)
        x_1 = func.relu(self.fc1(x_1))
        x_1 = self.fc2(x_1)

        x_2 = self.flatten(x)
        x_2 = func.relu(self.fc_2_1(x_2))
        x_2 = self.fc_2_2(x_2)

        return x_1 + x_2


setattr(__main__, "ConvNet", ConvNet)

device = 'cpu'
model_ = torch.load('model_5_7_14_27_0.993125_final')
model_.eval()


def get_class_of_demension(idx):
    classes = ['NonDemented', 'VeryMildDemented', 'MildDemented', 'ModerateDemented']
    return classes[idx]


def get_segmented_map(image_attr: np.array,
                      color_map: str = 'positive',
                      borders: tuple = (20, 20)) -> np.array:
    """arg: color_map: [positive, all]"""
    if color_map != 'all':
        for i in range(len(image_attr)):
            for j in range(len(image_attr[i])):
                flag_zero = False
                if color_map == 'positive':
                    if max(image_attr[i][j]) != image_attr[i][j][1]:
                        flag_zero = True
                    else:
                        if sum(image_attr[i][j]) - max(image_attr[i][j]) > borders[1]:
                            flag_zero = True
                elif color_map == 'negative':
                    if max(image_attr[i][j]) == image_attr[i][j][1] or max(image_attr[i][j]) == image_attr[i][j][2]:
                        flag_zero = True
                    else:
                        if sum(image_attr[i][j]) - max(image_attr[i][j]) > borders[0]:
                            flag_zero = True
                if flag_zero:
                    image_attr[i][j] = [0, 0, 0]
    return image_attr


def show_pack_of_images(images, labels):
    f, axes = plt.subplots(1, len(images), figsize=(30, 5))
    for i, axis in enumerate(axes):
        img = images[i]
        axes[i].imshow(img)
        axes[i].set_title(labels[i])
    plt.show()


def create_color_map_igrad(net, img_path: str) -> tuple:
    integrated_gradients = IntegratedGradients(net)
    img = cv2.cvtColor(cv2.resize(cv2.imread(img_path, 0), (180, 180)), cv2.COLOR_GRAY2RGB)
    img_tensor = torch.from_numpy(np.array(img).astype(np.float32)).to('cpu')
    img_tensor = img_tensor.permute(2, 0, 1) / 255
    img_tensor = img_tensor.unsqueeze(0)

    output = model_(img_tensor)
    prob = func.sigmoid(output)
    probability = float(np.max(prob.detach().numpy()))
    prediction_score, pred_label_idx = torch.topk(output, 1)
    pred_label_idx.squeeze_()
    predicted_label = pred_label_idx.item()

    attributions_ig = integrated_gradients.attribute(img_tensor, target=pred_label_idx, n_steps=200)

    imgs = [(img_tensor.squeeze().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8),
            (np.transpose(attributions_ig.squeeze().cpu().detach().numpy(), (1, 2, 0)) * 255).astype(np.uint8)]
    imgs.extend([get_segmented_map(imgs[1].copy(), 'negative'), get_segmented_map(imgs[1].copy(), 'positive')])
    labels = [get_class_of_demension(predicted_label), 'all', 'negative', 'positive']

    return imgs, labels, probability


def get_results_model(image_path, model):
    images, labels, probability = create_color_map_igrad(model, image_path)

    img = images[3].copy()
    original = images[0].copy()

    result = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY);
    result = cv2.blur(result, (5, 5));

    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    ret, result = cv2.threshold(result, 0.3 * max_val, 255, cv2.THRESH_BINARY)

    contours, hierarchy = cv2.findContours(result, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

    for element in contours:
        if 150 > len(element) > 35:
            color = (255, 0, 0)
            x, y, w, h = cv2.boundingRect(element)
            cv2.rectangle(original, (x - 2, y - 2), (x + w + 1, y + h + 1), color, 1)

    return original, labels[0], probability