import os
from typing import Union

import tensorflow as tf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split, KFold

import utility as ut
from variables import *


# Read the data.
train_data = pd.read_csv(os.path.join(DATA_PATH, ".".join([DATA_TRAIN, DATA_EXT])), header = 0)

# Get the labels.
Y = train_data.pop(LABEL)

sample_weights = np.ones(Y.shape[0])
for i in range(10, 24):
	sample_weights[train_data["_".join(("hour", str(i)))] == 1] = 1.5


#  -- For classification -- #
# CLASSES = np.unique(Y)
# N_CLASSES = len(CLASSES)
# Y = Y.replace(dict(zip(CLASSES, range(0, len(CLASSES)))))
 

# Data shape parameters.
N_FEATURES = train_data.shape[1]
N_SAMPLES = train_data.shape[0]

# Split the training data.
X_train, X_val, Y_train, Y_val = train_test_split(train_data, Y, shuffle = True, random_state = 7919)

def build_and_compile(input_: tuple = (WB_SIZE, N_FEATURES), 
                      loss_func: str = "mae") -> tf.keras.Model:
	"""
	Build and compile a TensorFLow LSTM network.
	
	Parameters
	----------
	input_ :
		Shape of the trainining data. Should specify 
		`(batch_size` or `window_size, n_features)`
	loss_func :
		Loss function to use for training.
  
	Returns
	-------
	`tf.keras.Model` :
		A compiled TensorFlow model.
	"""
 
	# Seqential keras model.
	model = tf.keras.models.Sequential([
		tf.keras.layers.LSTM(50, input_shape = input_, return_sequences = True),
		tf.keras.layers.LSTM(50, return_sequences = False),
		tf.keras.layers.GaussianNoise(1.0),
		tf.keras.layers.Dense(1024, activation = "relu"),
		tf.keras.layers.Dropout(0.7),
		tf.keras.layers.Dense(128, activation = "relu"),
		tf.keras.layers.Dropout(0.2),
		tf.keras.layers.Dense(64, activation = "relu"),
		tf.keras.layers.GaussianNoise(0.2),
		# tf.keras.layers.Dense(32, activation = "relu"),
		# tf.keras.layers.GaussianNoise(0.7),
		tf.keras.layers.Dense(1, activation = "relu")
	])
 
	# Compile the model.
	model.compile(
		loss = loss_func,
		optimizer = "adam"
	)
 
	return model

def train(model: tf.keras.Model,
          train_data: np.ndarray,
          train_labels: np.ndarray,
          val_data: np.ndarray,
          val_labels: np.ndarray,
		  epochs: int = 200,
		  sample_weights: np.array = None,
		  cross_val = False) -> pd.DataFrame:
	"""
	Trains the TensorFlow `model`.
 
	Parameters
	----------
	model :
		A TensorFlow compiled model.
	train_data :
		The data to be trained. Shape must be consistent with what is passed during model compilation.
	train_labels :
		The ground truth predictions.
	val_data :
		The data to be used as validation.
	val_labels : 
		The ground truth validation predictions.
	epochs :
		Total number of epochs to train.
	sample_weights :
		Weights for `train_data` to use during training.
  
	Returns
	-------
	pd.DataFrame:
		Training information.
	"""
 
	# Check for overfitting.
	early_stopping = tf.keras.callbacks.EarlyStopping(
		monitor =  "val_loss",
		min_delta = 0.001,
		patience = 100, 
		restore_best_weights = False)
 
	history = model.fit(
		train_data.reshape(-1, WB_SIZE, N_FEATURES), 
  		train_labels,
		sample_weight = sample_weights,
		validation_data = (val_data.reshape(-1, WB_SIZE, N_FEATURES), val_labels),
		verbose = 1, 
		epochs = epochs,
		callbacks = early_stopping)
 
	return pd.DataFrame(history.history)

# def cross_validate(train_data: pd.DataFrame,
#                    train_labels: pd.DataFrame,
#                    epochs: int = 50,
#                    sample_weights: np.array = None,
#                    folds: int = 2) -> pd.DataFrame:
    
# 	splits = KFold(n_splits = folds, shuffle = True)
# 	print("Starting cross validation.")
# 	accuracy = list()
# 	val_loss = list()
# 	models = list()
# 	for i, (train_index, test_index) in enumerate(splits.split(train_data, train_labels)):
# 		print(f"Iteration {i}\n")
# 		X_train, X_val, Y_train, Y_val = train_data[train_index], train_data[test_index], train_data[train_index], train_labels[test_index]
  
# 		model = build_and_compile((WB_SIZE, N_FEATURES), "mae")
  
# 		history_df = train(model, X_train, Y_train, epochs)
# 		# train_stats(history_df, i)
		
# 		scores = model.evaluate(X_val.reshape(-1, WB_SIZE, N_FEATURES), Y_val)
# 		print(f"Validation loss: {scores}\n")
#         #of {scores[0]} {model.metrics_names[1]} of {scores[1] * 100:.2f}%")

# 		# accuracy.append(scores[1] * 100)
# 		val_loss.append(scores)
# 		models.append(model)
  
# 	return models[np.argmin(val_loss)]
	 
