#
# based on a version coded ANN of FIMO's group
#
# modified by Mr. Nam @ 17-Feb-2022
# just find the best parameters for ANN model
#

import os, sys, itertools
import numpy as np
import pandas as pd

from sklearn.neural_network import MLPRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error

hidden_layer_sizes = [(100, 200, 300)]
test_size          = 0.22
n_iters            = 100
scoring            = 'rmse'
err_thresh         = 0
table_idx_prams    = []
table_error_prams  = []

#
# range params
#
param_grid = {'hidden_layer_sizes': hidden_layer_sizes,
	'activation': ['identity', 'logistic', 'tanh', 'relu'],
	'solver': ['lbfgs', 'adam', 'sgd'],
	'alpha': [0.0001],
	'batch_size': ['auto'],
	'learning_rate_init': [0.001],
	'max_iter': [200],
	'shuffle': [False],
}
keys, values = zip(*param_grid.items())
permutations_dicts = [dict(zip(keys, v)) for v in itertools.product(*values)]

#
# just for testing network
#
data = pd.read_csv('../data/HANOI_month=7_lead=0.txt', delimiter='\s+')
X    = data[['nino34','ao']]
y    = data['rainfall']
x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=None, shuffle=False) # shuffle=False disable random

# 
# check the number of loops
#
iterations = 0
if n_iters > len(permutations_dicts):
	iterations = len(permutations_dicts)
else:
	iterations = n_iters

#
# just loop
#
for i in range(iterations):
	print('iter {}'.format(i))
	cur_params = permutations_dicts[i]
	# print(cur_params)
	model = MLPRegressor(**cur_params)

	try:
		model.fit(x_train, y_train)
		y_train_pred = model.predict(x_train)
		y_test_pred = model.predict(x_test)

		if scoring == 'rmse':
			train_error = np.sqrt(mean_squared_error(y_train_pred, y_train))
			test_error = np.sqrt(mean_squared_error(y_test_pred, y_test))
		elif scoring == 'mae':
			train_error = mean_absolute_error(y_train_pred, y_train)
			test_error = mean_absolute_error(y_test_pred, y_test)
		else:
			print('invalid input for metric name')

		cur_err_drop = test_error - train_error
		if cur_err_drop <= 0 or cur_err_drop < err_thresh:
			print('found param set that satisfies the current threshold with error drop between test and train:', cur_err_drop)
			table_idx_prams.append(i)
			table_error_prams.append(cur_err_drop)

	except:
		print('param set produced a model that gives invalid values, therefore is disqualified')
		pass

#
# print out*
#
sys.stdout=open('list_prams.txt','w')

for i, j in enumerate(table_idx_prams):
	print(permutations_dicts[j],':',table_error_prams[i])

sys.stdout.close()

sys.exit()