from model_utils import *
import pickle
from tqdm import tqdm


class StyleEmdedding(torch.nn.Module):
    """
    AAAI 2019 model 2
    """

    def __init__(self, args, reny_dataloader):
        super(StyleEmdedding, self).__init__()
        self.args = args
        self.training_all_except_encoder = True
        self.reny_dataloader = reny_dataloader
        self.encoder = EncoderRNN(args, args.number_of_tokens, args.hidden_dim, args.number_of_layers)
        self.decoder = DecoderRNN(args, args.style_dim + args.content_dim, args.number_of_tokens, args.number_of_layers)
        self.loss = torch.nn.NLLLoss(ignore_index=self.args.tokenizer.pad_token_id)
        self.loss_classif = torch.nn.NLLLoss()
        if self.args.complex_proj_content:
            self.proj_style = nn.Sequential(nn.Linear(1, args.style_dim), nn.LeakyReLU(),
                                            nn.Linear(args.style_dim, args.style_dim), nn.LeakyReLU(),  # emb labels
                                            nn.Linear(args.style_dim, args.style_dim))  # emb labels
            self.proj_content = nn.Sequential(nn.Linear(args.hidden_dim, args.hidden_dim), nn.LeakyReLU(),
                                              nn.Linear(args.hidden_dim, args.content_dim), nn.LeakyReLU(),
                                              nn.Linear(args.hidden_dim, args.content_dim))
        else:
            self.proj_style = nn.Sequential(nn.Linear(1, args.style_dim), nn.LeakyReLU(),
                                            nn.Linear(args.style_dim, args.style_dim))  # emb labels
            self.proj_content = nn.Sequential(nn.Linear(args.hidden_dim, args.hidden_dim), nn.LeakyReLU(),
                                              nn.Linear(args.hidden_dim, args.content_dim))

        # D_\gamma
        self.d_gamma = ClassifierGamma(args.content_dim + args.number_of_styles,
                                       args.number_of_styles)

        # Style classifier
        self.style_classifier = Classifier(args.content_dim, args.number_of_styles, args.use_complex_classifier)

        # Loss : paper multipliers
        self.mul_mi = self.args.mul_mi

    def forward(self, input_tensor, labels, teacher_ratio):
        ###############################
        # Update style classifier loss:
        ###############################
        loss_gen, loss_mi, reny, loss_gamma, loss_style_classifier, loss_gen_reny, loss_h_sz, loss_h_s, loss_gen_reny = torch.tensor(
            0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(self.args.device)
        gradient_encoder, gradient_decoder, gradient_content_proj, gradient_style_proj, gradient_reny = torch.tensor(
            0.0).to(self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(self.args.device)

        reny_pos_class, reny_neg_class, reny_one_labels = torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device)
        pos_class, neg_class, one_labels = torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device)
        count = 0
        # epoch_iterator = tqdm(self.reny_dataloader, desc="Reny Training")
        if self.training_all_except_encoder:
            for step, batch in enumerate(self.reny_dataloader):
                if count == self.args.reny_training + 1:
                    break
                count += 1
                inputs_reny = batch['line'].to(self.args.device)
                labels_reny = batch['label'].to(self.args.device)
                free_params(self.style_classifier)
                frozen_params(self.d_gamma)
                frozen_params(self.proj_style)
                frozen_params(self.proj_content)
                frozen_params(self.encoder)
                frozen_params(self.decoder)

                encoder_hidden = self.encoder.initHidden()
                encoder_output, encoder_hidden = self.encoder(inputs_reny.clone(), encoder_hidden)

                # Content space
                content = self.proj_content(encoder_hidden)
                style_content_pred = self.style_classifier(content)
                loss_style_classifier = self.loss_classif(style_content_pred, labels_reny)

                reny_pos_class = torch.sum(torch.exp(style_content_pred)[:, 1])
                reny_neg_class = torch.sum(torch.exp(style_content_pred)[:, 0])
                reny_one_labels = torch.sum(labels_reny)
                if self.training:
                    loss_style_classifier.backward()
                    torch.nn.utils.clip_grad_norm_(self.style_classifier.parameters(), self.args.max_grad_norm)
                    self.args.optimizer.step()
                    self.style_classifier.zero_grad()

                ################
                # Update D_gamma
                ################
                if not self.args.no_minimization_of_mi_training:
                    loss_gamma = 0
                    free_params(self.d_gamma)
                    frozen_params(self.style_classifier)
                    frozen_params(self.proj_style)
                    frozen_params(self.proj_content)
                    frozen_params(self.encoder)
                    frozen_params(self.decoder)

                    inputs_reny_dgamma = inputs_reny.clone()
                    encoder_hidden = self.encoder.initHidden()
                    encoder_output, encoder_hidden = self.encoder(inputs_reny_dgamma, encoder_hidden)

                    # Content space
                    content = self.proj_content(encoder_hidden)

                    style_content_pred = self.style_classifier(content)
                    label_content_pred = style_content_pred.topk(1, dim=-1)[-1].squeeze(-1)

                    if self.args.use_complex_gamma_training:
                        raise NotImplementedError

                    else:
                        label_v = torch.tensor(
                            [[1., 0.] if el == 0 else [0., 1.] for el in label_content_pred.tolist()]).to(
                            self.args.device)
                        label_u = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in labels_reny.tolist()]).to(
                            self.args.device)

                        u = torch.cat([content, label_u.unsqueeze(0).float().repeat(4, 1, 1)],
                                      dim=-1)  # 4 is for bid + 2 layers
                        v = torch.cat([content, label_v.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)

                        d_gamma_content_pred_u = self.d_gamma(u)  # - log
                        d_gamma_content_pred_v = self.d_gamma(v)  # - log # TODO :)

                        loss_gamma = - torch.mean(d_gamma_content_pred_u[:, 0]) / 2 - torch.mean(
                            d_gamma_content_pred_v[:, 1]) / 2

                    if self.training:
                        loss_gamma.backward()
                        gradient_reny += comput_gradient_norm(self.d_gamma)
                        torch.nn.utils.clip_grad_norm_(self.d_gamma.parameters(), self.args.max_grad_norm)
                        self.args.optimizer.step()
                        self.d_gamma.zero_grad()

                ################
                # Update Decoder
                ################
                free_params(self.decoder)
                frozen_params(self.d_gamma)
                frozen_params(self.style_classifier)
                frozen_params(self.proj_style)
                frozen_params(self.proj_content)
                frozen_params(self.encoder)

                input_tensor_reny_gen_golden = inputs_reny.clone()
                inputs_reny_gen = inputs_reny.clone()
                if self.args.add_noise:
                    inputs_reny_gen = corrupt_input(self, inputs_reny_gen)

                encoder_hidden = self.encoder.initHidden()
                encoder_output, encoder_hidden = self.encoder(inputs_reny_gen, encoder_hidden)

                # Style space
                style = self.proj_style(labels_reny.float().unsqueeze(1))
                style = torch.cat([style.unsqueeze(0) for _ in range(4)])
                # Content space
                content = self.proj_content(encoder_hidden)

                # Concatenate style and other
                encoder_hidden = torch.cat([style, content], dim=-1)

                # Sentence Generation
                loss_gen_reny = 0
                decoder_input = torch.ones(self.args.batch_size, 1).to(
                    self.args.device).long() * self.args.tokenizer.sep_token_id
                decoder_hidden = encoder_hidden
                use_teacher_forcing = True if random.random() < teacher_ratio else False

                decoded_words = [decoder_input]

                if use_teacher_forcing and self.training:
                    # Teacher forcing: Feed the target as the next input
                    for di in range(self.args.max_length):
                        decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden)
                        loss_gen_reny += self.loss(decoder_output.squeeze(1),
                                                   input_tensor_reny_gen_golden[:, di]) / self.args.max_length
                        decoder_input = input_tensor_reny_gen_golden[:, di].unsqueeze(-1)  # Teacher forcing

                else:
                    # Without teacher forcing: use its own predictions as the next input
                    for di in range(self.args.max_length):
                        decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden)
                        topv, topi = decoder_output.topk(1)
                        decoder_input = topi.squeeze().detach().unsqueeze(-1)  # detach from history as input
                        decoded_words.append(topi.squeeze(-1))
                        loss_gen_reny += self.loss(decoder_output.squeeze(1),
                                                   input_tensor_reny_gen_golden[:, di]) / self.args.max_length

                if self.training:
                    loss_gen_reny.backward()
                    torch.nn.utils.clip_grad_norm_(self.decoder.parameters(), self.args.max_grad_norm)
                    self.args.optimizer.step()
                    self.decoder.zero_grad()

        ###############################
        # Update Genloss + \lambda * MI
        ###############################
        input_tensor_gen_golden = input_tensor.clone()
        input_tensor_gen = input_tensor.clone()
        if self.args.add_noise:
            input_tensor_gen = corrupt_input(self, input_tensor_gen)
        input_tensor_mi = input_tensor_gen.clone()
        if self.training:
            frozen_params(self.d_gamma)
            frozen_params(self.style_classifier)
            frozen_params(self.decoder)
            free_params(self.proj_style)
            free_params(self.proj_content)
            free_params(self.encoder)

        # H(S)
        encoder_hidden_mi = self.encoder.initHidden()
        encoder_output_mi, encoder_hidden_mi = self.encoder(input_tensor_mi, encoder_hidden_mi)

        # Content space
        content_mi = self.proj_content(encoder_hidden_mi)
        style_content_pred_mi = self.style_classifier(content_mi)

        if not self.args.no_minimization_of_mi_training:
            len_ = len(labels.tolist())
            loss_h_s = (-1) * torch.log(torch.sum(torch.exp(style_content_pred_mi)[:, 0]) / len_) * (
                    1 - torch.sum(labels).item() / len_) - torch.log(
                torch.sum(torch.exp(style_content_pred_mi)[:, 1]) / len_) * (torch.sum(labels).item() / len_)
            pos_class = torch.sum(torch.exp(style_content_pred_mi)[:, 1])
            neg_class = torch.sum(torch.exp(style_content_pred_mi)[:, 0])
            one_labels = torch.sum(labels)
            # loss_h_s = - torch.mean(style_content_pred_mi[:, 0]) * (1 - torch.mean(labels.float()).item()) - torch.mean(
            # style_content_pred_mi[:, 1]) * torch.mean(labels.float()).item()  # we use Jensen
            # print(loss_h_s)
            # Compute Reny
            label_content_pred = style_content_pred_mi.topk(1, dim=-1)[-1].squeeze(-1)  # TODO ?
            label_v = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in label_content_pred.tolist()]).to(
                self.args.device)
            label_u = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in labels.tolist()]).to(self.args.device)
            u = torch.cat([content_mi, label_u.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)
            v = torch.cat([content_mi, label_v.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)
            d_gamma_content_pred_u = self.d_gamma(u)
            d_gamma_content_pred_v = self.d_gamma(v)
            R = torch.mean((torch.exp(d_gamma_content_pred_u[:, 0]) / torch.exp(d_gamma_content_pred_v[:, 1])) ** (
                    self.args.alpha - 1))

            # Remove biais from gradients
            reny = torch.abs(torch.log(R) / (self.args.alpha - 1))

            loss_h_sz = self.loss_classif(style_content_pred_mi, labels)

            # MI
            if self.args.no_reny:
                loss_mi = self.args.mul_mi * torch.abs(loss_h_s - loss_h_sz)
            else:
                loss_mi = self.args.mul_mi * torch.abs(loss_h_s - loss_h_sz + reny)
        else:
            loss_mi = - self.args.mul_mi * self.loss_classif(style_content_pred_mi, labels)

        #########################
        ######## Genloss ########
        #########################
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(input_tensor_gen, encoder_hidden)

        # Style space
        style = self.proj_style(labels.float().unsqueeze(1))
        style = torch.cat([style.unsqueeze(0) for _ in range(4)])
        # Content space
        content = self.proj_content(encoder_hidden)

        # Concatenate style and other
        encoder_hidden = torch.cat([style, content], dim=-1)

        # Sentence Generation
        loss_gen = 0
        decoder_input = torch.ones(self.args.batch_size, 1).to(
            self.args.device).long() * self.args.tokenizer.sep_token_id
        decoder_hidden = encoder_hidden
        use_teacher_forcing = True if random.random() < teacher_ratio else False

        decoded_words = [decoder_input]

        if use_teacher_forcing and self.training:
            # Teacher forcing: Feed the target as the next input
            for di in range(self.args.max_length):
                decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden)
                loss_gen += self.loss(decoder_output.squeeze(1), input_tensor_gen_golden[:, di]) / self.args.max_length
                decoder_input = input_tensor_gen_golden[:, di].unsqueeze(-1)  # Teacher forcing

        else:
            # Without teacher forcing: use its own predictions as the next input
            for di in range(self.args.max_length):
                decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden)
                topv, topi = decoder_output.topk(1)
                decoder_input = topi.squeeze().detach().unsqueeze(-1)  # detach from history as input
                decoded_words.append(topi.squeeze(-1))
                loss_gen += self.loss(decoder_output.squeeze(1), input_tensor_gen_golden[:, di]) / self.args.max_length

        # Compute All Losses for MI
        if self.training:
            loss = torch.abs(loss_gen + loss_mi)
            loss.backward()
            gradient_encoder, gradient_decoder, gradient_content_proj, gradient_style_proj = comput_gradient_norm(
                self.encoder), comput_gradient_norm(self.decoder), comput_gradient_norm(
                self.proj_content), comput_gradient_norm(self.proj_style)
            torch.nn.utils.clip_grad_norm_(self.proj_style.parameters(), self.args.max_grad_norm)
            torch.nn.utils.clip_grad_norm_(self.proj_content.parameters(), self.args.max_grad_norm)
            torch.nn.utils.clip_grad_norm_(self.encoder.parameters(), self.args.max_grad_norm)
            torch.nn.utils.clip_grad_norm_(self.decoder.parameters(), self.args.max_grad_norm)
            self.args.optimizer.step()
            self.args.scheduler.step()
            self.proj_style.zero_grad()
            self.proj_content.zero_grad()
            self.encoder.zero_grad()

        decoded_words = torch.cat(decoded_words, dim=-1)  # memory ineficient but nicer
        losses_dic = {'loss_gen': loss_gen, 'loss_mi': loss_mi, 'reny': reny, 'loss_gamma': loss_gamma,
                      'loss_h_sz': loss_h_sz, 'loss_h_s': loss_h_s, 'loss_style_classifier': loss_style_classifier,
                      'gradient_encoder': gradient_encoder, 'gradient_decoder': gradient_decoder,
                      'loss_gen_reny': loss_gen_reny, 'pos_class': pos_class, 'neg_class': neg_class,
                      'one_labels': one_labels,
                      'reny_pos_class': reny_pos_class, 'reny_neg_class': reny_neg_class,
                      'reny_one_labels': reny_one_labels,
                      'gradient_content_proj': gradient_content_proj, 'gradient_style_proj': gradient_style_proj,
                      'gradient_reny': gradient_reny / self.args.reny_training}
        return (losses_dic, decoded_words, input_tensor_gen)

    def predict_latent_space(self, input_tensor):
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(input_tensor, encoder_hidden)

        # Ensure Style not in content space
        content = self.proj_content(encoder_hidden)
        return content

    def forward_transfert(self, inputs, labels_to_transfert, pos_style, neg_style):
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(inputs, encoder_hidden)
        # Style space
        style = self.proj_style(labels_to_transfert.float().unsqueeze(1))
        style = torch.cat([style.unsqueeze(0) for _ in range(4)])

        # Content space
        content = self.proj_content(encoder_hidden)

        # Concatenate style and other
        encoder_hidden = torch.cat([style, content], dim=-1)

        # Sentence Generation
        loss_gen_reny = 0
        decoder_input = torch.ones(self.args.batch_size, 1).to(
            self.args.device).long() * self.args.tokenizer.sep_token_id
        decoder_hidden = encoder_hidden

        decoded_words = [decoder_input]

        for di in range(self.args.max_length):
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden)
            topv, topi = decoder_output.topk(1)
            decoder_input = topi.squeeze().detach().unsqueeze(-1)  # detach from history as input
            decoded_words.append(topi.squeeze(-1))
        decoded_words = torch.cat(decoded_words, dim=-1)  # clean data
        return decoded_words


