from __future__ import absolute_import, division, print_function, unicode_literals

# TensorFlow and tf.keras
import tensorflow as tf
from tensorflow import keras
from tensorflow import feature_column
from tensorflow.keras import layers
from sklearn.model_selection import train_test_split


# Helper libraries
import numpy as np
import pandas as pd
import pathlib
import os
import sys
import argparse

import nnutils


# setup and parse cli arguments
parser = argparse.ArgumentParser(description='Back propagation NN that trains '
                                 'on data generated by the MagRss app')
parser.add_argument('--train_dir', type=str, default='m1')
parser.add_argument('--train_file', type=str, default='master_0.csv')
parser.add_argument('--dbm', action='store_true', help='set this flag to '
                    'include dbm entries into the featureset')
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--batch_size', type=int, default=20)
parser.add_argument('--early_stopping', nargs=2, default=argparse.SUPPRESS, 
                    metavar=('MONITOR', 'PATIENCE'))
parser.add_argument('--features', nargs='+', type=str, 
                    default=argparse.SUPPRESS)
parser.add_argument('--layer', nargs=4, dest='layers', action='append', 
                    metavar=('LAYER_TYPE', 'NUMBER_OF_NODES', 
                             'KERNEL_REGULARIZER', 'DROPUT'))
parser.add_argument('--tensorboard', action='store_true', 
                    help='set this flag to log history/scalar data from the ' 
                         'model\'s run')
args = vars(parser.parse_args())
if args['layers'] is None:
    args['layers'] = [
        ['dense', '128', '0.001', '0.2'], ['dense', '128', '0.001', '0.2']] 
print(args)

print('---------------------------------------------------------')
train_files = nnutils._get_clean_train_filepaths(args['train_dir'])
# print(train_files)
# print(args['train_file'])
# exit(0)
for train_file in train_files:
    if str(train_file) == args['train_file']:
        df = pd.read_csv(train_file)
df.fillna(0, inplace = True)

print(df.columns)

# include/exclude cid dbm fields
if not args['dbm']:
    dbm_columns = [column for column in df.columns if 'dbm' in column]
    df.drop(columns=dbm_columns, inplace=True)

# set of features to be used for training
if 'features' in args:
    f_col = [column for column in df.columns
             if column not in args['features'] 
             and 'dbm' not in column and column != 'label']
    df.drop(columns=f_col, inplace=True)
print(df.columns)

train, val = train_test_split(df, test_size=0.2)
# train, val = train_test_split(train, test_size=0.2)


# setup data arrays from pandas dataframe
batch_size = args['batch_size']
train_ds = nnutils.dataframe_to_dataset_input_fn(
    df=train, batch_size=batch_size)
val_ds = nnutils.dataframe_to_dataset_input_fn(
    df=val, shuffle=False, batch_size=batch_size)
# test_ds = nnutils.dataframe_to_dataset_input_fn(
#     df=test, shuffle=False, batch_size=batch_size)
steps_per_epoch = len(train.index) // batch_size
validation_steps = len(val) // batch_size

# setup the features as numeric
feature_columns = [feature_column.numeric_column(str(f)) for f in train.columns
                   if 'label' not in f]
feature_layer = tf.keras.layers.DenseFeatures(feature_columns)

lr_schedule = tf.keras.optimizers.schedules.InverseTimeDecay(
    0.001, 
    decay_steps = steps_per_epoch * 1000, 
    decay_rate=1, 
    staircase=False)
optimizer = tf.keras.optimizers.Adam(lr_schedule)

# setup callbacks
callbacks = []
if 'early_stopping' in args:
    monitor = str(args['early_stopping'][0])
    patience = int(args['early_stopping'][1])
    callbacks.append(
        tf.keras.callbacks.EarlyStopping(monitor=monitor, patience=patience))

# if args['tensorboard']:
#     callbacks.append(nnutils.get_tensorboard_log_callback(train.columns, args))


layers = args['layers'].copy()
h = {}
model = nnutils.get_sequential_model(feature_layer=feature_layer, layers=layers)
model.compile(optimizer=optimizer, loss='binary_crossentropy', 
    metrics=['accuracy', 'binary_crossentropy'])
h['drop'] = model.fit(
    train_ds, validation_data=val_ds, epochs=args['epochs'], 
    steps_per_epoch=steps_per_epoch, validation_steps=validation_steps,
    callbacks=callbacks, verbose=2)
# print(h['drop'])
# model.summary()
# nnutils.plot_train_history(h['drop'], 'drop')
pandas.read_csv tensorflow.keras.layers.DenseFeatures sklearn.model_selection.train_test_split tensorflow.keras.optimizers.schedules.InverseTimeDecay tensorflow.keras.optimizers.Adam tensorflow.keras.callbacks.EarlyStopping