# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in copliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

import sys
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.tri import triplot
from matplotlib.pyplot import savefig
import datetime
import math

from keras.utils import np_utils

#from sklearn.metrics import classification_report 

from utils.profiler import profile
from utils.list_utils import filter_list_by_index
from utils import my_numpy_utils as my_np

#
#  MimoSequence: a class to manage multiple-input multiple-output (MIMO) data
#
#  attributes:
#    Arbitrary-length data:
#    input: list of arbitrary length input sequences, sequence_i begin a (seq_length_i, n_input_dim) npArray 
#    target: list of of arbitrary length target sequences, sequence_i begin a (seq_length_i, n_output_dim) npArray
#    output: list of of arbitrary length output sequences, sequence_i begin a (seq_length_i, n_output_dim) npArray
#
#    Fixed-length data:
#    block_input: numpy.array of input sequences (dimension n_sequences x max_length x n_input_dim)
#    block_target: numpy.array of target sequences (dimension n_sequences x max_length x n_output_dim)
#    block_output: numpy.array of output sequences (dimension n_sequences x max_length x n_output_dim)
#    
#    for fixed-length data, sequences of length < max_length are right padded with pad values: 
#      masked_input_val = 0
#      padded_target_val = [0 0 ... 0 1]
#    These samples become effectively an additional class "No signal" or "PAD"
#    
#    Sequences longer than max_length can be split into multiple sequences by split_long(...)
#    before calling generate_block_sequences(...).  This augments n_sequences.
#
class MimoSequence:
  def __init__(self, tag, n_input_dim=1, 
                          n_output_dim = 1, 
                          max_length=None, 
                          labels=None, 
                          add_pad_label=False,                          
                          out_dec_factor = 1):
  
    self.tag = tag

    self.n_input_dim = n_input_dim
    self.n_output_dim = n_output_dim
    self.max_length = max_length
    self.labels = labels
    if add_pad_label:
      self.add_pad_label()
    
    self.out_dec_factor = out_dec_factor
    
    self.input = []
    self.target = []
    self.output = []
    
    self.block_input = []
    self.block_target = []
    self.block_output = []
        
    self.sequence_tags = []
    self.sequence_lengths = []

    #self.n_sequences = 0
    
    self.sample_weights = []
    self.padded_target_val = []
    self.masked_input_val = 0
        
    self.n_labels = len(self.labels)
      
  def clone_structure(self, tag):
    new_mimo = MimoSequence(tag, n_input_dim=self.n_input_dim, 
                             n_output_dim=self.n_output_dim, 
                             max_length=self.max_length, labels=self.labels, 
                             out_dec_factor=self.out_dec_factor)
    return new_mimo
  
  def add_pad_label(self):
    self.labels = ['PAD'] + self.labels
    self.n_labels = len(self.labels)
    
  def append_input(self, input, tag):
    self.input.append(input)
    self.sequence_lengths.append(len(input))
    self.sequence_tags.append(tag)
  
  # expects (self.n_labels-1) 0-based targets and generates 
  # self.n_labels 0-based targets with 0 reserved for 'PAD'
  def append_target(self, target, pad_to_length=False):
    cur_sequence = len(self.target)
    if pad_to_length:
      target = np.tile(target, 
                (self.sequence_lengths[cur_sequence]//self.out_dec_factor,));
    #self.target.append(target)
    self.target.append(target+1)
        
  def n_sequences(self):
#     assert(len(self.input) == len(self.target) and
#           len(self.input) == len(self.sequence_lengths))
    assert(len(self.input) == len(self.sequence_lengths))
    if len(self.target) > 0:
      assert(len(self.input) == len(self.target))    
    if len(self.output) > 0:
      assert(len(self.output) == len(self.target))
    return len(self.input)
  
  def my_display(self):
    print('mimo_sequence_old attributes')
    print('Tag: ' + self.tag)
    print('Labels:')
    print(self.labels);
    print('Number of sequences: ' + str(self.n_sequences()))
    print('Sequence tags')
    for x in self.sequence_tags:
      print(x)
    print('Sequence lengths')
    for x in self.sequence_lengths:
      print(x)
    print('Input (1-100)')
    print(self.input[1:100])  
    print('Target (1-100)')
    print(self.target[1:100])
    
  #  from concatenated input and target with variable sequence lengths,
  #  in a (sum_sequence_lengths, n_input_dim) numpy.array
  #  extract the individual sequences into a list of  
  #  (sequence_length, n_input_dim)  numpy.array 
  def extract_sequences(self, concat_input, concat_target, sequence_lengths):    

      self.sequence_lengths = sequence_lengths
      n_sequences = len(self.sequence_lengths)
            
      #self.n_features = self.input.shape[1]
      if concat_input.ndim == 1:
        self.n_input_dim = 1
      else:
        self.n_input_dim = concat_input.shape[1]
        
      if concat_target.ndim == 1:
        self.n_output_dim = 1 
      else:
        self.n_output_dim = concat_target.shape[1]
        
      #self.n_sequences = len(self.sequence_lengths)
      if self.max_length is None:
        self.max_length = max(self.sequence_lengths)
      
      #self.n_labels = len(self.labels)
      
      #self.input = np.zeros((self.n_sequences, self.max_length, self.n_features))
      #self.target = np.tile(self.padded_output_val, (self.n_sequences, self.max_length, 1));
             
      offset = 0
      for i in range(n_sequences):
          thisLength = self.sequence_lengths[i]
          truncLength = thisLength
          if truncLength > self.max_length:
            truncLength = self.max_length;
          #self.input[i, 0:truncLength, :] = concat_input[offset:offset+truncLength, :]
          self.input.append(concat_input[offset:offset+truncLength, :])
          #self.target.append(np_utils.to_categorical(concat_target[offset:offset+truncLength,]+1, self.n_labels+1)) 
          self.target.append(concat_target[offset:offset+truncLength,]) 
          offset += thisLength

  #
  #   Ensure that all sequences longer than base_length*1.5 are split into
  #   shorter sequence segments. Initial splits are of length base_length*1.5 
  #   and the final segments will be of minimal length base_length*0.5.   
  def split_long(self, base_length):
    if base_length == sys.maxsize:
      max_length = max(self.sequence_lengths)
      print('split_long: no splitting done max_length: ' + str(max_length))
      return self

    base_length_150 = math.floor(base_length*1.5)
    base_length_200 = math.floor(base_length*2.0)
    max_length = base_length_150 
    print('split_long: max_length: ' + str(max_length))
    
    split_mimo = self.clone_structure(self.tag + '_split_' + str(base_length_150))
    split_mimo.max_length = base_length_150
    
    n_sequences = self.n_sequences()
    print('split_long: n_sequences before: ' + str(self.n_sequences()))
        
    offset = 0
    j_sequences = 0
    for i in range(n_sequences):
      cur_length = self.sequence_lengths[i]
      n_segments = 0
      #current_input = np.asarray(self.input[i])
      current_input = self.input[i]
      current_target = self.target[i]
      while cur_length > 0:
        if cur_length >= base_length_200:
          cur_segment_length = base_length_150
        elif cur_length > base_length_150:
          cur_segment_length = cur_length//2;
        else: 
          cur_segment_length = cur_length;
          
        split_mimo.input.append(current_input[0:cur_segment_length])
        # consider only target length 1 for now
        #split_mimo.target.append(self.target[i])
        split_mimo.target.append(current_target[0:cur_segment_length//self.out_dec_factor])
        
        # and disregard output
        #split_mimo.output.append(self.output[i])

        #don't change id if no split
        #split_mimo.sequence_tags.append( + '_' + str(n_segments))
        curr_tag = self.sequence_tags[i]
        if n_segments > 0:
            curr_tag = curr_tag  + '_' + str(n_segments)
        split_mimo.sequence_tags.append(curr_tag)

        split_mimo.sequence_lengths.append(cur_segment_length)
        n_segments += 1 
        j_sequences += 1 
        current_input = np.delete(current_input, np.s_[0:cur_segment_length:1], axis=0) 
        current_target = np.delete(current_target, np.s_[0:cur_segment_length//self.out_dec_factor:1], axis=0) 
        cur_length = cur_length - cur_segment_length
#         print(self.sequence_tags[i] + '_' + str(n_segments) + ' input length: ' + str(cur_segment_length)
#               + ' target length: ' + str(cur_segment_length//self.out_dec_factor))
        print(curr_tag + ' input length: ' + str(cur_segment_length)
              + ' target length: ' + str(cur_segment_length//self.out_dec_factor))

    print('split_long: n_sequences after: ' + str(self.n_sequences()))
    return split_mimo
  
  
  
#   from the arbitrary sequence lengths, extract individual sequences into a 
#   nFeatures x nTimesteps x nSequences npArray
#   It is assumed that self.target for n classes ranges from 0.. n-1 
#   But output class 0 is reserved for padded of variable length sequences
#   So, the padding value on the input and output is 0
#   The output one-hot class for padded values is 1000 (e.g., for a three-class problem)
#   and the first class of the three classes is encoded as 0100 
  def generate_block_sequences(self):
      if len(self.target)==0:
        self.generate_block_input()
        return
            
      if self.max_length is None:
        self.max_length = max(self.sequence_lengths)
  
      #self.n_labels = len(self.labels)
      n_sequences = self.n_sequences()
      
      self.block_input = np.zeros((n_sequences, self.max_length, self.n_input_dim))
      #self.padded_target_val = np.zeros(self.n_labels+1)
      self.padded_target_val = np.zeros(self.n_labels)
      self.padded_target_val[0] = 1;
      self.block_target = np.tile(self.padded_target_val, (n_sequences, self.max_length//self.out_dec_factor, 1));
      
      for i in range(n_sequences):
          thisLength = self.sequence_lengths[i]
          self.block_input[i, 0:thisLength, :] = self.input[i]
          self.block_target[i, 0:(thisLength//self.out_dec_factor), :] = np_utils.to_categorical(self.target[i], self.n_labels) 
      self.calc_sample_weights()  
                 
#   from the arbitrary sequence lengths, extract individual sequences into a 
#   nFeatures x nTimesteps x nSequences npArray
  def generate_block_input(self):
      if self.max_length is None:
        self.max_length = max(self.sequence_lengths)
  
      self.n_labels = len(self.labels)
      n_sequences = self.n_sequences()
      
      self.block_input = np.zeros((n_sequences, self.max_length, self.n_input_dim))
      
      for i in range(n_sequences):
          thisLength = self.sequence_lengths[i]
          self.block_input[i, 0:thisLength, :] = self.input[i]
                 
  def filter_by_index(self, filterRange, keepTag, complTag=None):  
    keep = self.clone_structure(keepTag)
    compl = self.clone_structure(complTag)
    
    keep.input, compl.input = filter_list_by_index(self.input, filterRange)
    keep.target, compl.target = filter_list_by_index(self.target, filterRange)
    keep.output, compl.output = filter_list_by_index(self.output, filterRange)
    keep.sequence_tags, compl.sequence_tags = filter_list_by_index(self.sequence_tags, filterRange)
    keep.sequence_lengths, compl.sequence_lengths = filter_list_by_index(self.sequence_lengths, filterRange)   
    
    return keep, compl
    
         
  
  #@profile
  def read(self, inputFile):  
    pass
  
  #
  #  weight valid samples according to target value:
  #   if target == padded_target_val, weight = 0
  #   otherwise weight = 1
  #
  def calc_sample_weights(self):
    self.sample_weights = 1.0 * (1 - np.all((self.block_target-self.padded_target_val)==0, axis=2))  
    
  
  def calculate_pad_size(self, n_sequences, base):
    n_to_pad = n_sequences % base
    if n_to_pad > 0:
      #n_to_pad =  max(n_sequences, base) - n_to_pad
      n_to_pad =  base - n_to_pad
    return n_to_pad 

  
  def pad_number_of_sequences(self, n_sequences):
    n_to_pad = self.calculate_pad_size(self.n_sequences(), n_sequences)
    if n_to_pad == 0:
      print('pad_number_of_sequences: no padding required')
      return

    # pad input and target
    pad_input = [np.zeros((self.n_input_dim))] * n_to_pad
    pad_target = [np.zeros((self.n_output_dim))] * n_to_pad
    self.input = self.input + pad_input
    self.target = self.target + pad_target
    self.sequence_lengths = self.sequence_lengths + [0] * n_to_pad
    self.sequence_tags = self.sequence_tags + ['pad'] * n_to_pad
    
    # pad block_input and block_target
    pad_input = np.zeros((n_to_pad, self.max_length, self.n_input_dim))
    #pad_target = np.tile(self.padded_target_val, (n_to_pad, self.max_length, 1));
    pad_target = np.tile(self.padded_target_val, (n_to_pad, self.max_length//self.out_dec_factor, 1));
    self.block_input = np.concatenate((self.block_input, pad_input), axis=0)
    self.block_target = np.concatenate((self.block_target, pad_target), axis=0)
    self.calc_sample_weights()
    
  # right pad output by n_samples and then
  # truncate left by n_samples 
  def shift_target_left(self, n_samples):
    #self.plot_output('OriginalOutput.pdf')
    pad_output = np.tile(self.padded_target_val, (self.target.shape[0], n_samples, 1));
    self.target = np.concatenate((self.target, pad_output), axis=1)
    #np.delete(self.target, list(range(n_samples)), axis=1) 
    self.target = np.delete(self.target, np.s_[0:n_samples:1], axis=1) 

    zero_input = np.zeros((self.n_sequences, n_samples, self.n_input_dim))
    self.input[:, (self.max_length-n_samples):self.max_length, :] = zero_input
    
    print(self.target)
    #self.plot_output('ShiftedOutput.pdf')
    
  def shift_target_right(self, n_samples):
    #self.plot_output('OriginalOutput.pdf')
    self.target = np.delete(self.target, np.s_[(self.max_length-n_samples):self.max_length:1], axis=1) 
    pad_output = np.tile(self.padded_target_val, (self.target.shape[0], n_samples, 1));
    self.target = np.concatenate((pad_output, self.target), axis=1)

    #zero_input = np.zeros((self.n_sequences, n_samples, self.n_input_dim))
    #self.input[:, (self.max_length-n_samples):self.max_length, :] = zero_input
    print(self.target)
    self.plot_output('ShiftedOutput')
    
  def pad_last_output(self):
    self.plot_output('OriginalOutput.pdf')
    last_sample = self.block_output[:, -1, :]
    new_output = np.tile(last_sample, (1, self.block_output.shape[1], 1))
    
    self.output = new_output
    self.plot_output('ShiftedOutput.pdf')
    gothere = True
    
  def add_input_gaussian_noise(self, sigma):
    if sigma <= 0:
        return    
    maskedInput = 1.0 * (1 - np.all((self.block_input-self.masked_input_val)==0, axis=2))  
    maskedInputTiled = np.tile(maskedInput[..., None], (1, 1, self.n_input_dim));
    
    randomInputNoise = np.random.normal(loc=0.0, scale=sigma, size=self.block_input.shape) * maskedInputTiled
    self.block_input += randomInputNoise
        
  def plot_target(self, tag=''):
    #nPlots = trainMimo.n_sequences()
    nSeq = min(self.n_sequences(), 6)
    print('Plotting Results')
    for iSeq in range(nSeq):
      plt.subplot(nSeq, 1, iSeq+1)
      currExOutput = my_np.from_categorical(self.block_target[iSeq, :, :])
      lossLine_t, = plt.plot(currExOutput, 'r', label="Target", linewidth = 1)
      plt.ylabel(self.sequence_tags[iSeq], fontsize=10)  
      plt.ylim((0, self.n_labels))
      plt.autoscale(enable=True, axis='x', tight=True)      
      if iSeq==0:
        plt.xlabel('Sample')
      
    savefig(self.tag + '_Target' + tag + '.pdf', bbox_inches='tight')
    plt.clf()
                
      
  def plot_output(self, tag=''):
    #nPlots = trainMimo.n_sequences()
    nSeq = min(self.n_sequences(), 6)
    print('Plotting Results')
    for iSeq in range(nSeq):
      plt.subplot(nSeq, 1, iSeq+1)
      currPrOutput = my_np.prob_to_class(self.block_output[iSeq, :, :])
      lossLine_t, = plt.plot(currPrOutput, 'r', label="Output", linewidth = 1)
      plt.ylabel(self.sequence_tags[iSeq], fontsize=10)  
      plt.ylim((0, self.n_labels))
      plt.autoscale(enable=True, axis='x', tight=True)      
      if iSeq==0:
        plt.xlabel('Sample')
      
    savefig(self.tag + '_Output' + tag + '.pdf', bbox_inches='tight')
    plt.clf()
                
      
  def plot_target_and_output(self, tag=''):
    #nPlots = trainMimo.n_sequences()
    nSeq = min(self.n_sequences(), 6)
    print('Plotting Results')
    for iSeq in range(nSeq):
      plt.subplot(nSeq, 1, iSeq+1)
      currExOutput = my_np.from_categorical(self.block_target[iSeq, :, :])
      lossLine_t, = plt.plot(currExOutput, 'r', label="Target", linewidth = 1)
      currPrOutput = my_np.prob_to_class(self.block_output[iSeq, :, :])
      lossLine_o, = plt.plot(currPrOutput, 'b-', label="Output", linewidth = 1)
      plt.ylabel(self.sequence_tags[iSeq], fontsize=10)
      plt.ylim((0, self.n_labels))
      plt.autoscale(enable=True, axis='x', tight=True)      
      if iSeq==0:
        first_legend = plt.legend(handles=[lossLine_t], loc=1)    
        ax = plt.gca().add_artist(first_legend)
        plt.legend(handles=[lossLine_o], loc=4)
        plt.xlabel('Sample')
      
    savefig(self.tag + '_TargetAndOutput' + tag + '.pdf', bbox_inches='tight')
    plt.clf()

  def plot_input(self, tag=''):
    nSeq = min(self.n_sequences(), 6)
    nFeatures = min(self.n_input_dim, 6)
    print('Plotting Input')
    for iSeq in range(nSeq):
      for jFeature in range(nFeatures):
        plt.subplot(nFeatures, 1, jFeature+1)
        currInput = self.block_input[iSeq, :, jFeature]
        plt.plot(currInput, 'b', linewidth = 1)
        #plt.ylabel(self.sequence_tags[iSeq])
        plt.ylabel('in' + str(jFeature+1))
        plt.xlabel('Sample')
        plt.autoscale(enable=True, axis='x', tight=True)      
    
      savefig(self.tag + '_'+ self.sequence_tags[iSeq] + '_Input' + tag + '.pdf', bbox_inches='tight')
      plt.clf()
    print('done')

  #
  #  note that with one-hot outputs,  (n_labels+1) probabilities 
  #  are written per output, so there are 
  #
  def write_predictions(self, filename):
    f = open(filename, 'w')
    for iSeq in range(len(self.sequence_tags)):
      f.write(self.sequence_tags[iSeq] + ';')
      seq_length_limit = min(self.sequence_lengths[iSeq]//self.out_dec_factor, self.max_length//self.out_dec_factor)
      for i in range(seq_length_limit):
        for j in range(self.block_output[iSeq, :, :].shape[1]):
          f.write(str(self.block_output[iSeq, i, j]) + ';')
      f.write('\n')
    f.close()  
    
  def set_block_output(self, block_output):
    self.block_output = block_output
    

  def do_classification(self):
    #---------------------------- target = np.argmax(self.block_target, axis=-1)
    #------------------------ prediction = np.argmax(self.block_output, axis=-1)
#------------------------------------------------------------------------------ 
    #-------------------------------------- flat_target = np.reshape(target, -1)
    #------------------------------ flat_prediction = np.reshape(prediction, -1)
    #----------------- flat_sample_weights = np.reshape(self.sample_weights, -1)
    #--------------------- class_report_str = classification_report(flat_target,
                                             #----------------- flat_prediction,
                                             #-- labels=range(0, self.n_labels),
                                             #-------- target_names=self.labels,
                                             # sample_weight=flat_sample_weights,
                                             #------------------------ digits=3)
#------------------------------------------------------------------------------ 
    #--------------------------------------------------- print(class_report_str)
    
    pass
    
  def do_classification_by_first_target_last_output(self):
#===============================================================================
#     predictions = np.argmax(self.block_output, axis=-1)
#     target = np.argmax(self.block_target, axis=-1)
# 
#     # for classification stats, select last prediction, 
#     # first target and first sample_weight of each sequence
#     #last_predictions = predictions[:, -1]
#     n_seq = len(self.sequence_lengths)
#     last_predictions = np.zeros((n_seq,), dtype=np.int)
#     for i in range(0, n_seq):
#       last_predictions[i] = predictions[i, self.sequence_lengths[i]//self.out_dec_factor -1]
#     
#     first_target = target[:, 0]
#     first_sample_weight = self.sample_weights[:, 0]
#     class_report_str = classification_report(first_target, 
#                                                 last_predictions, 
#                                                 labels=range(0, self.n_labels), 
#                                                 target_names=self.labels, 
#                                                 sample_weight=first_sample_weight, 
#                                                 digits=3)    
#         
#     print(class_report_str)
#     return last_predictions 
# #     file = open(test_mimo.tag + '_answers.txt', 'w')
# #     for i in range(0, n_seq):
# #       file.write('{0},{1}\n'.format(test_mimo.sequence_tags[i], test_mimo.labels[last_predictions[i]]))
# #     file.close()
#     
#     #self.shut_down()
#===============================================================================
    pass
          
    
    
if __name__ == "__main__":
  MimoSequence(sys.argv[1])
  #read("D:\\data\\studies\\xValEigen\\Perinatal\\xValEigen_Agg3_MI60_20161116T192046\\doLearnLSTM_EigenAll_klPc_20161220T170106\\Rep1\\RR1\\NN1\\auxTest1.nc")
