import torch
import torch.nn as nn
import time
import argparse
from typing import Tuple, Union, Optional
from torch.utils.data import Dataset, DataLoader, Subset
import random
import os
import datetime
from torch.distributions.gumbel import Gumbel
import pickle
import pandas as pd 
from torch.distributions.categorical import Categorical
import numpy as np
from torch.utils import data
from functools import lru_cache
import warnings
import math
import torch.nn.functional as F

import matplotlib.pyplot as plt
warnings.filterwarnings("ignore", category=UserWarning)
torch.cuda.empty_cache()

class SaveBestModel:
    """
    Class to save the best model while training. If the current epoch's 
    validation loss is less than the previous least less, then save the
    model state.
    """
    def __init__(
        self, best_valid_loss=0.0
    ):
        self.best_valid_loss = best_valid_loss
        
    def __call__(
        self, current_valid_loss, 
        epoch, model, optimizer, criterion
    ):
        if current_valid_loss < self.best_valid_loss:
            self.best_valid_loss = current_valid_loss
            print(f"\nBest validation loss: {self.best_valid_loss}")
            print(f"\nSaving best model for epoch: {epoch+1}\n")
            # Saving checkpoint
            checkpoint_dir = os.path.join("checkpoint")
            if not os.path.exists(checkpoint_dir):
                  os.makedirs(checkpoint_dir)
            torch.save({
                  'epoch': epoch,
                  'loss': criterion,
                  'model_train': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  }, '{}.pkl'.format(checkpoint_dir+ "best_model"))

            
def save_model(epochs, model, optimizer, criterion):
        """
        Function to save the trained model to disk.
        """
        print(f"Saving final model...")
        checkpoint_dir = os.path.join("checkpoint")
        if not os.path.exists(checkpoint_dir):
                  os.makedirs(checkpoint_dir)
        torch.save({
                  'epoch': epochs,
                  'loss': criterion,
                  'model_train': model.state_dict(),
                  'optimizer': optimizer.state_dict(),
                  }, '{}.pkl'.format(checkpoint_dir+ "final_model"))

        
def save_plots(train_loss, valid_loss):
        # loss plots
        plt.figure(figsize=(10, 7))
        plt.plot(
            train_loss, color='orange', linestyle='-', 
            label='train loss'
        )
        plt.plot(
            valid_loss, color='red', linestyle='-', 
            label='validataion loss'
        )
        plt.xlabel('Epochs')
        plt.ylabel('Loss')
        plt.legend()
        plt.savefig('loss.png')   

def info_value_of_dtype(dtype: torch.dtype):
  """
  Returns the `finfo` or `iinfo` object of a given PyTorch data type. Does not allow torch.bool.

  Adapted from allennlp by allenai:
    https://github.com/allenai/allennlp/blob/master/allennlp/nn/util.py
  """
  
  if dtype == torch.bool:
    raise TypeError("Does not support torch.bool")
  elif dtype.is_floating_point:
    return torch.finfo(dtype)
  else:
    return torch.iinfo(dtype)


def min_value_of_dtype(dtype: torch.dtype):
  """
  Returns the minimum value of a given PyTorch data type. Does not allow torch.bool.

  Adapted from allennlp by allenai:
    https://github.com/allenai/allennlp/blob/master/allennlp/nn/util.py
  """

  return info_value_of_dtype(dtype).min


def masked_log_softmax(
  x: torch.Tensor,
  mask: torch.Tensor,
  dim: int = -1,
  eps: float = 1e-45
) -> torch.Tensor:
  """
  Apply softmax to x with masking.

  Adapted from allennlp by allenai:
    https://github.com/allenai/allennlp/blob/master/allennlp/nn/util.py

  Args:
    x - Tensor of arbitrary shape to apply softmax over.
    mask - Binary mask of same shape as x where "False" indicates elements
      to disregard from operation.
    dim - Dimension over which to apply operation.
    eps - Stability constant for log operation. Added to mask to avoid NaN
      values in log.
  Outputs:
    Tensor with same dimensions as x.
  """

  x = x + (mask.float() + eps).log()
  return torch.nn.functional.log_softmax(x, dim=dim)
 
def masked_max(
  x: torch.Tensor,
	mask: torch.Tensor,
	dim: int,
	keepdim: bool = False
) -> Tuple[torch.Tensor, torch.Tensor]:
  """
  Apply max to x with masking.

  Adapted from allennlp by allenai:
    https://github.com/allenai/allennlp/blob/master/allennlp/nn/util.py

  Args:
    x - Tensor of arbitrary shape to apply max over.
    mask - Binary mask of same shape as x where "False" indicates elements
      to disregard from operation.
    dim - Dimension over which to apply operation.
    keepdim - If True, keeps dimension dim after operation.
  Outputs:
    A ``torch.Tensor`` of including the maximum values.
  """

  x_replaced = x.masked_fill(~mask, min_value_of_dtype(x.dtype))
  max_value, max_index = x_replaced.max(dim=dim, keepdim=keepdim)
  return max_value, max_index

def convert_binary_mask_to_infinity_mask(mask: torch.Tensor) -> torch.Tensor:
  """
  Convert the 0 and 1 elements in a binary mask to -inf and 0 for the
    transformer.

  Args:
    mask: Binary mask tensor.
  Outputs:
    Infinity mask tensor with same size as mask.
  """

  return mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))


class AverageMeter(object):
  """
  Computes and stores the average and current value

  Adapted from pointer-networks-pytorch by ast0414:
    https://github.com/ast0414/pointer-networks-pytorch
  """

  def __init__(self):
    self.history = []
    self.reset(record=False)

  def reset(
    self,
    record: bool = True
  ):
    if record:
      self.history.append(self.avg)
    self.val = 0
    self.avg = 0
    self.sum = 0
    self.count = 0

  def update(
    self,
    val: Union[float, int],
    n: int = 1
  ):
    self.val = val
    self.sum += val * n
    self.count += n
    self.avg = self.sum / self.count

def masked_accuracy(
  output: torch.Tensor,
  target: torch.Tensor,
  mask: torch.Tensor
) -> float:
  """
  Compute accuracy of softmax output with mask applied over values.

  Adapted from pointer-networks-pytorch by ast0414:
    https://github.com/ast0414/pointer-networks-pytorch
  """

  with torch.no_grad():
        correct = (output == target) & mask
        accuracy = correct.sum().float() / mask.sum().float()
        return accuracy.item()