class ClassificationStyleEmdedding(torch.nn.Module):
    """
    AAAI 2019 model 2
    """

    def __init__(self, args, reny_dataloader):
        super(ClassificationStyleEmdedding, self).__init__()
        self.args = args
        self.reny_dataloader = reny_dataloader
        self.encoder = EncoderRNN(args, args.number_of_tokens, args.hidden_dim, args.number_of_layers)
        self.loss_classif = torch.nn.NLLLoss()

        self.proj_content = nn.Sequential(nn.Linear(args.hidden_dim, args.hidden_dim), nn.LeakyReLU(),
                                          nn.Linear(args.hidden_dim, args.content_dim))

        # D_\gamma
        self.d_gamma = ClassifierGamma(args.content_dim + args.number_of_styles,
                                       args.number_of_styles)

        # TODO : Classifier : here we assume binary for both cases
        self.advers_classifier = Classifier(args.content_dim, args.number_of_styles, args.use_complex_classifier)
        self.downstream_classifier = Classifier(args.content_dim, args.number_of_styles, args.use_complex_classifier)

        # Loss : paper multipliers
        self.mul_mi = self.args.mul_mi

    def forward(self, input_tensor, protected_labels, downstream_labels):
        loss_mi, reny, loss_gamma, loss_style_classifier, loss_h_sz, loss_h_s, loss_adv_classifier = torch.tensor(
            0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(self.args.device)
        gradient_encoder, gradient_content_proj, gradient_reny = torch.tensor(0.0).to(self.args.device), torch.tensor(
            0.0).to(self.args.device), torch.tensor(0.0).to(self.args.device)

        reny_pos_class, reny_neg_class, reny_one_labels = torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device)
        pos_class, neg_class, one_labels = torch.tensor(0.0).to(self.args.device), torch.tensor(0.0).to(
            self.args.device), torch.tensor(0.0).to(self.args.device)

        count = 0
        for step, batch in enumerate(self.reny_dataloader):
            if count == self.args.reny_training + 1:
                break
            count += 1
            # TODO : make both adv and downstream labels
            inputs_reny = batch['line'].to(self.args.device)
            labels_reny_downstream = batch['downstream_labels'].to(self.args.device)
            labels_reny_protected = batch['label'].to(self.args.device)
            free_params(self.advers_classifier)
            frozen_params(self.downstream_classifier)
            frozen_params(self.d_gamma)
            frozen_params(self.proj_content)
            frozen_params(self.encoder)

            encoder_hidden = self.encoder.initHidden()
            encoder_output, encoder_hidden = self.encoder(inputs_reny.clone(), encoder_hidden)

            # Content space
            content = self.proj_content(encoder_hidden)
            adv_pred = self.advers_classifier(content)
            loss_adv_classifier = self.loss_classif(adv_pred, labels_reny_protected)

            reny_pos_class = torch.sum(torch.exp(adv_pred)[:, 1])
            reny_neg_class = torch.sum(torch.exp(adv_pred)[:, 0])
            reny_one_labels = torch.sum(labels_reny_protected)
            if self.training:
                loss_adv_classifier.backward()
                torch.nn.utils.clip_grad_norm_(self.advers_classifier.parameters(), self.args.max_grad_norm)
                self.args.optimizer.step()
                self.advers_classifier.zero_grad()

            ################
            # Update D_gamma
            ################
            if not self.args.no_minimization_of_mi_training:
                loss_gamma = 0
                free_params(self.d_gamma)
                frozen_params(self.downstream_classifier)
                frozen_params(self.advers_classifier)
                frozen_params(self.proj_content)
                frozen_params(self.encoder)

                inputs_reny_dgamma = inputs_reny.clone()
                encoder_hidden = self.encoder.initHidden()
                encoder_output, encoder_hidden = self.encoder(inputs_reny_dgamma, encoder_hidden)

                # Content space
                content = self.proj_content(encoder_hidden)

                adv_content_pred = self.advers_classifier(content)
                label_content_pred = adv_content_pred.topk(1, dim=-1)[-1].squeeze(-1)

                label_v = torch.tensor(
                    [[1., 0.] if el == 0 else [0., 1.] for el in label_content_pred.tolist()]).to(
                    self.args.device)
                label_u = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in labels_reny_protected.tolist()]).to(
                    self.args.device)

                u = torch.cat([content, label_u.unsqueeze(0).float().repeat(4, 1, 1)],
                              dim=-1)  # 4 is for bid + 2 layers
                v = torch.cat([content, label_v.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)

                d_gamma_content_pred_u = self.d_gamma(u)  # - log
                d_gamma_content_pred_v = self.d_gamma(v)  # - log # TODO :)

                loss_gamma = - torch.mean(d_gamma_content_pred_u[:, 0]) / 2 - torch.mean(
                    d_gamma_content_pred_v[:, 1]) / 2

                if self.training:
                    loss_gamma.backward()
                    gradient_reny += comput_gradient_norm(self.d_gamma)
                    torch.nn.utils.clip_grad_norm_(self.d_gamma.parameters(), self.args.max_grad_norm)
                    self.args.optimizer.step()
                    self.d_gamma.zero_grad()

            ##########################
            # Update downstream taks #
            ##########################
            free_params(self.downstream_classifier)
            frozen_params(self.d_gamma)
            frozen_params(self.advers_classifier)
            frozen_params(self.proj_content)
            frozen_params(self.encoder)

            inputs_reny_gen = inputs_reny.clone()
            if self.args.add_noise:
                inputs_reny_gen = corrupt_input(self, inputs_reny_gen)

            encoder_hidden = self.encoder.initHidden()
            encoder_output, encoder_hidden = self.encoder(inputs_reny_gen, encoder_hidden)

            # Content space
            content = self.proj_content(encoder_hidden)
            down_pred = self.downstream_classifier(content)
            loss_down_classifier = self.loss_classif(down_pred, labels_reny_downstream)

            if self.training:
                loss_down_classifier.backward()
                torch.nn.utils.clip_grad_norm_(self.downstream_classifier.parameters(), self.args.max_grad_norm)
                self.args.optimizer.step()
                self.downstream_classifier.zero_grad()

        ###############################
        # Update Genloss + \lambda * MI
        ###############################
        input_tensor_gen = input_tensor.clone()
        if self.args.add_noise:
            input_tensor_gen = corrupt_input(self, input_tensor_gen)
        input_tensor_mi = input_tensor_gen.clone()
        if self.training:
            frozen_params(self.d_gamma)
            frozen_params(self.advers_classifier)
            frozen_params(self.downstream_classifier)
            free_params(self.proj_content)
            free_params(self.encoder)

        # H(S)
        encoder_hidden_mi = self.encoder.initHidden()
        encoder_output_mi, encoder_hidden_mi = self.encoder(input_tensor_mi, encoder_hidden_mi)

        # Content space
        content_mi = self.proj_content(encoder_hidden_mi)
        adv_content_pred_mi = self.advers_classifier(content_mi)

        if not self.args.no_minimization_of_mi_training:
            len_ = len(protected_labels.tolist())
            loss_h_s = (-1) * torch.log(torch.sum(torch.exp(adv_content_pred_mi)[:, 0]) / len_) * (
                    1 - torch.sum(protected_labels).item() / len_) - torch.log(
                torch.sum(torch.exp(adv_content_pred_mi)[:, 1]) / len_) * (torch.sum(protected_labels).item() / len_)
            pos_class = torch.sum(torch.exp(adv_content_pred_mi)[:, 1])
            neg_class = torch.sum(torch.exp(adv_content_pred_mi)[:, 0])
            one_labels = torch.sum(protected_labels)

            # Compute Reny
            label_content_pred = adv_content_pred_mi.topk(1, dim=-1)[-1].squeeze(-1)  # TODO ?
            label_v = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in label_content_pred.tolist()]).to(
                self.args.device)
            label_u = torch.tensor([[1., 0.] if el == 0 else [0., 1.] for el in protected_labels.tolist()]).to(
                self.args.device)
            u = torch.cat([content_mi, label_u.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)
            v = torch.cat([content_mi, label_v.unsqueeze(0).float().repeat(4, 1, 1)], dim=-1)
            d_gamma_content_pred_u = self.d_gamma(u)
            d_gamma_content_pred_v = self.d_gamma(v)
            R = torch.mean((torch.exp(d_gamma_content_pred_u[:, 0]) / torch.exp(d_gamma_content_pred_v[:, 1])) ** (
                    self.args.alpha - 1))

            # Remove biais from gradients
            reny = torch.abs(torch.log(R) / (self.args.alpha - 1))

            loss_h_sz = self.loss_classif(adv_content_pred_mi, protected_labels)

            # MI
            if self.args.no_reny:
                loss_mi = self.args.mul_mi * torch.abs(loss_h_s - loss_h_sz)
            else:
                loss_mi = self.args.mul_mi * torch.abs(loss_h_s - loss_h_sz + reny)
        else:
            loss_mi = - self.args.mul_mi * self.loss_classif(adv_content_pred_mi, protected_labels)

        ###################
        # Downstream loss #
        ###################
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(input_tensor_gen, encoder_hidden)

        # Content space
        content = self.proj_content(encoder_hidden)
        down_pred = self.downstream_classifier(content)
        loss_down_classifier = self.loss_classif(down_pred, downstream_labels)

        # Compute All Losses for MI
        if self.training:
            loss = torch.abs(loss_down_classifier + loss_mi)
            loss.backward()
            gradient_encoder, gradient_content_proj = comput_gradient_norm(
                self.encoder), comput_gradient_norm(self.proj_content)
            torch.nn.utils.clip_grad_norm_(self.proj_content.parameters(), self.args.max_grad_norm)
            torch.nn.utils.clip_grad_norm_(self.encoder.parameters(), self.args.max_grad_norm)
            self.args.optimizer.step()
            self.args.scheduler.step()
            self.proj_content.zero_grad()
            self.encoder.zero_grad()

        losses_dic = {'loss_down_classifier': loss_down_classifier,
                      'loss_mi': loss_mi, 'reny': reny,
                      'loss_gamma': loss_gamma,
                      'loss_h_sz': loss_h_sz,
                      'loss_h_s': loss_h_s,
                      'loss_style_classifier': loss_style_classifier,
                      'gradient_encoder': gradient_encoder,
                      'pos_class': pos_class,
                      'neg_class': neg_class,
                      'one_labels': one_labels,
                      'reny_pos_class': reny_pos_class,
                      'reny_neg_class': reny_neg_class,
                      'reny_one_labels': reny_one_labels,
                      'gradient_content_proj': gradient_content_proj,
                      'loss_adv_classifier': loss_adv_classifier,
                      'gradient_reny': gradient_reny / self.args.reny_training}
        return losses_dic,

    def predict_latent_space(self, input_tensor):
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(input_tensor, encoder_hidden)

        # Ensure Style not in content space
        content = self.proj_content(encoder_hidden)
        return content

    def predict_downstream(self,inputs):
        encoder_hidden = self.encoder.initHidden()
        encoder_output, encoder_hidden = self.encoder(inputs, encoder_hidden)

        # Content space
        content = self.proj_content(encoder_hidden)
        down_preds = self.downstream_classifier(content)

        return down_preds
