import matplotlib.pyplot as plt
import numpy as np
import torch.optim.optimizer
from torch.utils.data import DataLoader
from tqdm import tqdm
from data.hyperkvasir import KvasirSegmentationDataset, KvasirMNVset
from evaluation.metrics import iou
from losses.consistency_losses import *
from perturbation.model import ModelOfNaturalVariation
from training.vanilla_trainer import VanillaTrainer
from utils import logging
from models.ensembles import TrainedEnsemble


class ConsistencyTrainer(VanillaTrainer):
    def __init__(self, id, config):
        super(ConsistencyTrainer, self).__init__(id, config)
        self.criterion = ConsistencyLoss().to(self.device)
        self.nakedcloss = NakedConsistencyLoss()

    def train_epoch(self):
        self.model.train()
        losses = []
        for x, y, fname in self.train_loader:
            image = x.to("cuda")
            mask = y.to("cuda")
            aug_img, aug_mask = self.mnv(image, mask)
            self.optimizer.zero_grad()
            output = self.model(image)
            aug_output = self.model(aug_img)
            mean_iou = torch.mean(iou(output, mask))
            loss = self.criterion(aug_mask, mask, aug_output, output, mean_iou)
            loss.backward()
            self.optimizer.step()
            losses.append(np.abs(loss.item()))
        return np.mean(losses)

    def train(self):

        best_val_loss = 1000
        best_consistency = 0
        print("Starting Segmentation training")
        for i in range(self.epochs):
            training_loss = np.abs(self.train_epoch())
            val_loss, ious, closs = self.validate(epoch=i, plot=False)
            gen_ious = self.validate_generalizability(epoch=i, plot=False)
            mean_iou = float(torch.mean(ious))
            gen_iou = float(torch.mean(gen_ious))
            consistency = 1 - np.mean(closs)
            test_iou = np.mean(self.test().numpy())

            self.config["lr"] = [group['lr'] for group in self.optimizer.param_groups]
            logging.log_full(epoch=i, id=self.id, config=self.config, result_dict=
            {"train_loss": training_loss, "val_loss": val_loss,
             "iid_val_iou": mean_iou, "iid_test_iou": test_iou, "ood_iou": gen_iou,
             "consistency": consistency}, type="consistency")

            self.scheduler.step(i)
            # self.mnv.step()
            print(
                f"Epoch {i} of {self.epochs} \t"
                f" lr={[group['lr'] for group in self.optimizer.param_groups]} \t"
                f" loss={training_loss} \t"
                f" val_loss={val_loss} \t"
                f" ood_iou={gen_iou}\t"
                f" val_iou={mean_iou} \t"
                f" gen_prop={gen_iou / mean_iou} \t,"
                f" consistency={consistency}"
            )

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                np.save(
                    f"experiments/Data/Augmented-Pipelines/{self.model_str}/{self.id}",
                    test_iou)
                print(f"Saving new best model. IID test-set mean iou: {test_iou}")
                torch.save(self.model.state_dict(),
                           f"Predictors/Augmented/{self.model_str}/{self.id}")
                print("saved in: ", f"Predictors/Augmented/{self.model_str}/{self.id}")

            if consistency > best_consistency:
                best_consistency = consistency
                torch.save(self.model.state_dict(),
                           f"Predictors/Augmented/{self.model_str}/maximum_consistency{self.id}")
        torch.save(self.model.state_dict(),
                   f"Predictors/Augmented/{self.model_str}/{self.id}_last_epoch")

    def test(self):
        self.model.eval()
        ious = torch.empty((0,))
        with torch.no_grad():
            for x, y, fname in self.test_loader:
                image = x.to("cuda")
                mask = y.to("cuda")
                output = self.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()))
        return ious


