# -*- coding: utf-8 -*-
"""
Created on Wed Jul 14 15:31:56 2021

@author: Mike
"""

# -*- coding: utf-8 -*-
"""
Created on Mon Aug  3 12:19:42 2020

@author: Mike
"""

import numpy as np
import pandas as pd
from datetime import datetime

from tensorflow.keras.layers import Dense, Input, LSTM, Embedding, Dropout, Activation, Flatten#, CuDNNLSTM

#TF 2 compatibility
#import keras
#from tensorflow.compat.v1.keras.layers import CuDNNLSTM

from tensorflow.keras.layers import Bidirectional, GlobalMaxPool1D, BatchNormalization, Conv1D, GlobalMaxPooling1D, MaxPooling1D
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, CSVLogger
from tensorflow.keras.models import model_from_json
from tensorflow.keras import metrics, regularizers
#from tensorflow.keras.utils import multi_gpu_model
from tensorflow.keras.optimizers import Nadam, Adam

from sklearn.metrics import mean_squared_error, mean_absolute_error
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from sklearn.model_selection import train_test_split as split

import tensorflow.keras.backend as K
import tensorflow.keras.callbacks as Kc



# TF2: Mixed precision

#from tensorflow.keras.mixed_precision import experimental as mixed_precision
#mixed_precision.set_policy('mixed_float16')

# TF2: Disable eager execution
import tensorflow as tf
"""
tf.compat.v1.disable_eager_execution()
tf.get_logger().setLevel('INFO')
"""
import time
from datetime import datetime
from random import randrange

#from PPM_CODE.Eval_helpers import*
#from PPM_CODE.HPO_searchspace import *


def GenSimpleModel(data_objects, model_params):
    # Clear the TF graph
    K.clear_session()
    
    # Load the data to check dimensionality
    x_train, y_train = data_objects["x_train"], data_objects["y_train"]
    print("Input data shape:",x_train.shape)
    
    ############################################
    
    # Number of block layers
    BLOCK_LAYERS = model_params["BLOCK_LAYERS"]
    
    # Number of cells
    FULLY_CONNECTED = model_params["FULLY_CONNECTED"]
    
    #CNN related params
    DROPOUT_RATE = model_params["DROPOUT_RATE"]
    
    
    #####################################
    
    input_dim = (x_train.shape[1], x_train.shape[2])
        
    model = Sequential()
    
    if BLOCK_LAYERS == 1:
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=False))
    
        
    # LSTM with batchNorm and recurrent dropout
    if BLOCK_LAYERS == 2:
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=False))
    
    # LSTM with batchNorm and recurrent dropout
    if BLOCK_LAYERS == 3:
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=False))
        
    
    # LSTM with batchNorm and recurrent dropout
    if BLOCK_LAYERS == 4:
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=True))
        model.add(BatchNormalization())
        
        model.add(LSTM(FULLY_CONNECTED,  implementation=2, 
                             recurrent_dropout=DROPOUT_RATE, 
                             input_shape=input_dim,
                             return_sequences=False))
        
    model.add(Dense(1)) #, dtype='float32' #Only the softmax is adviced to be float32 
    
    print(model.summary())
    
    return model, 0
