"""Adversarial adaptation to train target encoder."""

import os

import torch
import torch.optim as optim
from torch import nn

from utils import make_variable
from core.test import eval

def train_tgt(src_encoder, src_classifier, tgt_encoder, critic, src_data_loader, tgt_data_loader, params):
    """Train encoder for target domain."""
    ####################
    # 1. setup network #
    ####################

    # setup criterion and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer_tgt = optim.Adam(tgt_encoder.parameters(),
                               lr=params.tgt_learning_rate,
                               betas=(params.beta1, params.beta2),
                               weight_decay=2.5e-5)
    optimizer_critic = optim.Adam(critic.parameters(),
                                  lr=params.critic_learning_rate,
                                  betas=(params.beta1, params.beta2),
                                  weight_decay=2.5e-5)
    len_data_loader = min(len(src_data_loader), len(tgt_data_loader))

    ####################
    # 2. train network #
    ####################

    for epoch in range(params.num_epochs):

        # set train state for Dropout and BN layers
        tgt_encoder.train()
        critic.train()

        # zip source and target data pair
        data_zip = enumerate(zip(src_data_loader, tgt_data_loader))
        for step, ((images_src, _), (images_tgt, _)) in data_zip:
            ###########################
            # 2.1 train discriminator #
            ###########################

            # make images variable
            images_src = make_variable(images_src)
            images_tgt = make_variable(images_tgt)

            # zero gradients for optimizer
            optimizer_critic.zero_grad()

            # extract and concat features
            feat_src = src_encoder(images_src)
            feat_tgt = tgt_encoder(images_tgt)
            feat_concat = torch.cat((feat_src, feat_tgt), 0)

            # predict on discriminator
            pred_src = critic(feat_src.detach())
            pred_tgt = critic(feat_tgt.detach())

            # prepare real and fake label
            label_src = make_variable(torch.ones(feat_src.size(0)).long())
            label_tgt = make_variable(torch.zeros(feat_tgt.size(0)).long())
            print(label_src.shape, label_src.max(), label_src.min())
            print(label_tgt.shape, label_tgt.max(), label_tgt.min())
            
            # compute loss for critic
            loss_critic = criterion(pred_src, label_src) + criterion(pred_tgt, label_tgt)
            loss_critic.backward()

            # optimize critic
            optimizer_critic.step()
           
            print(pred_src.shape, pred_src.max(), pred_src.min())
            print(pred_tgt.shape, pred_tgt.max(), pred_tgt.min())
            pred_cls_src = torch.squeeze(pred_src.max(1)[1])
            pred_cls_tgt = torch.squeeze(pred_tgt.max(1)[1])
            acc = ((pred_cls_src == label_src).float().mean() + (pred_cls_tgt == label_tgt).float().mean()) / 2

            ############################
            # 2.2 train target encoder #
            ############################

            # zero gradients for optimizer
            optimizer_critic.zero_grad()
            optimizer_tgt.zero_grad()

            # extract and target features
            feat_tgt = tgt_encoder(images_tgt)

            # predict on discriminator
            pred_tgt = critic(feat_tgt)

            # prepare fake labels
            label_tgt = make_variable(torch.ones(feat_tgt.size(0)).long())

            # compute loss for target encoder
            loss_tgt = criterion(pred_tgt, label_tgt)
            loss_tgt.backward()

            # optimize target encoder
            optimizer_tgt.step()

            #######################
            # 2.3 print step info #
            #######################
            if ((step + 1) % params.log_step == 0):
                print("Epoch [{:3d}/{:3d}] Step [{:3d}/{:3d}]:"
                      "critic_loss={:.5f} tgt_loss={:.5f} acc={:.5f}"
                      .format(epoch + 1,
                              params.num_epochs,
                              step + 1,
                              len_data_loader,
                              loss_critic.item(),
                              loss_tgt.item(),
                              acc.item()))

        #############################
        # 2.4 eval training model #
        #############################
        if ((epoch + 1) % params.eval_step == 0):
            print ("eval model on source data")
            eval(tgt_encoder, src_classifier, src_data_loader)
            print ("eval model on target data")
            eval(tgt_encoder, src_classifier, tgt_data_loader)

        #############################
        # 2.5 save model parameters #
        #############################
        if ((epoch + 1) % params.save_step == 0):
            torch.save(critic.state_dict(), os.path.join(
                params.model_root,
                "{}-{}-critic-{}.pt".format(params.src_dataset, params.tgt_dataset, epoch + 1)))
            torch.save(tgt_encoder.state_dict(), os.path.join(
                params.model_root,
                "{}-{}-target-encoder-{}.pt".format(params.src_dataset, params.tgt_dataset, epoch + 1)))

    torch.save(critic.state_dict(), os.path.join(
        params.model_root,
        "{}-{}-critic-final.pt".format(params.src_dataset, params.tgt_dataset)))
    torch.save(tgt_encoder.state_dict(), os.path.join(
        params.model_root,
        "{}-{}-target-encoder-final.pt".format(params.src_dataset, params.tgt_dataset)))
    return tgt_encoder