import os

from torch.utils.data import DataLoader

from torchvision import transforms

import torchvision.utils as vutils

import torch

import torch.nn as nn

from data.data import InpaintingDataset, ToTensor

from model.net import InpaintingModel_GMCNN

from options.train_options import TrainOptions

from util.utils import getLatest

from data.hyperkvasir import KvasirSegmentationDataset

import torch.nn.functional as F

from inpainting.model.basemodel import BaseModel

from inpainting.model.basenet import BaseNet

from inpainting.model.loss import WGANLoss, IDMRFLoss

from inpainting.model.layer import init_weights, PureUpsampling, ConfidenceDrivenMaskLayer, SpectralNorm

from inpainting.model.net import *

config = {"model": "DeepLab",

          "device": "cuda",

          "lr": 0.00001,

          "batch_size": 8,

          "epochs": 250}

DATASET_PATH = ""

class gmcnn_inpainter_trainer:

    def __init__(self, config):

        self.model = InpaintingModel_GMCNN(in_channels=3, opt=config)

        self.dataset = KvasirSegmentationDataset(DATASET_PATH, augment=True)  # todo make inpainting dataset

        self.dataloader = DataLoader(self.dataset, batch_size=config["batch_size"], shuffle=True)

        self.device = config["device"]

        self.recloss = nn.L1Loss()

        self.aeloss = nn.L1Loss()

        self.confidence_mask_layer = ConfidenceDrivenMaskLayer()

        self.netGM = GMCNN(3, out_channels=3, cnum=32, act=F.elu, norm=F.instance_norm).cuda()

        init_weights(self.netGM)

        self.model_names = ['GM']

        self.netD = None

        self.optimizer_G = torch.optim.Adam(self.netGM.parameters(), lr=opt.lr, betas=(0.5, 0.9))

        self.optimizer_D = None

        self.wganloss = None

        self.recloss = nn.L1Loss()

        self.aeloss = nn.L1Loss()

        self.mrfloss = None

        # self.lambda_adv = opt.lambda_adv

        # self.lambda_rec = opt.lambda_rec

        # self.lambda_ae = opt.lambda_ae

        # self.lambda_gp = opt.lambda_gp

        # self.lambda_mrf = opt.lambda_mrf

        self.G_loss = None

        self.G_loss_reconstruction = None

        self.G_loss_mrf = None

        self.G_loss_adv, self.G_loss_adv_local = None, None

        self.G_loss_ae = None

        self.D_loss, self.D_loss_local = None, None

        self.GAN_loss = None

        self.gt, self.gt_local = None, None

        self.mask, self.mask_01 = None, None

        self.rect = None

        self.im_in, self.gin = None, None

        self.completed, self.completed_local = None, None

        self.completed_logit, self.completed_local_logit = None, None

        self.gt_logit, self.gt_local_logit = None, None

        self.pred = None

        self.netD = GlobalLocalDiscriminator(3, cnum=64, act=F.elu,

                                             spectral_norm=True,

                                             g_fc_channels=512 // 16 * 512 // 16 * 64 * 4,

                                             l_fc_channels=512 // 16 * 512 // 16 * 64 * 4).to(self.device)

        init_weights(self.netD)

        self.optimizer_D = torch.optim.Adam(filter(lambda x: x.requires_grad, self.netD.parameters()), lr=config["lr"],

                                            betas=(0.5, 0.9))

        self.wganloss = WGANLoss()

        self.mrfloss = IDMRFLoss()

    def train(self):

        for epoch in range(config["epochs"]):

            self.train_epoch(epoch)

            ret_loss = self.model.get_current_losses()

            self.model.save_networks(epoch + 1)

    def train_epoch(self, epoch):

        for img, mask, fname in self.dataloader:

            img, mask = img.to(self.device), mask.to(self.device)

            img_in = img * (1 - mask)

            self.gen_in = torch.cat((img_in, mask), 1)

            self.model.setInput(img, mask)

            self.model.optimize_parameters()

            self.pred = self.netGM(self.gin)

            self.completed = self.pred * self.mask_01 + self.gt * (1 - self.mask_01)

            self.completed_local = self.completed

            for i in range(5):  # train discriminator 5 times interleaved

                self.optimizer_D.zero_grad()

                self.optimizer_G.zero_grad()

                self.forward_D()

                self.backward_D()

                self.optimizer_D.step()

            self.optimizer_G.zero_grad()

            self.forward_G()

            self.backward_G()

            self.optimizer_G.step()

            # TODO come back here to finish gmcnn

    def forward_G(self):

        self.G_loss_reconstruction = self.recloss(self.completed * self.mask, self.gt.detach() * self.mask)

        self.G_loss_reconstruction = self.G_loss_reconstruction / torch.mean(self.mask_01)

        self.G_loss_ae = self.aeloss(self.pred * (1 - self.mask_01), self.gt.detach() * (1 - self.mask_01))

        self.G_loss_ae = self.G_loss_ae / torch.mean(1 - self.mask_01)

        self.G_loss = self.lambda_rec * self.G_loss_reconstruction + self.lambda_ae * self.G_loss_ae

        self.completed_logit, self.completed_local_logit = self.netD(self.completed, self.completed_local)

        self.G_loss_mrf = self.mrfloss((self.completed_local + 1) / 2.0, (self.gt_local.detach() + 1) / 2.0)

        self.G_loss = self.G_loss + self.lambda_mrf * self.G_loss_mrf

        self.G_loss_adv = -self.completed_logit.mean()

        self.G_loss_adv_local = -self.completed_local_logit.mean()

        self.G_loss = self.G_loss + self.lambda_adv * (self.G_loss_adv + self.G_loss_adv_local)

    def forward_D(self):

        self.completed_logit, self.completed_local_logit = self.netD(self.completed.detach(),

                                                                     self.completed_local.detach())

        self.gt_logit, self.gt_local_logit = self.netD(self.gt, self.gt_local)

        # hinge loss

        self.D_loss_local = nn.ReLU()(1.0 - self.gt_local_logit).mean() + nn.ReLU()(

            1.0 + self.completed_local_logit).mean()

        self.D_loss = nn.ReLU()(1.0 - self.gt_logit).mean() + nn.ReLU()(1.0 + self.completed_logit).mean()

        self.D_loss = self.D_loss + self.D_loss_local

    def backward_G(self):

        self.G_loss.backward()

    def backward_D(self):

        self.D_loss.backward(retain_graph=True)