import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader
from torch.utils.data.sampler import WeightedRandomSampler
import math

from network import Network
from input_pipeline import get_datasets
from utils import evaluate, make_weights_for_balanced_classes


"""
The purpose of this script is to train a simple
CNN on mnist and svhn using deep CORAL method.
"""


BATCH_SIZE = 128
NUM_EPOCHS = 15
EMBEDDING_DIM = 64

DELAY = 1000  # number of steps before turning on additional losses
GROWTH_STEPS = 1000  # number of steps of linear growth of additional losses
# so domain adaptation losses are in full strength after `DELAY + GROWTH_STEPS` steps

BETA = 4.0
DEVICE = torch.device('cuda:0')
SOURCE_DATA = 'svhn'  # 'svhn' or 'mnist'


def train_and_evaluate():

    svhn, mnist = get_datasets(is_training=True)
    source_dataset = svhn if SOURCE_DATA == 'svhn' else mnist
    target_dataset = mnist if SOURCE_DATA == 'svhn' else svhn

    weights = make_weights_for_balanced_classes(source_dataset, num_classes=10)
    sampler = WeightedRandomSampler(weights, len(weights))
    source_loader = DataLoader(source_dataset, BATCH_SIZE, sampler=sampler, pin_memory=True, drop_last=True)
    target_loader = DataLoader(target_dataset, BATCH_SIZE, shuffle=True, pin_memory=True, drop_last=True)

    val_svhn, val_mnist = get_datasets(is_training=False)
    val_svhn_loader = DataLoader(val_svhn, BATCH_SIZE, shuffle=False, drop_last=False)
    val_mnist_loader = DataLoader(val_mnist, BATCH_SIZE, shuffle=False, drop_last=False)
    print('\nsource dataset is', SOURCE_DATA, '\n')

    num_steps_per_epoch = math.floor(min(len(svhn), len(mnist)) / BATCH_SIZE)
    embedder = Network(image_size=(32, 32), embedding_dim=EMBEDDING_DIM).to(DEVICE)
    classifier = nn.Linear(EMBEDDING_DIM, 10).to(DEVICE)
    model = nn.Sequential(embedder, classifier)
    model.train()

    optimizer = optim.Adam(lr=1e-3, params=model.parameters(), weight_decay=1e-3)
    scheduler = CosineAnnealingLR(optimizer, T_max=num_steps_per_epoch * NUM_EPOCHS - DELAY, eta_min=1e-6)

    cross_entropy = nn.CrossEntropyLoss()
    coral = Coral()

    text = 'e:{0:2d}, i:{1:3d}, classification loss: {2:.3f}, ' +\
        'coral loss: {3:.3f}, lr: {4:.6f}'
    i = 0  # iteration

    for e in range(NUM_EPOCHS):
        model.train()
        for (x_source, y_source), (x_target, _) in zip(source_loader, target_loader):

            x_source = x_source.to(DEVICE)
            x_target = x_target.to(DEVICE)
            y_source = y_source.to(DEVICE)

            x = torch.cat([x_source, x_target], dim=0)
            embeddings = embedder(x)
            a, b = torch.split(embeddings, BATCH_SIZE, dim=0)
            logits = classifier(a)
            usual_loss = cross_entropy(logits, y_source)
            coral_loss = coral(a, b)

            if i > DELAY:
                growth = torch.clamp(torch.tensor((i - DELAY)/GROWTH_STEPS).to(DEVICE), 0.0, 1.0)
                loss = usual_loss + growth * BETA * coral_loss
            else:
                loss = usual_loss

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

            if i > DELAY:
                scheduler.step()
            lr = scheduler.get_lr()[0]

            log = (e, i, usual_loss.item(), coral_loss.item(), lr)
            print(text.format(*log))
            i += 1

        result1 = evaluate(model, cross_entropy, val_svhn_loader, DEVICE)
        result2 = evaluate(model, cross_entropy, val_mnist_loader, DEVICE)
        print('\nsvhn loss {0:.3f} and accuracy {1:.3f}'.format(*result1))
        print('mnist loss {0:.3f} and accuracy {1:.3f}\n'.format(*result2))


class Coral(nn.Module):

    def __init__(self):
        super(Coral, self).__init__()

    def forward(self, a, b):
        """
        Arguments:
            a: a float tensor with shape [n, d].
            b: a float tensor with shape [m, d].
        Returns:
            a float tensor with shape [].
        """
        d = a.size(1)
        a = a - a.mean(0)
        b = b - b.mean(0)

        cs = torch.matmul(a.t(), a)
        ct = torch.matmul(b.t(), b)
        # they have shape [d, d]

        normalizer = 4*d*d
        return ((cs - ct)**2).sum()/normalizer


train_and_evaluate()
