#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr 11 09:30:20 2022

@author: amirahijazi
"""
#%% Packages 
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
import pickle
import glob
import pandas as pd 
from pathlib import Path
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 matplotlib.pyplot as plt
from Functions import *
warnings.filterwarnings("ignore", category=UserWarning)
torch.cuda.empty_cache()

plt.style.use('ggplot')

#%%%

#%% Preparing the dataset
TOKENS = {
  '<sos>': 0,
  '<eos>': 1
}
max_nb_jobs = 81
data_folder = Path(__file__).parent / "Training_Data"
print(f"Looking in: {data_folder.resolve()}")

Files = sorted(data_folder.glob("*.pkl"))
Files_training = [str(f) for f in Files]

print(f"Found files: {Files_training}")

class CGDataset(Dataset):
  def __init__(self, max_jobs: int, Files):
      self.max_jobs = max_jobs
      self.Files = Files
      objects = []
      for file in Files:
          inputFile = open(file, 'rb')
          while True:
              try:
                  objects.append(pickle.load(inputFile))
              except EOFError:
                  break
              
          inputFile.close()
      random.shuffle(objects)
      
      #objects = [obj for obj in objects if type(obj['y']) is list]
      
      self.x_instances = [instance['x'] for instance in objects]
      y_instances = [instance['y'] for instance in objects]
      self.y_instances = [l +[0] if type(l) is list else [0] for l in y_instances]  
      self.n_samples = len(self.x_instances)

  def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
    x_instance, y_instance, length = self.pad_item(self.x_instances[idx], self.y_instances[idx])
    return x_instance, y_instance, length

  def __len__(self) -> int:
    return self.n_samples

  def pad_item(
    self,
    x_instances: list,
    y_instances: list
  ) -> Tuple[torch.tensor, torch.Tensor]:
    n_tokens = len(TOKENS)

    x_padded = np.zeros((self.max_jobs + n_tokens, 4 + n_tokens),
      dtype=np.float32)
    y_padded = np.ones((self.max_jobs), dtype=np.int64) \
      * TOKENS['<eos>']
    #y_padded = np.full((self.max_jobs,), fill_value=-1, dtype=np.int64)

    x_padded[TOKENS['<sos>'], 3] = 1.0
    x_padded[TOKENS['<eos>'], 4] = 1.0
    x_padded[n_tokens:n_tokens + len(x_instances), :3] = x_instances
    x_padded[n_tokens + len(x_instances):, 5] = 1.0
    y_padded[:len(y_instances)] = np.array([t + n_tokens for t in y_instances])

    x_padded = torch.tensor(x_padded, dtype=torch.float32)
    y_padded = torch.tensor(y_padded, dtype=torch.int64)
    length = torch.tensor(len(x_instances) + 2, dtype=torch.int64)
    return x_padded, y_padded, length

dataset = CGDataset(max_nb_jobs, Files_training)

#%%
class PointerNetwork(nn.Module):
  """
  From "Pointer Networks" by Vinyals et al. (2017)

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

  Args:
    n_hidden: The number of features to expect in the inputs.
  """

  def __init__(
    self,
    n_hidden: int
  ):
    super().__init__()
    self.n_hidden = n_hidden
    self.w1 = nn.Linear(n_hidden, n_hidden, bias=False)
    self.w2 = nn.Linear(n_hidden, n_hidden, bias=False)
    self.v = nn.Linear(n_hidden, 1, bias=False)

  def forward(
    self,
    x_decoder: torch.Tensor,
    x_encoder: torch.Tensor,
    mask: torch.Tensor,
    eps: float = 1e-16
  ) -> torch.Tensor:
    """
    Args:
      x_decoder: Encoding over the output tokens.
      x_encoder: Encoding over the input tokens.
      mask: Binary mask over the softmax input.
    Shape:
      x_decoder: (B, Ne, C)
      x_encoder: (B, Nd, C)
      mask: (B, Nd, Ne)
    """

    # (B, Nd, Ne, C) <- (B, Ne, C)
    encoder_transform = self.w1(x_encoder).unsqueeze(1).expand(
      -1, x_decoder.shape[1], -1, -1)
    # (B, Nd, 1, C) <- (B, Nd, C)
    decoder_transform = self.w2(x_decoder).unsqueeze(2)
    # (B, Nd, Ne) <- (B, Nd, Ne, C), (B, Nd, 1, C)
    prod = self.v(torch.tanh(encoder_transform + decoder_transform)).squeeze(-1)
    # (B, Nd, Ne) <- (B, Nd, Ne)
    log_score = masked_log_softmax(prod, mask, dim=-1, eps=eps)
    return log_score

    
#%%
class PointerTransformer(nn.Module):
  def __init__(
    self,                   
    c_inputs,
    c_embed,
    n_heads,
    n_layers_encoder,
    n_layers_decoder,
    dropout,
    c_hidden
  ):
    super().__init__()
    self.c_hidden = c_hidden
    self.c_inputs = c_inputs
    self.c_embed = c_embed
    self.n_heads = n_heads
    self.n_layers_decoder = n_layers_decoder
    self.n_layers_encoder = n_layers_encoder
    self.dropout = dropout

    self.embedding = nn.Linear(c_inputs, c_embed, bias=False)
    self.dropoutlayer = nn.Dropout(p=dropout)
    encoder_layers = nn.TransformerEncoderLayer(c_embed, n_heads, c_hidden, dropout)
    self.encoder = nn.TransformerEncoder(encoder_layers, n_layers_encoder)
    decoder_layers = nn.TransformerDecoderLayer(c_embed, n_heads, c_hidden, dropout)
    self.decoder = nn.TransformerDecoder(decoder_layers, n_layers_decoder)
    self.pointer = PointerNetwork(n_hidden=c_embed)
    

  def forward(
    self,
    batch_data: torch.Tensor,
    batch_lengths: torch.Tensor,
    batch_labels: Optional[torch.Tensor] = None,
  ) -> Tuple[torch.Tensor, torch.Tensor,  torch.Tensor]:
    # assumes batch-first inputs
    batch_size = batch_data.shape[0]
    max_seq_len = batch_data.shape[1]
    c_embed = self.c_embed
    n_heads = self.n_heads
   
    #batch_indices = batch_indices.unsqueeze(2)
    x_embed = self.embedding(batch_data) 
    #div_term = torch.exp(torch.arange(0, c_embed, 2) * (-math.log(10000.0) / c_embed))
    #pe = torch.zeros(batch_size, max_seq_len,c_embed)
    #pe[:, :, 0::2] = torch.sin(batch_indices * div_term)
    #pe[:, :, 1::2] = torch.cos(batch_indices * div_term)
    #x_embed = x_embed + pe
    x_embed = self.dropoutlayer(x_embed)
    encoder_outputs = self.encoder(x_embed.permute(1, 0, 2)) #torch.Size([84, 16, 64])

     # make mask
    range_tensor = torch.arange(max_seq_len, device=batch_lengths.device,
      dtype=batch_lengths.dtype).expand(batch_size, max_seq_len - len(TOKENS), max_seq_len) #([16, 81, 84])
    each_len_tensor = batch_lengths.view(-1, 1, 1).expand(-1, max_seq_len - len(TOKENS), max_seq_len)  #[16, 81, 84])
    mask_tensor = (range_tensor < each_len_tensor) #([16, 81, 84])
    mask_tensor[:,:,0] = False

    if batch_labels is not None:
      _bl = torch.cat((torch.zeros_like(batch_labels[:, :1]), batch_labels[:, :-1]), dim=1).permute(1, 0).unsqueeze(-1) #(seq_len, batch_size, 1)
      _bl = _bl.expand(-1, batch_size, c_embed) #(seq_len, batch_size, c_embed)
      decoder_input = torch.gather(encoder_outputs, dim=0, index=_bl) #torch.Size([81, 16, 64])
      decoder_mask = mask_tensor.repeat((n_heads, 1, 1))
      dm = convert_binary_mask_to_infinity_mask(decoder_mask)
   
      tgt_mask = nn.Transformer.generate_square_subsequent_mask(len(decoder_input)).to(decoder_input.device)
   
      decoder_outputs = self.decoder(decoder_input, encoder_outputs,
        tgt_mask=tgt_mask, memory_mask=dm)
 
      # pass through pointer network
      log_pointer_scores = self.pointer(
        decoder_outputs.permute(1, 0, 2),
        encoder_outputs.permute(1, 0, 2),
        mask_tensor)
      _, masked_argmaxs = masked_max(log_pointer_scores, mask_tensor, dim=-1)
      return log_pointer_scores, masked_argmaxs
    else:
      # 
      log_pointer_scores = []
      masked_argmaxs = []
      decoder_input = encoder_outputs[:1]
   
      for _ in range(max_seq_len - len(TOKENS)):
        # pass through decoder network
        decoder_mask = mask_tensor[:, :len(decoder_input)].repeat((n_heads, 1, 1))
        dm = convert_binary_mask_to_infinity_mask(decoder_mask)
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(len(decoder_input)).to(dm.device)
        decoder_outputs = self.decoder(decoder_input, encoder_outputs,
          tgt_mask=tgt_mask, memory_mask=dm)
        
        # pass through pointer network
        mask_subset = mask_tensor[:, :len(decoder_outputs)]
        log_pointer_score = self.pointer(
          decoder_outputs.permute(1, 0, 2),
          encoder_outputs.permute(1, 0, 2),
          mask_subset)
        _, masked_argmax = masked_max(log_pointer_score, mask_subset, dim=-1)

        # append new predictions
        log_pointer_scores.append(log_pointer_score[:, -1, :])
        new_maxes = masked_argmax[:, -1]
        masked_argmaxs.append(new_maxes)
        
        # mask out predicted inputs
        new_max_mask = torch.zeros((mask_tensor.shape[0], mask_tensor.shape[2]),
          dtype=torch.bool, device=mask_tensor.device)
        new_max_mask = new_max_mask.scatter(1, new_maxes.unsqueeze(1), True)
        new_max_mask = new_max_mask.unsqueeze(1).expand(-1, mask_tensor.shape[1], -1)
        mask_tensor[new_max_mask] = False

        # prepare inputs for next iteration
        next_indices = torch.stack(masked_argmaxs, dim=0).unsqueeze(-1).expand(-1, batch_size, c_embed)
        decoder_input = torch.cat((encoder_outputs[:1], 
          torch.gather(encoder_outputs, dim=0, index=next_indices)), dim=0)
      log_pointer_scores = torch.stack(log_pointer_scores, dim=1)
      masked_argmaxs = torch.stack(masked_argmaxs, dim=1)
      return log_pointer_scores, masked_argmaxs
  

n_rows_train = math.floor(len(dataset)*0.8)
n_rows_val = math.floor(len(dataset)*0.2)


torch.random.manual_seed(231)
train_dataset = Subset(dataset,list(range(n_rows_train)))
val_dataset = Subset(dataset, list(range(n_rows_train, len(dataset))))


#%%
###################
# Hyper-parameters
###################
class DotDict(dict):
    def __init__(self, **kwds):
        self.update(kwds)
        self.__dict__ = self
args = DotDict()        
args.c_inputs = 6
args.c_embed = 64
args.c_hidden = 64
args.n_heads = 8
args.n_layers_encoder = 2
args.n_layers_decoder = 2
args.dropout = 0.0
args.use_cuda = True
args.n_workers = 8
args.n_epochs = 150
args.batch_size = 16
args.lr = 0.0001
args.weight_decay = 1e-06
log_interval = 500

###################
# Hardware : CPU / GPU(s)
###################

device = torch.device("cpu"); gpu_id = -1 

if torch.cuda.is_available():
    device = torch.device("cuda")
    
print(device)

train_loader = DataLoader(train_dataset, batch_size=args.batch_size,shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle = False)

try:
    del model
except:
    pass

model = PointerTransformer(c_inputs=args.c_inputs, c_embed=args.c_embed, n_heads=args.n_heads,
  n_layers_encoder=args.n_layers_encoder,n_layers_decoder=args.n_layers_decoder, dropout=args.dropout, c_hidden=args.c_hidden).to(device)

# if trained with multiple GPUs
print(torch.cuda.device_count())
if torch.cuda.device_count()>1:
    model = nn.DataParallel(model)

optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay = args.weight_decay)
save_best_model = SaveBestModel()
# Logs
os.system("mkdir logs")
time_stamp=datetime.datetime.now().strftime("%y-%m-%d--%H-%M-%S")
file_name = 'logs'+'/'+time_stamp + "-n{}".format(max_nb_jobs) + "-gpu{}".format(gpu_id) + ".txt"
file = open(file_name,"w",1) 
file.write(time_stamp+'\n\n') 
for arg in vars(args):
    file.write(arg)
    hyper_param_val="={}".format(getattr(args, arg))
    file.write(hyper_param_val)
    file.write('\n')
file.write('\n\n') 

file.write('# train '+str(len(train_dataset)))
file.write('\n')
file.write('# val '+str(len(val_dataset)))
file.write('\n\n') 
all_strings = []
epoch_ckpt = 0
tot_time_ckpt = 0

criterion = torch.nn.NLLLoss()

train_loss = AverageMeter()
train_accuracy = AverageMeter()
val_loss = AverageMeter()
val_accuracy = AverageMeter()
best_val_loss_epoch = float('inf')

# begin training
for epoch in range(args.n_epochs):
  mean_schedule_cost_val = 0
  mean_DP = 0  
  model.train()
  for bat, (batch_data, batch_labels, batch_lengths) in enumerate(train_loader):
    batch_data = batch_data.to(device)
    batch_labels = batch_labels.to(device)
    batch_lengths = batch_lengths.to(device)
    optimizer.zero_grad()
    log_pointer_scores, pointer_argmaxs = model(batch_data, batch_lengths, batch_labels=batch_labels)
    loss = criterion(log_pointer_scores.view(-1, log_pointer_scores.shape[-1]), batch_labels.reshape(-1))

    assert not np.isnan(loss.item()), 'Model diverged with loss = NaN'
    loss.backward()
    optimizer.step()

    train_loss.update(loss.item(), batch_data.size(0))
    mask = batch_labels != TOKENS['<eos>']
    acc = masked_accuracy(pointer_argmaxs, batch_labels, mask)
    train_accuracy.update(acc, mask.int().sum().item())

    if bat % log_interval == 0:
        file.write(f'Epoch {epoch}: '
               f'Train [{bat * len(batch_data):9d}/{len(train_dataset):9d} '
               f'Loss: {train_loss.avg:.6f}\tAccuracy: {train_accuracy.avg:3.4%}\n'
               f'Labels: {batch_labels}\n'
               f'pointerargmax : {pointer_argmaxs}')
        file.write('\n\n')
  model.eval()

  for bat, (batch_data, batch_labels, batch_lengths) in enumerate(val_loader):
    batch_data = batch_data.to(device)
    batch_labels = batch_labels.to(device)
    batch_lengths = batch_lengths.to(device)
    log_pointer_scores, pointer_argmaxs = model(batch_data, batch_lengths,
      batch_labels=batch_labels)
 
    loss = criterion(log_pointer_scores.view(-1, log_pointer_scores.shape[-1]), batch_labels.reshape(-1))
    
    assert not np.isnan(loss.item()), 'Model diverged with loss = NaN'

    val_loss.update(loss.item(), batch_data.size(0))
    mask = batch_labels != TOKENS['<eos>']
    acc = masked_accuracy(pointer_argmaxs, batch_labels, mask)
    val_accuracy.update(acc, mask.int().sum().item())
    
  file.write(f'Epoch {epoch}: Val\tLoss: {val_loss.avg:.6f} '
               f'\tAccuracy: {val_accuracy.avg:3.4%}\n')
  
  save_best_model(
        val_accuracy.avg, epoch, model, optimizer, criterion
    )
  train_loss.reset()
  train_accuracy.reset()
  val_loss.reset()
  val_accuracy.reset()

  
# save the trained model weights for a final time
save_model(args.n_epochs, model, optimizer, criterion)
# save the loss and aocuracy plots
save_plots(train_loss.history, val_loss.history)
file.write('\n')
file.write('TRAINING COMPLETE')
  
# display metrics and curves
idx_best_train_loss = np.argmin(train_loss.history)
best_train_loss = train_loss.history[idx_best_train_loss]
idx_best_train_accuracy = np.argmax(train_accuracy.history)
best_train_accuracy = train_accuracy.history[idx_best_train_accuracy]
idx_best_val_loss = np.argmin(val_loss.history)
best_val_loss = val_loss.history[idx_best_val_loss]
idx_best_val_accuracy = np.argmax(val_accuracy.history)
best_val_accuracy = val_accuracy.history[idx_best_val_accuracy]
file.write('Best Scores:\n')
file.write(f'train_loss: {best_train_loss:.4f} (ep: {idx_best_train_loss})\n')
file.write(f'train_accuracy {best_train_accuracy:3.2%} (ep: {idx_best_train_accuracy})\n')
file.write(f'val_loss: {best_val_loss:.4f} (ep: {idx_best_val_loss})\n')
file.write(f'val_accuracy: {best_val_accuracy:3.2%} (ep: {idx_best_val_accuracy})\n')
x_epochs = list(range(args.n_epochs))

fig, ax = plt.subplots(1, 2, figsize=(10, 8))
ax[0].plot(x_epochs, train_loss.history, 'C2')
ax[0].plot(x_epochs, val_loss.history, 'C1')

_ = ax[0].set_xlabel('Epoch', fontsize = 12)
_ = ax[0].set_ylabel('Loss', fontsize = 12)
ax[1].plot(x_epochs, train_accuracy.history, 'C2')
ax[1].plot(x_epochs, val_accuracy.history, 'C1')
#_ = ax[1].set_title('Train vs. Val Accuracy')
_ = ax[1].set_xlabel('Epoch', fontsize = 12)
_ = ax[1].set_ylabel('Accuracy', fontsize = 12)
fig.legend(ax[0], labels = ['Training', 'Validation'], fontsize = 12, loc='upper right')
fig.suptitle('Model Loss and Accuracy', fontsize = 20)
plt.savefig("modellossaccuracy.png")
torch.cuda.empty_cache()


  



       