# -*- coding: utf-8 -*-
"""
Created on Sat Oct  8 18:36:26 2022

@author: amhijazi
"""

#%% 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 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 matplotlib.pyplot as plt
from Functions import *
warnings.filterwarnings("ignore", category=UserWarning)
torch.random.manual_seed(231)

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

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

gpu_id = '0' # select a single GPU  
#gpu_id = '2,3' # select multiple GPUs  
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)  
if torch.cuda.is_available():
    device = torch.device("cuda")
    print('GPU name: {:s}, gpu_id: {:s}'.format(torch.cuda.get_device_name(0),gpu_id))   
    
print(device)

###################
# Network definition
# Notation : 
#            bsz : batch size
#            nb_nodes : number of nodes/cities
#            dim_emb : embedding/hidden dimension
#            nb_heads : nb of attention heads
#            dim_ff : feed-forward dimension
#            nb_layers : number of encoder/decoder layers
###################

#%% calculating schedule cost  
def compute_schedule(x, schedule):

    bsz, num_jobs = schedule.shape
    
    # Mask invalid jobs (<= 2)
    mask = schedule > 2
    filtered = schedule.clone()
    filtered[~mask] = -1

    # Gather proc and weights
    proc = torch.gather(x[:, :, 0], 1, filtered.clamp(min=0))
    weights = torch.gather(x[:, :, 1], 1, filtered.clamp(min=0))
    # Compute ratio and sort
    ratio = proc / weights
    ratio[filtered == -1] = float('inf')
    sorted_idx = torch.argsort(ratio, dim=1)
    sorted_jobs = torch.gather(filtered, 1, sorted_idx)
    valid_mask = sorted_jobs >= 0
    proc_sorted = torch.gather(proc, 1, sorted_idx)
    weights_sorted = torch.gather(weights, 1, sorted_idx)
    # Compute cumulative sums and contributions
    cumsum_proc = torch.cumsum(proc_sorted * valid_mask, dim=1)
    contrib = cumsum_proc * weights_sorted * valid_mask

    # Schedule cost
    C = contrib.sum(dim=1)

    # Reduced cost
    dual_values = torch.gather(x[:, :, 2], 1, sorted_jobs.clamp(min=0))
    dual_sum = (dual_values * valid_mask).sum(dim=1)
    L = C - dual_sum - x[:, 2, 2]

    # Convert schedules to lists
    Schedules = [(sorted_jobs[b][valid_mask[b]] - 2).tolist() for b in range(bsz)]

    return L, C, Schedules



#%%
TOKENS = {
  '<sos>': 0,
  '<eos>': 1
}
#max_nb_jobs = args.max_nb_jobs
class CGDataset(Dataset):
  def __init__(self, max_jobs: int, test_dataset):
      self.max_jobs = max_jobs
     
      self.x_instances = test_dataset
      self.n_samples = self.x_instances.shape[0]
  

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

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

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

    x_padded = np.zeros((self.max_jobs + n_tokens, 4 + n_tokens),
      dtype=np.float32)
    
    
    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
    
    x_padded = torch.tensor(x_padded, dtype=torch.float32)
    length = torch.tensor(len(x_instances) + 2, dtype=torch.int64)
    return x_padded, length
#%%
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,
    max_len = 1000
  ):
    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)
   
    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,
    beam_width: int = 1
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
    # 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) 
    
    encoder_outputs = self.encoder(x_embed.permute(1, 0, 2))

    # 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)
    each_len_tensor = batch_lengths.view(-1, 1, 1).expand(-1, max_seq_len - len(TOKENS), max_seq_len)
    mask_tensor = (range_tensor < each_len_tensor)
 
   
    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]
          if torch.all((new_maxes == TOKENS['<eos>'])):
             break
          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[:, 1] = False
          new_max_mask[:, 0] = 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)
    masked_argmaxs = torch.stack(masked_argmaxs, dim=1)
    
    return masked_argmaxs

 
torch.random.manual_seed(231)

def test(model, batch_data, batch_lengths, M):
    model.eval()
    pointer_argmaxs = model(batch_data.to(device), batch_lengths.to(device), None, 1)
    ReducedCost, Cost, schedule = compute_schedule(batch_data, pointer_argmaxs)
    return ReducedCost, Cost, schedule

