#!/usr/bin/env python

# Copyright 2017 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance 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.

# original code from: https://github.com/GoogleCloudPlatform/training-data-analyst/tree/master/blogs/timeseries
# modified by: Marianne Linhares, monteirom@google.com, May 2017

# tensorflow
import tensorflow as tf
import tensorflow.contrib.learn as tflearn
import tensorflow.contrib.layers as tflayers
from tensorflow.contrib.learn.python.learn import learn_runner
import tensorflow.contrib.metrics as metrics
import tensorflow.contrib.rnn as rnn

# helpers
import numpy as np
import csv

# enable tensorflow logs
tf.logging.set_verbosity(tf.logging.INFO)

"""
Explanation: Shakespeare generation using RNNs + Estimators
This notebook illustrates how to:
   1. Creating a Recurrent Neural Network in TensorFlow
   2. Creating a Custom Estimator in tf.contrib.learn
Dependecies
End of explanation
"""


# load file
file_path = 'shakespeare.txt'
with open(file_path,'r') as f:
    data = f.read()
    print("Data length:", len(data))

# make all letters lower case 
# it makes the problem simpler
# but we'll loose information
data = data.lower()

"""
Explanation: Getting the data
In this notebook we'll use shakespeare data, but you can basically choose any text file you like!
End of explanation
"""


# vocabulary
vocab = list(set(data))
vocab_size = len(vocab)
print vocab

# embeding words to one-hot
def text_to_onehot(data_, vocab):
    data = np.zeros((len(data_), len(vocab)))
    cnt = 0
    for s in data_:
        v = [0.0] * len(vocab)
        v[vocab.index(s)] = 1.0
        data[cnt, :] = v
        cnt += 1
    return data

def onehot_to_text(array, vocab):
    return vocab[ array.index(1) ]

data_embed = text_to_onehot(data, vocab)

"""
Explanation: Preprocess data
End of explanation
"""


BATCH_SIZE = 20
TIMESERIES_COL = 'data'

# how many characters should be used by the LSTM as input
N_INPUTS = 10

# how many characters should the LSTM predict
N_OUTPUTS = 1

# -------- read data and convert to needed format -----------
def read_dataset(filename, mode=tf.estimator.ModeKeys.TRAIN):  
  def _input_fn():
    num_epochs = 100 if mode == tf.estimator.ModeKeys.TRAIN else 1

    # could be a path to one file or a file pattern.
    input_file_names = tf.train.match_filenames_once(filename)
    filename_queue = tf.train.string_input_producer(
        input_file_names, num_epochs=num_epochs, shuffle=True)

    reader = tf.TextLineReader()
    _, value = reader.read_up_to(filename_queue, num_records=BATCH_SIZE)

    value_column = tf.expand_dims(value, -1)
    print('readcsv={}'.format(value_column))
    
    # all_data is a list of tensors
    all_data = tf.decode_csv(value_column, record_defaults=DEFAULTS)  
    inputs = all_data[:len(all_data)-N_OUTPUTS]  # first few values
    label = all_data[len(all_data)-N_OUTPUTS : ] # last few values
    
    # from list of tensors to tensor with one more dimension
    inputs = tf.concat(inputs, axis=1)
    label = tf.concat(label, axis=1)
    print(inputs)
    print('inputs={}'.format(inputs))
    
    return {TIMESERIES_COL: inputs}, label   # dict of features, label
  return _input_fn


"""
Explanation: Read datasets
End of explanation
"""


LSTM_SIZE = 3  # number of hidden layers in each of the LSTM cells

def simple_rnn(features, targets, mode, params):
  # 0. Reformat input shape to become a sequence
  x = tf.split(features[TIMESERIES_COL], N_INPUTS, 1)
    
  # 1. configure the RNN
  lstm_cell = rnn.BasicLSTMCell(LSTM_SIZE, forget_bias=1.0)
  outputs, _ = rnn.static_rnn(lstm_cell, x, dtype=tf.float32)

  # slice to keep only the last cell of the RNN
  outputs = outputs[-1]
  #print 'last outputs={}'.format(outputs)
  
  # output is result of linear activation of last layer of RNN
  weight = tf.Variable(tf.random_normal([LSTM_SIZE, N_OUTPUTS]))
  bias = tf.Variable(tf.random_normal([N_OUTPUTS]))
  predictions = tf.matmul(outputs, weight) + bias
    
  # 2. Define the loss function for training/evaluation
  #print 'targets={}'.format(targets)
  #print 'preds={}'.format(predictions)
  loss = tf.losses.mean_squared_error(targets, predictions)
  eval_metric_ops = {
      "rmse": tf.metrics.root_mean_squared_error(targets, predictions)
  }
  
  # 3. Define the training operation/optimizer
  train_op = tf.contrib.layers.optimize_loss(
      loss=loss,
      global_step=tf.contrib.framework.get_global_step(),
      learning_rate=0.01,
      optimizer="SGD")

  # 4. Create predictions
  predictions_dict = {"predicted": predictions}
  
  # 5. return ModelFnOps
  return tflearn.ModelFnOps(
      mode=mode,
      predictions=predictions_dict,
      loss=loss,
      train_op=train_op,
      eval_metric_ops=eval_metric_ops)

def serving_input_fn():
    feature_placeholders = {
        TIMESERIES_COL: tf.placeholder(tf.float32, [None, N_INPUTS])
    }
  
    features = {
      key: tf.expand_dims(tensor, -1)
      for key, tensor in feature_placeholders.items()
    }

    return tflearn.utils.input_fn_utils.InputFnOps(
      features,
      None,
      feature_placeholders
    )

"""
Explanation: RNN Model
End of explanation
"""


nn = tf.contrib.learn.Estimator(model_fn=simple_rnn)

# ---------- Training -------------
print('---------- Training ------------')
nn.fit(input_fn=get_train(), steps=10000)

# ---------- Evaluating -------------
print('---------- Evaluating ------------')
ev = nn.evaluate(input_fn=get_valid())
print(ev)

# ---------- Testing ----------------
print('---------- Testing ------------')
predictions = []
for p in nn.predict(input_fn=get_test()):
    print(p)
    predictions.append(p["predicted"])

"""
Explanation: Running model
End of explanation
"""


# read test csv
def read_csv(filename):
    with open(filename, 'rt') as csvfile:
        reader = csv.reader(csvfile)
        data = []
        for row in reader:
            data.append([float(x) for x in row])
        return data

test_data = read_csv('test.csv')

# update predictions with features
# preds = test_data[:INPUTS] concat with predictions
preds = [] 
for i in range(len(predictions)):
    preds.append(list(test_data[i][:N_INPUTS]) + list(predictions[i]))

# visualizing predictions
for d in test_data: sns.tsplot(d[N_INPUTS:])
for p in preds: sns.tsplot(p[N_INPUTS:], color="red")
plt.show()    
    
# visualizing all the series
for d in test_data: sns.tsplot(d)
for p in preds: sns.tsplot(p, color="red")
plt.show()

"""
Explanation: Visualizing predictions
End of explanation
"""