class ConsistencyTrainerUsingAugmentation(ConsistencyTrainer):
    """
        Uses vanilla data augmentation with p=0.5 instead of a a custom loss
    """

    def __init__(self, id, config):
        super(ConsistencyTrainerUsingAugmentation, self).__init__(id, config)
        self.jaccard = vanilla_losses.JaccardLoss()
        self.prob = 0
        self.dataset = KvasirMNVset("Datasets/HyperKvasir", "train", inpaint=config["use_inpainter"])
        self.train_loader = DataLoader(self.dataset, batch_size=config["batch_size"], shuffle=True)

    def get_iou_weights(self, image, mask):
        self.model.eval()
        with torch.no_grad():
            output = self.model(image)
        return torch.mean(iou(output, mask))

    def get_consistency(self, image, mask, augmented, augmask):
        self.model.eval()
        with torch.no_grad():
            output = self.model(image)
        self.model.train()
        return torch.mean(self.nakedcloss(output, mask, augmented, augmask))

    def train_epoch(self):
        self.model.train()
        losses = []
        plotted = False
        for x, y, fname, flag in self.train_loader:
            image = x.to("cuda")
            mask = y.to("cuda")
            self.optimizer.zero_grad()
            output = self.model(image)
            loss = self.jaccard(output, mask)
            loss.backward()
            self.optimizer.step()
            losses.append(np.abs(loss.item()))
        return np.mean(losses)


class AdversarialConsistencyTrainer(ConsistencyTrainer):
    """
        Adversariall samples difficult
    """

    def __init__(self, id, config):
        super(ConsistencyTrainer, self).__init__(id, config)
        self.mnv = ModelOfNaturalVariation(T0=1).to(self.device)
        self.num_adv_samples = 10
        self.naked_closs = NakedConsistencyLoss()
        self.criterion = ConsistencyLoss().to(self.device)

    def sample_adversarial(self, image, mask, output):
        self.model.eval()
        aug_img, aug_mask = None, None  #
        max_severity = -10
        with torch.no_grad():
            for i in range(self.num_adv_samples):
                adv_aug_img, adv_aug_mask = self.mnv(image, mask)
                adv_aug_output = self.model(adv_aug_img)
                severity = self.naked_closs(adv_aug_mask, mask, adv_aug_output, output)

                if severity > max_severity:
                    max_severity = severity
                    aug_img = adv_aug_img
                    aug_mask = adv_aug_mask
        self.model.train()
        return aug_img, aug_mask

    def train_epoch(self):
        self.model.train()
        losses = []
        for x, y, fname in self.train_loader:
            image = x.to("cuda")
            mask = y.to("cuda")
            self.optimizer.zero_grad()
            output = self.model(image)
            aug_img, aug_mask = self.sample_adversarial(image, mask, output)
            # aug_img, aug_mask = self.mnv(image, mask)
            aug_output = self.model(aug_img)
            mean_iou = torch.mean(iou(output, mask))
            loss = self.criterion(aug_mask, mask, aug_output, output, mean_iou)
            loss.backward()
            self.optimizer.step()
            losses.append(np.abs(loss.item()))
        return np.mean(losses)


class StrictConsistencyTrainer(ConsistencyTrainer):
    def __init__(self, id, config):
        super(StrictConsistencyTrainer, self).__init__(id, config)
        self.criterion = StrictConsistencyLoss()


class ConsistencyTrainerUsingControlledAugmentation(ConsistencyTrainer):
    """
        Uses vanilla data augmentation with p=0.5 instead of a a custom loss and has two samples
    """

    def __init__(self, id, config):
        super(ConsistencyTrainerUsingControlledAugmentation, self).__init__(id, config)
        self.jaccard = vanilla_losses.JaccardLoss()
        self.mnv = ModelOfNaturalVariation(1)

    def train_epoch(self):
        self.model.train()
        losses = []
        plotted = False
        for x, y, fname in self.train_loader:
            image = x.to("cuda")
            mask = y.to("cuda")
            aug_img, aug_mask = self.mnv(image, mask)
            img_batch = torch.cat((image, aug_img))
            mask_batch = torch.cat((mask, aug_mask))
            self.optimizer.zero_grad()
            output = self.model(img_batch)
            loss = self.jaccard(output, mask_batch)
            loss.backward()
            self.optimizer.step()
            losses.append(np.abs(loss.item()))
        return np.mean(losses)


class EnsembleConsistencyTrainer(ConsistencyTrainer):
    def __init__(self, id, config):
        super(EnsembleConsistencyTrainer, self).__init__(id, config)
        self.model = TrainedEnsemble("Singular")
        self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr)
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=50, T_mult=2)