import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data.dataloader import DataLoader

# from perturbation.gmcnn.model.net_with_dropout import InpaintingModel_GMCNN_Given_Mask
from data.hyperkvasir import *
from models.inpainters import SegGenerator
from perturbation.perturbator import RandomDraw


class Inpainter:
    # wrapper around gmcnn
    def __init__(self, path_to_state_dict):
        super(Inpainter, self).__init__()
        self.config = None
        self.model = SegGenerator()
        self.model.load_state_dict(torch.load(path_to_state_dict))
        self.model.to("cuda")
        self.model.eval()
        self.perturbator = RandomDraw()

    def forward(self, img, mask, masked_image=None):
        mask = mask.unsqueeze(1)
        masked_image = img * (1 - mask)
        polyp = self.model(masked_image)
        merged = (1 - mask) * img + (polyp * mask)
        return merged, polyp

    def add_polyp(self, img, old_mask_a):
        new_polyp_mask = np.expand_dims(self.perturbator(rad=0.25), -1)
        old_mask = np.expand_dims(old_mask_a, -1)
        total_mask = np.clip(new_polyp_mask + old_mask, 0, 1)
        masked = img * (1 - new_polyp_mask)
        with torch.no_grad():
            polyp = self.model(torch.Tensor(masked).to("cuda").T.unsqueeze(0))

        cpu_polyp = polyp.cpu().squeeze(0).T.numpy()
        inpainted = masked + (cpu_polyp * new_polyp_mask)
        return inpainted.astype(np.float32), total_mask.astype(np.float32)

    def get_test(self, split="test"):
        for i, (image, mask, masked_image, part, fname) in enumerate(
                DataLoader(KvasirSyntheticDataset("Datasets/HyperKvasir", split="test"),
                           batch_size=4)):
            merged, mask = self.add_polyp(image, mask)
            plt.title("Inpainted image")
            plt.imshow(merged[0].T)
            plt.show()
            # plt.savefig(f"perturbation/inpaint_examples/{i}")


if __name__ == '__main__':
    inpainter = Inpainter("Predictors/Inpainters/no-pretrain-deeplab-generator-4990")
    inpainter.get_test()