def train_stats(history_df: pd.DataFrame, it: int = None) -> None:
	"""
	Produces training statistics once training has run its course.
	
	Parameters
	----------
	history_df :
		The history as returned by Keras `fit` method.
	it :
		To be used with cross validation. Specifies the name of the learning curve based on the cross validation itertation `it`. 
  
	Returns
	-------
	`None`
	
	"""
	# Learning curve.
	plt.rcParams["figure.dpi"] = 160
	history_df.loc[:, ["loss", "val_loss"]].plot()
	plt.title("Model Loss")
	plt.ylabel("Loss")
	plt.xlabel("Epoch")
	name = TRAIN_FIG_SAVE_NAME
	if it is not None:
		name = "_".join([name, str(it)])
	plt.savefig(os.path.join(TRAIN_FIG_SAVE_PATH, ".".join([name, FIG_EXT])))
 
	# Stats 
	print(f"Minimum validation loss: {history_df['val_loss'].min()}")
	# plt.plot(f"Accuracy: {history_df['train_accuracy']}")
	# plt.plot(f"Validation Accuracy: {history_df['val_accuracy']}")
 
	return None

def main():
	model = build_and_compile((WB_SIZE, N_FEATURES))
	# model = cross_validate(np.array(train_data), np.array(Y))
	history_df = train(model, np.array(X_train), np.array(Y_train), np.array(X_val), np.array(Y_val))
	
	# train_stats(history_df)
 
	# Save trained model (better to use checkpoints).
	model.save(os.path.join(NN_MODEL_SAVE_PATH, NN_MODEL_SAVE_NAME))

if __name__ == "__main__":
	main()

[((21, 17, 21, 36), 'numpy.ones', 'np.ones', ({(21, 25, 21, 35): 'Y.shape[0]'}, {}), '(Y.shape[0])', True, 'import numpy as np\n'), ((37, 33, 37, 101), 'sklearn.model_selection.train_test_split', 'train_test_split', (), '', False, 'from sklearn.model_selection import train_test_split, KFold\n'), ((117, 18, 121, 31), 'tensorflow.keras.callbacks.EarlyStopping', 'tf.keras.callbacks.EarlyStopping', (), '', True, 'import tensorflow as tf\n'), ((132, 8, 132, 37), 'pandas.DataFrame', 'pd.DataFrame', ({(132, 21, 132, 36): 'history.history'}, {}), '(history.history)', True, 'import pandas as pd\n'), ((183, 1, 183, 24), 'matplotlib.pyplot.title', 'plt.title', ({(183, 11, 183, 23): '"""Model Loss"""'}, {}), "('Model Loss')", True, 'import matplotlib.pyplot as plt\n'), ((184, 1, 184, 19), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(184, 12, 184, 18): '"""Loss"""'}, {}), "('Loss')", True, 'import matplotlib.pyplot as plt\n'), ((185, 1, 185, 20), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(185, 12, 185, 19): '"""Epoch"""'}, {}), "('Epoch')", True, 'import matplotlib.pyplot as plt\n'), ((201, 27, 201, 44), 'numpy.array', 'np.array', ({(201, 36, 201, 43): 'X_train'}, {}), '(X_train)', True, 'import numpy as np\n'), ((201, 46, 201, 63), 'numpy.array', 'np.array', ({(201, 55, 201, 62): 'Y_train'}, {}), '(Y_train)', True, 'import numpy as np\n'), ((201, 65, 201, 80), 'numpy.array', 'np.array', ({(201, 74, 201, 79): 'X_val'}, {}), '(X_val)', True, 'import numpy as np\n'), ((201, 82, 201, 97), 'numpy.array', 'np.array', ({(201, 91, 201, 96): 'Y_val'}, {}), '(Y_val)', True, 'import numpy as np\n'), ((206, 12, 206, 64), 'os.path.join', 'os.path.join', ({(206, 25, 206, 43): 'NN_MODEL_SAVE_PATH', (206, 45, 206, 63): 'NN_MODEL_SAVE_NAME'}, {}), '(NN_MODEL_SAVE_PATH, NN_MODEL_SAVE_NAME)', False, 'import os\n'), ((60, 2, 60, 73), 'tensorflow.keras.layers.LSTM', 'tf.keras.layers.LSTM', (), '', True, 'import tensorflow as tf\n'), ((61, 2, 61, 52), 'tensorflow.keras.layers.LSTM', 'tf.keras.layers.LSTM', (), '', True, 'import tensorflow as tf\n'), ((62, 2, 62, 36), 'tensorflow.keras.layers.GaussianNoise', 'tf.keras.layers.GaussianNoise', ({(62, 32, 62, 35): '1.0'}, {}), '(1.0)', True, 'import tensorflow as tf\n'), ((63, 2, 63, 50), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((64, 2, 64, 30), 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', ({(64, 26, 64, 29): '0.7'}, {}), '(0.7)', True, 'import tensorflow as tf\n'), ((65, 2, 65, 49), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((66, 2, 66, 30), 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', ({(66, 26, 66, 29): '0.2'}, {}), '(0.2)', True, 'import tensorflow as tf\n'), ((67, 2, 67, 48), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n'), ((68, 2, 68, 36), 'tensorflow.keras.layers.GaussianNoise', 'tf.keras.layers.GaussianNoise', ({(68, 32, 68, 35): '0.2'}, {}), '(0.2)', True, 'import tensorflow as tf\n'), ((71, 2, 71, 47), 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (), '', True, 'import tensorflow as tf\n')]