from modules.comdel import cmodel
from modules.components import generator,discriminator,utils
from torch.nn import L1Loss, MSELoss
from torch.nn import functional as F
import torch
import itertools


class network(cmodel):

    def __init__(self, config):
        cmodel.__init__(self,config)

        self.config = config
        self.loss_names = ['D','G_GAN','IDT','TC']
        self.visual_names = ['REAL_As', 'FAKE_Bs']
        
        self.branch_num = config['branch_num']
        self._G_names = [f'G{n}'for n in range(1, self.branch_num + 1)]
        self._D_names = [f'D{n}'for n in range(1, self.branch_num + 1)]
        self.model_names = self._G_names + self._D_names

        print(self.model_names)

        for name in self._G_names:
            setattr(self, 'net' + name,
                generator.create_generator(
                    generator_type=config['generator_type'],
                    input_nc=int(config['input_nc']),
                    output_nc=int(config['output_nc']),
                    generator_channels_base=int(config['generator_channels_base']),
                    norm=config['norm'],
                    init_type=config['init_type'],
                    init_gain=float(config['init_gain']),
                    gpu_ids=config['gpu_ids'],
                    num_downs=config['num_downs'])
            )

        for name in self._D_names:
            setattr(self, 'net' + name, 
                discriminator.create_discriminator(input_nc=int(config['output_nc']),
                    discriminator_channels_base=int(config['discriminator_channels_base']),
                    norm=config['norm'],
                    init_type=config['init_type'],
                    init_gain=float(config['init_gain']),
                    gpu_ids=config['gpu_ids'])
            )

        if config['status'] == 'train':
            self.REAL_B_POOLS = [utils.ImagePool(config['pool_size']) for _ in range(self.branch_num)]

            self.criterionGAN = utils.GANLoss(config['gan_mode'],self.config).to(self.device)
            self.criterionIDT = L1Loss()
            self.criterionTC = L1Loss()

            self.optimizer_G = torch.optim.Adam(itertools.chain(*[getattr(self, 'net' + name).parameters() for name in self._G_names]),lr=float(config['lr_generator']),betas=(0.5, 0.999))
            self.optimizer_D = torch.optim.Adam(itertools.chain(*[getattr(self, 'net' + name).parameters() for name in self._D_names]),lr=float(config['lr_discriminator']),betas=(0.5, 0.999))

            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

        if config['status'] == 'test':
            self.classifier = torch.load('classifier.pkl').to(self.device)

        self.REAL_As = [None for _ in range(self.branch_num)]
        self.REAL_Bs = [None for _ in range(self.branch_num)]
        self.FAKE_Bs = [None for _ in range(self.branch_num)]
        self.IDTs = [None for _ in range(self.branch_num)]

        self.LOSS_D = None
        self.LOSS_G_GAN = None
        self.LOSS_G = None
        self.LOSS_IDT = None
        self.LOSS_TC = None


    def set_input(self, inputs):
        self.REAL_As = [inputs['A'].to(self.device) for _ in range(self.branch_num)]
        self.REAL_Bs = [inputs['B'].to(self.device) for _ in range(self.branch_num)]

    def _sum_tensor_list(self, tl):
        assert len(tl) > 0
        ret = tl[0]
        for i in tl[1:]:
            ret += i
        return ret

    def forward(self):
        if self.config['status'] == 'train':
            self.FAKE_Bs = [getattr(self, f'netG{i+1}')(self.REAL_As[i]) for i in range(self.branch_num)]

        if self.config['status'] == 'test':
            with torch.no_grad():
                self.test_result = []
                self.FAKE_Bs = [getattr(self, f'netG{i+1}')(self.REAL_As[i]) for i in range(self.branch_num)]

                FAKE_Bs_PROB = [F.softmax(self.classifier(self.FAKE_Bs[i]), dim=1)[0][1] for i in range(self.branch_num)]
                FAKE_B_INDEX = FAKE_Bs_PROB.index(max(FAKE_Bs_PROB))
                self.test_result.append(['FAKE_B', self.FAKE_Bs[FAKE_B_INDEX]])

    def backward_D(self):
        _REAL_Bs = [self.REAL_B_POOLS[i].query(self.REAL_Bs[i]) for i in range(self.branch_num)]
        _LOSS_D_REAL = self._sum_tensor_list([self.criterionGAN(getattr(self, f'netD{i+1}')(_REAL_Bs[i]), True) for i in range(self.branch_num)])
        _LOSS_D_FAKE = self._sum_tensor_list([self.criterionGAN(getattr(self, f'netD{i+1}')(self.FAKE_Bs[i].detach()), False) for i in range(self.branch_num)])
        self.LOSS_D = (_LOSS_D_REAL + _LOSS_D_FAKE) * self.config['d_loss_rate'] # 0.5
        self.LOSS_D.backward()


    def COMPUTE_TC_LOSS(self,FAKE1,FAKE2):
        LOSS_TC = self.criterionTC(FAKE1,FAKE2.clone().detach())
        LOSS_TC = LOSS_TC + self.criterionTC(FAKE2,FAKE1.clone().detach())
        LOSS_TC = LOSS_TC * self.config['tc_loss_rate'] # 20.0
        return LOSS_TC


    def backward_G(self):
        self.IDTs = [getattr(self, f'netG{i+1}')(self.REAL_Bs[i]) for i in range(self.branch_num)]
        
        self.LOSS_G_GAN = self._sum_tensor_list([self.criterionGAN(getattr(self, f'netD{i+1}')(self.FAKE_Bs[i]), True) for i in range(self.branch_num)])

        _LOSS_TCs = []
        for i in range(self.branch_num):
            for j in range(self.branch_num - 1):
                _FAKE1 = self.FAKE_Bs[i]
                _FAKE2 = self.FAKE_Bs[(i + j + 1) % self.branch_num]
                _LOSS_TCs.append(self.COMPUTE_TC_LOSS(_FAKE1, _FAKE2))
        self.LOSS_TC = self._sum_tensor_list(_LOSS_TCs)
        
        self.LOSS_IDT = self._sum_tensor_list([self.criterionIDT(self.IDTs[i], self.REAL_Bs[i]) * self.config['idt_loss_rate'] for i in range(self.branch_num)]) # 5.0

        self.LOSS_G = self.LOSS_G_GAN + self.LOSS_TC + self.LOSS_IDT
        self.LOSS_G.backward()


    def optimize_parameters(self):
        self.forward()
        
        self.set_requires_grad([getattr(self, f'netD{i+1}') for i in range(self.branch_num)],False)
        self.optimizer_G.zero_grad()
        self.backward_G()
        self.optimizer_G.step()

        self.set_requires_grad([getattr(self, f'netD{i+1}') for i in range(self.branch_num)],True)
        self.optimizer_D.zero_grad()
        self.backward_D()
        self.optimizer_D.step()
