import random
from os import listdir
from os.path import join

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from tqdm import tqdm

from data.etis import EtisDataset
from data.hyperkvasir import KvasirSegmentationDataset
from models.segmentation_models import DeepLab
from evaluation.metrics import iou
from utils.logging import log_iou


class AdditiveNoise(nn.Module):
    def __init__(self, noise_factor):
        super().__init__()
        self.noise_factor = noise_factor

    def forward(self, img):
        return img + torch.randn(img.shape).to("cuda") * random.randint(0, 5) * self.noise_factor

    def __repr__(self):
        return self.__class__.__name__ + '(noise_factor={})'.format(self.noise_factor)


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

    def forward(self, img, mask):
        return img, mask


def apply_stressors(image, mask):
    img_augs = transforms.Compose([AdditiveNoise(0.05),
                                   transforms.RandomErasing(0.5, scale=(0.02, 0.08)),
                                   transforms.RandomInvert(p=0.1)])
    img_and_mask_augs = segmentation_stressors()
    image = img_augs(image)
    image, mask = img_and_mask_augs(image, mask)
    return image, mask


def stresstesttest():
    dataset = KvasirSegmentationDataset("Datasets/HyperKvasir")
    for x, y, fname in DataLoader(dataset):
        image, mask = apply_stressors(x.to("cuda"), y.to("cuda"))
        plt.imshow(image[0].permute(1, 2, 0).cpu().numpy())
        # plt.imshow(mask.cpu().numpy().squeeze().T, alpha=0.5)
        plt.show()
        input()


def perform_stresstest(modelpath, stressors=True):
    dataset = EtisDataset("Datasets/ETIS-LaribPolypDB")
    # dataset = KvasirSegmentationDataset("Datasets/HyperKvasir")
    for predictor_name in listdir(modelpath):
        if len(predictor_name.split("-")) == 3:
            model = DeepLab(1).to("cuda")
            model = torch.nn.Sequential(model, torch.nn.Sigmoid())

            model.eval()
            test = torch.load(join(modelpath, predictor_name))
            # print(test)
            model.load_state_dict(test)
            ious = torch.empty((0,))

            for x, y, fname in tqdm(DataLoader(dataset)):
                if stressors:
                    image, mask = apply_stressors(x.to("cuda"), y.to("cuda"))
                else:
                    image, mask = x.to("cuda"), y.to("cuda")
                with torch.no_grad():
                    output = model(image)
                    batch_ious = torch.Tensor([iou(output_i, mask_j) for output_i, mask_j in zip(output, mask)])
                    ious = torch.cat((ious, batch_ious.flatten()))
            log_iou("logs/kvasir-no-pretrain-baseline.log", -1, predictor_name.split("-")[-1], ious)


if __name__ == '__main__':
    perform_stresstest("Predictors/BaselineDeepLab", stressors=False)
    # EtisDataset("")