import os
import sys
sys.path.append(os.path.abspath('.'))
from logging import getLogger
from time import time
import numpy as np
from sklearn.metrics import accuracy_score
import torch
from utils.utils import AverageMeter, save


logger = getLogger('adda.trainer')


def train_source_cnn(
    source_cnn, train_loader, test_loader, criterion, optimizer,
    args=None
):
    try:
        best_score = None
        for epoch_i in range(1, 1 + args.epochs):
            start_time = time()
            training = train(
                source_cnn, train_loader, criterion, optimizer, args=args)
            validation = validate(
                source_cnn, test_loader, criterion, args=args)
            log = 'Epoch {}/{} '.format(epoch_i, args.epochs)
            log += '| Train/Loss {:.3f} Acc {:.3f} '.format(
                training['loss'], training['acc'])
            log += '| Val/Loss {:.3f} Acc {:.3f} '.format(
                validation['loss'], validation['acc'])
            log += 'Time {:.2f}s'.format(time() - start_time)
            logger.info(log)

            # save
            is_best = (best_score is None or validation['acc'] > best_score)
            best_score = validation['acc'] if is_best else best_score
            state_dict = {
                'model': source_cnn.state_dict(),
                'optimizer': optimizer.state_dict(),
                'epoch': epoch_i,
                'val/acc': best_score,
            }
            save(args.logdir, state_dict, is_best)
            logger.info('Best val. acc.: {}'.format(best_score))
    except KeyboardInterrupt as ke:
        logger.info('\n============ Summary ============= \n')
        logger.info('Best val. acc.: {}'.format(best_score))

    return source_cnn


def train_target_cnn(
    source_cnn, target_cnn, discriminator,
    criterion, optimizer, d_optimizer,
    source_train_loader, target_train_loader, target_test_loader,
    args=None
):
    validation = validate(source_cnn, target_test_loader, criterion, args=args)
    log_source = 'Source/Acc {:.3f} '.format(validation['acc'])

    try:
        best_score = None
        for epoch_i in range(1, 1 + args.epochs):
            start_time = time()
            training = adversarial(
                source_cnn, target_cnn, discriminator,
                source_train_loader, target_train_loader,
                criterion, criterion,
                optimizer, d_optimizer,
                args=args
            )
            validation = validate(
                target_cnn, target_test_loader, criterion, args=args)
            validation2 = validate(
                target_cnn, target_train_loader, criterion, args=args)
            log = 'Epoch {}/{} '.format(epoch_i, args.epochs)
            log += 'D/Loss {:.3f} Target/Loss {:.3f} '.format(
                training['d/loss'], training['target/loss'])
            log += '[Val] Target/Loss {:.3f} Target/Acc {:.3f} '.format(
                validation['loss'], validation['acc'])
            log += log_source
            log += '[Train] Target/Loss {:.3f} Target/Acc {:.3f} '.format(
                validation2['loss'], validation2['acc'])
            log += 'Time {:.2f}s'.format(time() - start_time)
            logger.info(log)

            # save
            is_best = (best_score is None or validation['acc'] > best_score)
            best_score = validation['acc'] if is_best else best_score
            state_dict = {
                'model': target_cnn.state_dict(),
                'optimizer': optimizer.state_dict(),
                'epoch': epoch_i,
                'val/acc': best_score,
            }
            save(args.logdir, state_dict, is_best)
            logger.info('Best val. acc.: {}'.format(best_score))
    except KeyboardInterrupt as ke:
        logger.info('\n============ Summary ============= \n')
        logger.info('Best val. acc.: {}'.format(best_score))


def adversarial(
    source_cnn, target_cnn, discriminator,
    source_loader, target_loader,
    criterion, d_criterion,
    optimizer, d_optimizer,
    args=None
):
    source_cnn.eval()
    target_cnn.encoder.train()
    discriminator.train()

    losses, d_losses = AverageMeter(), AverageMeter()
    n_iters = min(len(source_loader), len(target_loader))
    source_iter, target_iter = iter(source_loader), iter(target_loader)
    for iter_i in range(n_iters):
        source_data, source_target = source_iter.next()
        target_data, target_target = target_iter.next()
        source_data = source_data.to(args.device)
        target_data = target_data.to(args.device)
        bs = source_data.size(0)

        D_input_source = source_cnn.encoder(source_data)
        D_input_target = target_cnn.encoder(target_data)
        D_target_source = torch.tensor(
            [0] * bs, dtype=torch.long).to(args.device)
        D_target_target = torch.tensor(
            [1] * bs, dtype=torch.long).to(args.device)

        # train Discriminator
        D_output_source = discriminator(D_input_source)
        D_output_target = discriminator(D_input_target)
        D_output = torch.cat([D_output_source, D_output_target], dim=0)
        D_target = torch.cat([D_target_source, D_target_target], dim=0)
        d_loss = criterion(D_output, D_target)
        d_optimizer.zero_grad()
        d_loss.backward()
        d_optimizer.step()
        d_losses.update(d_loss.item(), bs)

        # train Target
        D_input_target = target_cnn.encoder(target_data)
        D_output_target = discriminator(D_input_target)
        loss = criterion(D_output_target, D_target_source)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.update(loss.item(), bs)
    return {'d/loss': d_losses.avg, 'target/loss': losses.avg}


def step(model, data, target, criterion, args):
    data, target = data.to(args.device), target.to(args.device)
    output = model(data)
    loss = criterion(output, target)
    return output, loss


def train(model, dataloader, criterion, optimizer, args=None):
    model.train()
    losses = AverageMeter()
    targets, probas = [], []
    for i, (data, target) in enumerate(dataloader):
        bs = target.size(0)
        output, loss = step(model, data, target, criterion, args)
        output = torch.softmax(output, dim=1)  # NOTE
        losses.update(loss.item(), bs)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        targets.extend(target.cpu().detach().numpy().tolist())
        probas.extend(output.cpu().detach().numpy().tolist())
    probas = np.asarray(probas)
    preds = np.argmax(probas, axis=1)
    acc = accuracy_score(targets, preds)
    return {
        'loss': losses.avg, 'acc': acc,
    }


def validate(model, dataloader, criterion, args=None):
    model.eval()
    losses = AverageMeter()
    targets, probas = [], []
    with torch.no_grad():
        for iter_i, (data, target) in enumerate(dataloader):
            bs = target.size(0)
            output, loss = step(model, data, target, criterion, args)
            output = torch.softmax(output, dim=1)  # NOTE: check
            losses.update(loss.item(), bs)
            targets.extend(target.cpu().numpy().tolist())
            probas.extend(output.cpu().numpy().tolist())
    probas = np.asarray(probas)
    preds = np.argmax(probas, axis=1)
    acc = accuracy_score(targets, preds)
    return {
        'loss': losses.avg, 'acc': acc,
    }