import math
import os
import random

import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from evaluator.dataset import ZEWDPCBaseDataset
from evaluator.evaluation_metrics import get_zew_dpc_metrics

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

class CosFace(nn.Module):
    '''
     Paper: CosFace: Large Margin Cosine Loss for Deep Face Recognition
     arxiv: https://arxiv.org/pdf/1801.09414.pdf
    '''
    def __init__(self, in_features, out_features, margin, scale) -> None:
        super(CosFace, self).__init__()
        self.scale = scale
        self.margin = margin
        self.weight = nn.Parameter(torch.randn((out_features, in_features)))

    def forward(self, features, label):
        cosine_similarity = F.linear(F.normalize(features), F.normalize(self.weight))  # cos(phi) = features/||features|| * weight/||weight||
        output = (label * (cosine_similarity-self.margin)) + ((1.0 - label) * cosine_similarity)
        output *= self.scale

        return output

def train(training_dataset: ZEWDPCBaseDataset, train_idx, val_idx, model, optimizer, batch_size, epochs, criterion, device):
    start_time = time.time()
    n_train_it, n_val_it = math.ceil(train_idx.shape[0] / batch_size), math.ceil(val_idx.shape[0] / batch_size)
    for epoch in range(epochs):
        model.train()
        train_predictions, val_predictions = [], []
        random_idx = np.random.permutation(train_idx.shape[0])
        train_idx = train_idx[random_idx]
        for batch in range(n_train_it):
            optimizer.zero_grad()
            data, y_true = [], []
            for i in train_idx[batch*batch_size:(batch+1)*batch_size]:
                sample = training_dataset[i]
                data.append(sample['image'])
                y_true.append(sample['label'])
            data = torch.stack(data, dim=0).to(device)
            y_true = torch.tensor(y_true).float().to(device)
            output = model(data, y_true, True)
            train_predictions.append(output.cpu().detach())
            loss = criterion(output, y_true)
            loss.backward()
            optimizer.step()
            # print(f"==> [Train] Epoch {epoch+1}/{epochs} | Batch {batch+1}/{n_train_it} | Loss {loss:.6f} | Passed time {(time.time() - start_time)/60:.2f} min.")
        
        with torch.no_grad():
            train_predictions = torch.concat(train_predictions, dim=0)
            y_true = torch.from_numpy(training_dataset._get_all_labels()[train_idx]).float()
            loss = criterion(train_predictions, y_true)
            train_predictions[train_predictions <= .5] = 0
            train_predictions[train_predictions > .5] = 1
            scores = get_zew_dpc_metrics(training_dataset._get_all_labels()[train_idx], train_predictions)
            print(f"=> [Train] Epoch {epoch+1}/{epochs} | Loss {loss:.6f} | F1 {scores['F1_score_macro']:.3f} | "
                  f"Acc {scores['accuracy_score']:.3f} | HamL {scores['hamming_loss']:.3f} | Passed time {(time.time() - start_time)/60:.2f} min.")
            model.eval()
            for batch in range(n_val_it):
                data, y_true = [], []
                for i in val_idx[batch*batch_size:(batch+1)*batch_size]:
                    sample = training_dataset[i]
                    data.append(sample['image'])
                    y_true.append(sample['label'])
                data = torch.stack(data, dim=0).to(device)
                y_true = torch.tensor(y_true).float().to(device)
                output = model(data, y_true, True)
                val_predictions.append(output.cpu())
                loss = criterion(output, y_true)
                # print(f"==> [Val] Epoch {epoch+1}/{epochs} | Batch {batch+1}/{n_train_it} | Loss {loss:.6f} | Passed time {(time.time() - start_time)/60:.2f} min.")
            
            val_predictions = torch.concat(val_predictions, dim=0)
            y_true = torch.from_numpy(training_dataset._get_all_labels()[val_idx]).float()
            loss = criterion(val_predictions, y_true)
            val_predictions[val_predictions <= .5] = 0
            val_predictions[val_predictions > .5] = 1
            scores = get_zew_dpc_metrics(training_dataset._get_all_labels()[val_idx], val_predictions)
            print(f"=> [Val] Epoch {epoch+1}/{epochs} | Loss {loss:.6f} | F1 {scores['F1_score_macro']:.3f} | "
                  f"Acc {scores['accuracy_score']:.3f} | HamL {scores['hamming_loss']:.3f} | Passed time {(time.time() - start_time)/60:.2f} min.")

device = 'cuda:0' if torch.cuda.is_available else 'cpu'
device = torch.device(device)

mean, std = torch.tensor([0.485, 0.456, 0.406]), torch.tensor([0.229, 0.224, 0.225])
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std)
])

training_dataset = ZEWDPCBaseDataset(
        images_dir="./data/public_training/images",
        labels_path="./data/public_training/labels.csv",
        shuffle_seed=seed,
        transform=transform
    )

n_samples = len(training_dataset)
random_idx = np.random.permutation(n_samples)
train_idx, val_idx = random_idx[:math.floor(n_samples*.9)], random_idx[math.floor(n_samples*.9):]

class Encoder(nn.Module):
    def __init__(self, in_features, out_features, margin, scale) -> None:
        super(Encoder, self).__init__()
        self.model = models.resnet18(True, False)
        self.model = nn.Sequential(*list(self.model.children())[:-1], nn.Flatten())
        self.cos_face = CosFace(in_features, out_features, margin, scale)
    
    def forward(self, img, label=None, training=None):
        feature_vector = self.model(img)
        if training:
            output = self.cos_face(feature_vector, label)
            return output
        
        return feature_vector

model = Encoder(512, 6, .5, 64).to(device)

optimizer, batch_size, epochs, criterion = torch.optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-4), 32, 12, nn.CrossEntropyLoss()
train(training_dataset, train_idx, val_idx, model, optimizer, batch_size, epochs, criterion, device)