import sys

import os

import copy

lib_path = 'I:/code'

if not os.path.exists(lib_path):

  lib_path = '/media/6T/.tianle/.lib'

if os.path.exists(lib_path) and lib_path not in sys.path:

  sys.path.append(lib_path)

import numpy as np

import sklearn

import torch

import torch.nn as nn

from dl.utils.visualization.visualization import plot_scatter

def cosine_similarity(x, y=None, eps=1e-8):

  """Calculate cosine similarity between two matrices; 

  Args:

    x: N*p tensor

    y: M*p tensor or None; if None, set y = x

    This function do not broadcast

  Returns:

    N*M tensor

  """

  w1 = torch.norm(x, p=2, dim=1, keepdim=True)

  if y is None:

    w2 = w1.squeeze(dim=1)

    y = x

  else:

    w2 = torch.norm(y, p=2, dim=1)

  w12 = torch.mm(x, y.t())

  return w12 / (w1*w2).clamp(min=eps)

def adjust_learning_rate(optimizer, lr, epoch, reduce_every=2):

  """Reduce learning rate by 10% every reduce_every iterations

  """

  lr = lr * (0.1 ** (epoch//reduce_every))

  for param_group in optimizer.param_groups:

    param_group['lr'] = lr

def predict(model, x, batch_size=None, train=True, num_heads=1):

  """Calculate model(x)

  Args:

    batch_size: default None; predict in one batch for small model and data.

    train: default True; if False, call torch.set_grad_enabled(False) first. 

    num_heads: default 1; if num_heads > 1, then calculate multi-head output

  """

  if batch_size is None:

    batch_size = x.size(0)

  y_pred = []

  if num_heads > 1:

    y_pred = [[] for i in range(num_heads)] # store decoder output

  prev = torch.is_grad_enabled()

  if train:

    model.train()

    torch.set_grad_enabled(True)

  else:

    model.eval()

    torch.set_grad_enabled(False)

  for i in range(0, len(x), batch_size):

    y_ = model(x[i:i+batch_size])

    if num_heads > 1:

      for i in range(num_heads):

        y_pred[i].append(y_[i])

    else:

      y_pred.append(y_)

  torch.set_grad_enabled(prev)

  if num_heads > 1:

    return [torch.cat(y, 0) for y in y_pred]

  else:

    return torch.cat(y_pred, 0)

def plot_data(model, x, y, title='', num_heads=1, batch_size=None):

  """Scatter plot for input layer and output with colors corresponding to labels

  """

  if isinstance(y, torch.Tensor):

    y = y.cpu().detach().numpy()

  plot_scatter(x, labels=y, title=f'Input {title}')

  y_pred = predict(model, x, batch_size, train=False, num_heads=num_heads)

  if num_heads > 1:

    for i in range(num_heads):

      plot_scatter(y_pred[i], labels=y, title=f'Head {i}')

  else:

    plot_scatter(y_pred, labels=y, title='Output')

def plot_data_multi_splits(model, xs, ys, num_heads=1, titles=['Training', 'Validation', 'Test'], batch_size=None):

  """Call plot_data on multiple data splits, typically x_train, x_val, x_test

  Args:

    Most arguments are passed to plot_data

    xs: a list of model input

    ys: a list of target labels

    titles: a list of titles for each data split

  """

  if len(xs) != len(titles): # Make sure titles are of the same length as xs

    titles = [f'Data split {i}' for i in range(len(xs))]

  for i, (x, y) in enumerate(zip(xs, ys)):

    if len(x) > 0 and len(x)==len(y):

      plot_data(model, x, y, title=titles[i], num_heads=num_heads, batch_size=batch_size)

    else:

      print(f'x for {titles[i]} is empty or len(x) != len(y)')

def get_label_prob(labels, verbose=True):

  """Get label distribution

  """

  if isinstance(labels, torch.Tensor):

    unique_labels = torch.unique(labels).sort()[0]

    label_prob = torch.stack([labels==i for i in unique_labels], dim=0).sum(dim=1)

    label_prob = label_prob.float()/len(labels)

  else:

    labels = np.array(labels) # if labels is a list then change it to np.array

    unique_labels = sorted(np.unique(labels))

    label_prob = np.stack([labels==i for i in unique_labels], axis=0).sum(axis=1)

    label_prob = label_prob / len(labels)

  if verbose:

    msg = '\n'.join(map(lambda x: f'{x[0]}: {x[1].item():.2f}', 

                        zip(unique_labels, label_prob)))

    print(f'label distribution:\n{msg}')

  return label_prob

def eval_classification(y_true, y_pred=None, model=None, x=None, batch_size=None, multi_heads=False, 

  cls_head=0, average='weighted', predict_func=None, pred_kwargs=None, verbose=True):

  """Evaluate classification results

  Args:

    y_true: true labels; numpy array or torch.Tensor

    y_pred: if None, then y_pred = model(x)

    model: torch.nn.Module type

    x: input tensor

    batch_size: used for predict(model, x, batch_size)

    multi_heads: If true, the model output a list; Assume the classification head is the first one

    cls_head: only used when multi_heads is True; specify which head is used for classification; default 0

    average: used for sklearn.metrics to calculate precision, recall, f1, auc and ap; default: 'weighted'

    predict_func: if not None, use predict_func(model, x, **pred_kwargs) instead of predict()

    pred_kwargs: dictionary arguments for predict_func

  """

  if isinstance(y_true, torch.Tensor): 

    y_true = y_true.cpu().detach().numpy().reshape(-1)

  num_cls = len(np.unique(y_true))

  auc = -1 # dummy variable for multi-class classification

  average_precision = -1 # dummy variable for multi-class classification

  y_score = None # only used to calculate auc and average_precision for binary classification; will be set later

  if y_pred is None: # Calculate y_pred = model(x) in batches

    if predict_func is None:

      # use predict() defined in this file

      num_heads = 2 if multi_heads else 1 # num_heads >= 2 is to make predict() to process the model as multi-output

      y_ = predict(model, x, batch_size, train=False, num_heads=num_heads)

      y_pred = y_[cls_head] if multi_heads else y_

    else:

      # use customized predict_func with variable keyworded arguments

      y_pred = predict_func(model, x, **pred_kwargs)

  if isinstance(y_pred, torch.Tensor):

    # either input argument is a torch.Tensor or calculate it from model(x) in the last chunk

    y_pred = y_pred.cpu().detach().numpy()

  if isinstance(y_pred, np.ndarray) and y_pred.ndim == 2 and y_pred.shape[1] > 1:

    # y_pred is the class score matrix: n_samples * n_classes

    if y_pred.shape[1] == 2: # for binary classification

      y_score = y_pred[:,1] - y_pred[:,0] # y_score is only useful for calculating auc and average precison 

    y_pred = y_pred.argmax(axis=-1) # only consider top 1 prediction

  if num_cls==2 and y_pred.dtype == np.dtype('float'): # last chunk had not been executed

    # For binary classification, argument y_pred can be the scores for belonging to class 1.

    y_score =  y_pred # Used for calculate auc and average_precision

    y_pred = (y_score > 0).astype('int')

  acc = sklearn.metrics.accuracy_score(y_true, y_pred)

  precision = sklearn.metrics.precision_score(y_true, y_pred, average=average)

  recall = sklearn.metrics.recall_score(y_true, y_pred, average=average)

  f1_score = sklearn.metrics.f1_score(y_true=y_true, y_pred=y_pred, average=average)

  adjusted_mutual_info = sklearn.metrics.adjusted_mutual_info_score(labels_true=y_true, labels_pred=y_pred)

  confusion_mat = sklearn.metrics.confusion_matrix(y_true, y_pred)

  msg = f'acc={acc:.3f}, precision={precision:.3f}, recall={recall:.3f}, fl={f1_score:.3f}, adj_MI={adjusted_mutual_info:.3f}'

  if num_cls == 2:

    # When y_pred is given as an int np.array or tensor, model(x) is not called; 

    # set y_score = y_pred to calculate auc and average precision approximately; 

    # it may not be 100% accurate because I assign y_pred (binary labels) to y_score (which should be probabilities)

    if y_score is None: 

      y_score = y_pred

    auc = sklearn.metrics.roc_auc_score(y_true=y_true, y_score=y_score, average=average)

    average_precision = sklearn.metrics.average_precision_score(y_true=y_true, y_score=y_score, average=average)

    msg = msg + f', auc={auc:.3f}, ap={average_precision:.3f}'

  msg = msg + f', confusion_mat=\n{confusion_mat}'

  if verbose:

    print(msg)

    print('report', sklearn.metrics.classification_report(y_true=y_true, y_pred=y_pred))

  return np.array([acc, precision, recall, f1_score, adjusted_mutual_info, auc, average_precision]), confusion_mat

def eval_classification_multi_splits(model, xs, ys, batch_size=None, multi_heads=False, cls_head=0, 

  average='weighted', return_result=True, split_names=['Train', 'Validataion', 'Test'],

  predict_func=None, pred_kwargs=None, verbose=True):

  """Call eval_classification on multiple data splits, e.g., x_train, x_val, x_test with given model

  Args:

    Most arguments are passed to eval_classification

    xs: a list of model input, e.g., [x_train, x_val, x_test]

    ys: a list of targets, e.g., [y_train, y_val, y_test]

    return_results: if True return results on non-empty data splits

    split_names: for print purpose; default: ['train', 'val', 'test']

  """

  res = []

  if len(xs) != len(split_names):

    split_names = [f'Data split {i}' for i in range(len(xs))]

  for i, (x, y) in enumerate(zip(xs, ys)):

    if len(x) > 0:

      print(split_names[i])

      metric = eval_classification(y_true=y, model=model, x=x, batch_size=batch_size, 

                          multi_heads=multi_heads, cls_head=cls_head, average=average,

                          predict_func=predict_func, pred_kwargs=pred_kwargs, verbose=verbose)

      res.append(metric)

  if return_result:

    return res

def run_one_epoch_single_loss(model, x, y_true, loss_fn=nn.CrossEntropyLoss(), train=True, optimizer=None, 

  batch_size=None, return_loss=True, epoch=0, print_every=10, verbose=True):

  """Run one epoch, i.e., model(x), but split into batches

  Args:

    model: torch.nn.Module

    x: torch.Tensor

    y_true: target torch.Tensor

    loss_fn: loss function

    train: if False, call model.eval() and torch.set_grad_enabled(False) to save time

    optimizer: needed when train is True

    batch_size: if None, batch_size = len(x)

    return_loss: if True, return epoch loss

    epoch: for print 

    print_every: print epoch_loss if print_every % epoch == 0

    verbose: if True, print batch_loss

  """

  is_grad_enabled = torch.is_grad_enabled()

  if train:

    model.train()

    torch.set_grad_enabled(True)

  else:

    model.eval()

    torch.set_grad_enabled(False)

  loss_history = []

  is_classification = isinstance(y_true.cpu(), torch.LongTensor)

  if is_classification:

    acc_history = []

  if batch_size is None:

    batch_size = len(x)

  for i in range(0, len(x), batch_size):

    y_pred = model(x[i:i+batch_size])

    loss = loss_fn(y_pred, y_true[i:i+batch_size])

    loss_history.append(loss.item())

    if is_classification:

      labels_pred = y_pred.topk(1, -1)[1].squeeze() # only calculate top 1 accuracy

      acc = (labels_pred == y_true[i:i+batch_size]).float().mean().item()

      acc_history.append(acc)

    if verbose:

      msg = 'Epoch{} {}/{}: loss={:.2e}'.format(

        epoch, i//batch_size, (len(x)+batch_size-1)//batch_size, loss.item())

      if is_classification:

        msg = msg + f', acc={acc:.2f}'

      print(msg)

    if train:

      optimizer.zero_grad()

      loss.backward()

      optimizer.step()

  torch.set_grad_enabled(is_grad_enabled)

  loss_epoch = np.mean(loss_history)

  if is_classification:

    acc_epoch = np.mean(acc_history)

  if epoch % print_every == 0:  

    msg = 'Epoch{} {}: loss={:.2e}'.format(epoch, 'Train' if train else 'Test', np.mean(loss_history))

    if is_classification:

      msg = msg + f', acc={np.mean(acc_history):.2f}'

    print(msg)

  if return_loss:

    if is_classification:

      return loss_epoch, acc_epoch, loss_history, acc_history

    else:

      return loss_epoch, loss_history

def train_single_loss(model, x_train, y_train, x_val=[], y_val=[], x_test=[], y_test=[], 

    loss_fn=nn.CrossEntropyLoss(), lr=1e-2, weight_decay=1e-4, amsgrad=True, batch_size=None, num_epochs=1, 

    reduce_every=200, eval_every=1, print_every=1, verbose=False, 

    loss_train_his=[], loss_val_his=[], loss_test_his=[], 

    acc_train_his=[], acc_val_his=[], acc_test_his=[], return_best_val=True):

  """Run a number of epochs to backpropagate

  Args:

    Most arguments are passed to run_one_epoch_single_loss

    lr, weight_decay, amsgrad are passed to torch.optim.Adam

    reduce_every: call adjust_learning_rate if cur_epoch % reduce_every == 0

    eval_every: call run_one_epoch_single_loss on validation and test sets if cur_epoch % eval_every == 0

    print_every: print epoch loss if cur_epoch % print_every == 0

    verbose: if True, print batch loss

    return_best_val: if True, return the best model on validation set for classification task 

  """

  def eval_one_epoch(x, targets, loss_his, acc_his, epoch, train=False):

    """Function within function; reuse parameters within proper scope

    """

    results = run_one_epoch_single_loss(model, x, targets, loss_fn=loss_fn, train=train, optimizer=optimizer, 

      batch_size=batch_size, return_loss=True, epoch=epoch, print_every=print_every, verbose=verbose)

    if is_classification:

      loss_epoch, acc_epoch, loss_history, acc_history = results

    else:

      loss_epoch, loss_history = results

    loss_his.append(loss_epoch)

    if is_classification:

      acc_his.append(acc_epoch)

  is_classification = isinstance(y_train.cpu(), torch.LongTensor)

  best_val_acc = -1 # best_val_acc >=0 after the first epoch for classification task

  for i in range(num_epochs):   

    if i == 0:

      optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), 

        lr=lr, weight_decay=weight_decay, amsgrad=amsgrad)

    # Should I create a new torch.optim.Adam instance every time I adjust learning rate? 

    adjust_learning_rate(optimizer, lr, i, reduce_every=reduce_every)

    eval_one_epoch(x_train, y_train, loss_train_his, acc_train_his, i, train=True)

    if i % eval_every == 0:

      if len(x_val)>0 and len(y_val)>0:

        eval_one_epoch(x_val, y_val, loss_val_his, acc_val_his, i, train=False) # Set train to be False is crucial!

        if is_classification:

          if acc_val_his[-1] > best_val_acc:

            best_val_acc = acc_val_his[-1]

            best_model = copy.deepcopy(model)

            best_epoch = i

            print('epoch {}, best_val_acc={:.2f}, train_acc={:.2f}'.format(

              best_epoch, best_val_acc, acc_train_his[-1]))

      if len(x_test)>0 and len(y_test)>0:

        eval_one_epoch(x_test, y_test, loss_test_his, acc_test_his, i, train=False) # Set train to be False

  if is_classification:

    if return_best_val and len(x_val)>0 and len(y_val)>0:

      return best_model, best_val_acc, best_epoch

    else:

      return model, acc_train_his[-1], i

def run_one_epoch_multiloss(model, x, targets, heads=[0,1], loss_fns=[nn.CrossEntropyLoss(), nn.MSELoss()], 

  loss_weights=[1,0], other_loss_fns=[], other_loss_weights=[], return_loss=True, batch_size=None, 

  train=True, optimizer=None, epoch=0, print_every=10, verbose=True):

  """Calculate a multi-head model with multiple losses including losses from the outputs and targets (head losses) 

  and regularizers on model parameters (non-head losses).

  Args:

    model: A model with multihead; for example, an AutoEncoder classifier, returns classification scores 

      (or regression target) and decoder output (reconstruction of input)

    x: input

    targets: a list of targets associated with multi-head output specified by argument heads; 

      e.g., for an autoencoder with two heads, targets = [y_labels, x]

      targets are not needed to pair with all heads output one-to-one; 

      use arguments heads to specify which heads are paired with targets;

      The elements of targets can be None, too; 

      the length of targets must be compatible with that of loss_weights, loss_fns, and heads

    heads: the index for the heads paired with targets for calculating losses; 

      if None, set heads = list(range(len(targets)))

    loss_fns: a list of loss functions for the corresponding head

    loss_weights: the (non-negative) weights for the above head-losses;

      heads, loss_fns, and loss_weights are closely related to each other; need to handle it carefully

    other_loss_fns: a list of loss functions as regularizers on model parameters

    other_loss_weights: the corresponding weights for other_loss_fns

    return_loss: default True, return all losses

    batch_size: default None; split data into batches

    train: default True; if False, call model.eval() and torch.set_grad_enabled(False) to save time

    optimizer: when train is True, optimizer must be given; default None, do not use for evaluation

    epoch: for print only

    print_every: print epoch losses if epoch % print_every == 0

    verbose: if True, print losses for each batch

  """

  is_grad_enabled = torch.is_grad_enabled()

  if train:

    model.train()

    torch.set_grad_enabled(True)

  else:

    model.eval()

    torch.set_grad_enabled(False)

  if batch_size is None:

    batch_size = len(x)

  if len(targets) < len(loss_weights):

    # Some losses do not require targets (using 'implicit' targets in the objective)

    # Add None so that targets for later use

    targets = targets + [None]*(len(loss_weights) - len(targets))

  is_classification = [] # record the indices of targets that is for classification

  has_unequal_size = [] # record the indices of targets that has a different size with input

  is_none = [] # record the indices of the targets that is None

  for j, y_true in enumerate(targets):

    if y_true is not None:

      if len(y_true) == len(x):

        if isinstance(y_true.cpu(), torch.LongTensor):

          # if targets[j] is LongTensor, treat it as classification task

          is_classification.append(j)

      else:

        has_unequal_size.append(j)

    else:

      is_none.append(j)

  loss_history = []

  if len(is_classification) > 0:

    acc_history = []

  if heads is None: # If head is not given, then assume the targets is paired with model output in order

    heads = list(range(len(targets)))

  for i in range(0, len(x), batch_size):

    y_pred = model(x[i:i+batch_size])

    loss_batch = []

    for j, w in enumerate(loss_weights):

      if w>0: # only execute when w>0

        if j in is_none:

          loss_j = loss_fns[j](y_pred[heads[j]]) * w

        elif j in has_unequal_size:

          loss_j = loss_fns[j](y_pred[heads[j]], targets[j]) * w # targets[j] is the same for all batches

        else:

          loss_j = loss_fns[j](y_pred[heads[j]], targets[j][i:i+batch_size]) * w

        loss_batch.append(loss_j)

    for j, w in enumerate(other_loss_weights):

      if w>0:

        # The implicit 'target' is encoded in the loss function itself

        # todo: in addition to argument model, make loss_fns handle other 'dynamic' arguments as well

        loss_j = other_loss_fns[j](model) * w 

        loss_batch.append(loss_j)

    loss = sum(loss_batch)

    loss_batch = [v.item() for v in loss_batch]

    loss_history.append(loss_batch)

    # Calculate accuracy

    if len(is_classification) > 0:

      acc_batch = []

      for k, j in enumerate(is_classification):

        labels_pred = y_pred[heads[j]].topk(1, -1)[1].squeeze()

        acc = (labels_pred == targets[j][i:i+batch_size]).float().mean().item()

        acc_batch.append(acc)

      acc_history.append(acc_batch)

    if verbose:

      msg = 'Epoch{} {}/{}: loss:{}'.format(epoch, i//batch_size, (len(x)+batch_size-1)//batch_size, 

        ', '.join(map(lambda x: f'{x:.2e}', loss_batch)))

      if len(is_classification) > 0:

        msg = msg + ', acc={}'.format(', '.join(map(lambda x: f'{x:.2f}', acc_batch)))

      print(msg)

    if train:

      optimizer.zero_grad()

      loss.backward()

      optimizer.step()

  torch.set_grad_enabled(is_grad_enabled)

  loss_epoch = np.mean(loss_history, axis=0)

  if len(is_classification) > 0:

    acc_epoch = np.mean(acc_history, axis=0)

  if epoch % print_every == 0:

    msg = 'Epoch{} {}: loss:{}'.format(epoch, 'Train' if train else 'Test', 

      ', '.join(map(lambda x: f'{x:.2e}', loss_epoch)))

    if len(is_classification) > 0:

      msg = msg + ', acc={}'.format(', '.join(map(lambda x: f'{x:.2f}', acc_epoch)))

    print(msg)

  if return_loss:

    if len(is_classification) > 0:

      return loss_epoch, acc_epoch, loss_history, acc_history

    else:

      return loss_epoch, loss_history

def train_multiloss(model, x_train, y_train, x_val=[], y_val=[], x_test=[], y_test=[], heads=[0, 1], 

  loss_fns=[nn.CrossEntropyLoss(), nn.MSELoss()], loss_weights=[1,0], other_loss_fns=[], other_loss_weights=[], 

  lr=1e-2, weight_decay=1e-4, batch_size=None, num_epochs=1, reduce_every=100, eval_every=1, print_every=1,

  loss_train_his=[], loss_val_his=[], loss_test_his=[], acc_train_his=[], acc_val_his=[], acc_test_his=[], 

  return_best_val=True, amsgrad=True, verbose=False):

  """Train a number of epochs

  Most of the parameters are passed to run_one_epoch_multiloss

  Args:

    lr, weight_decay, amsgrad are passed to torch.optim.Adam

    reduce_every: call adjust_learning_rate if i % reduce_every == 0; i is the current epoch

    eval_every: run_one_multiloss on validation and test set if i % eval_every == 0

    return_best_val: for classification task, if validation set is available, return the best model on validation set

    print_every: print epoch losses if i % print_every == 0

    verbose: if True, print batch losses

  """

  def eval_one_epoch(x, targets, loss_his, acc_his, epoch, train=False):

    """This is a function within a function; reuse some parameters in the scope of the "outer" function

    """

    results = run_one_epoch_multiloss(model, x, targets=targets, heads=heads, loss_fns=loss_fns, 

                loss_weights=loss_weights, other_loss_fns=other_loss_fns, other_loss_weights=other_loss_weights, 

                return_loss=True, batch_size=batch_size, train=train, optimizer=optimizer, epoch=epoch, 

                print_every=print_every, verbose=verbose)

    if is_classification:

      loss_epoch, acc_epoch, loss_history, acc_history = results

    else:

      loss_epoch, loss_history = results

    # loss_train_his += loss_history

    # acc_train_his += acc_history

    loss_his.append(loss_epoch)

    if is_classification:

      acc_his.append(acc_epoch)

  cls_targets = []

  for i, y_true in enumerate(y_train):

    if isinstance(y_true.cpu(), torch.LongTensor):

      cls_targets.append(i)

  is_classification = len(cls_targets) > 0

  best_val_acc = -1 # After the first iteration, best_val_acc >= 0 

  for i in range(num_epochs):   

    if i == 0: # I did not clear the caches after adjusting the learning rate later; this works, but is it better?

      optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, 

        weight_decay=weight_decay, amsgrad=amsgrad)

    adjust_learning_rate(optimizer, lr, i, reduce_every=reduce_every)

    eval_one_epoch(x_train, y_train, loss_train_his, acc_train_his, i, train=True)

    if i % eval_every == 0:

      if len(x_val)>0 and len(y_val)>0:

        # Must set train=False, otherwise leak data

        eval_one_epoch(x_val, y_val, loss_val_his, acc_val_his, i, train=False) 

        if is_classification:

          cur_val_acc = np.mean(acc_val_his[-1])

          if cur_val_acc > best_val_acc: # Use the mean accuracy for all classification tasks (in most case just one)

            best_val_acc = cur_val_acc

            best_model = copy.deepcopy(model)

            best_epoch = i

            print('epoch {}, best_val_acc={:.2f}, train_acc={:.2f}'.format(

              best_epoch, best_val_acc, np.mean(acc_train_his[-1])))

      if len(x_test)>0 and len(y_test)>0:

        eval_one_epoch(x_test, y_test, loss_test_his, acc_test_his, i, train=False)

  if is_classification:

    if return_best_val and len(x_val)>0 and len(y_val)>0:

      return best_model, best_val_acc, best_epoch

    else:

      return model, np.mean(acc_train_his[-1]), i