import matplotlib.pyplot as plt

from models.segmentation_models import *
from data.hyperkvasir import KvasirSegmentationDataset
from torch.utils.data import DataLoader
import torch
import torch.nn as nn
import copy


class SplicedReconstructor(nn.Module):
    def __init__(self):
        super(SplicedReconstructor, self).__init__()
        inductivenet = InductiveNet()
        inductivenet.load_state_dict(torch.load("Predictors/Augmented/InductiveNet/consistency_1"))
        self.decoder = copy.deepcopy(inductivenet.reconstruction_decoder)
        self.head = copy.deepcopy(inductivenet.reconstruction_head)
        del inductivenet
        deeplab = DeepLab()
        deeplab.load_state_dict(torch.load("Predictors/Augmented/DeepLab/consistency_1"))
        self.encoder = copy.deepcopy(deeplab.encoder)
        del deeplab

    def predict(self, x):
        features = self.encoder(x)
        reconstructor_output = self.decoder(*features)
        reconstructed = self.head(reconstructor_output)
        return reconstructed


if __name__ == '__main__':
    model = SplicedReconstructor().to("cuda").eval()

    for x, y, _ in DataLoader(KvasirSegmentationDataset("Datasets/HyperKvasir/", "test")):
        with torch.no_grad():
            reconstruction = model.predict(x.to("cuda")).cpu()
        fig, ax = plt.subplots(ncols=1, nrows=2, sharey=True, sharex=True, figsize=(2, 1), dpi=1000)
        fig.subplots_adjust(wspace=0, hspace=0)
        ax[0].imshow(reconstruction[0].T)
        ax[1].imshow(x[0].T)
        plt.show()
        print("Showing...")