"""
Experiment 1.   - create first proposed model
                - predict 200 days of validation period

Experiment 2.   - predict 200 days of testing period
                - improve results by increasing LSTM units
                - improve results by replacing Z-Score norm. with MaxAbs norm.
"""

from src.constants import *
from src.stock import Stock
from src.model import Model
from src.helper import *
from sklearn.metrics import mean_squared_error

"""
for experiment 1. and 2., 200 days will be predicted (200 + 60 will be fetched - required for calculation)
standard 10y dataset should be used for training
validation period ranges from 2018-2019
testing period ranges from 2020-2021
"""
# data_sets with fixed will have fixed normalization
data_sets = ['validation_period', 'testing_period', 'validation-fixed_period', 'testing-fixed_period']
# these stocks will be predicted
reference_stocks = ['BRK-A', 'GOOG', 'MSFT']

for ref_stock in reference_stocks:
    # get stock prices and create a dataframe
    stock = Stock(ref_stock, START, END, FORCE_FETCH_STOCK)
    stock.fetch()

    training_validation_dataset = stock.df['Adj Close']
    training_dataset_copy = training_validation_dataset.copy()

    for data_set in data_sets:
        # determine validation period
        if data_set == 'validation_period':
            testing_start = datetime.datetime(2018, 1, 1)
            testing_end = datetime.datetime(2019, 1, 15)
        else:
            testing_start = datetime.datetime(2020, 1, 1)
            testing_end = datetime.datetime(2021, 1, 12)

        # get testing data
        stock_test = Stock(ref_stock, testing_start, testing_end, FORCE_FETCH_STOCK)
        stock_test.fetch()
        test_closing_prices = stock_test.df['Adj Close']

        # save original prices for later
        original_prices = test_closing_prices

        # split training validation
        training_size = int(len(training_validation_dataset) * TRAIN_VALIDATION_SPLIT)

        # create joined dataset for normalization
        whole_dataset = training_validation_dataset.append(test_closing_prices)

        # generate results with fixed and unfixed normalization
        if "fixed" in data_set:
            # split data ranges
            train_data = training_validation_dataset[training_size:]
            validation_data = training_validation_dataset[:training_size]
            # normalize dataset using MaxAbsScaler
            scaler, training_data, testing_data = normalize_dataset(train_data, validation_data, test_closing_prices)
        else:
            # for unfixed use z-score normalization
            mean = training_validation_dataset[:training_size].mean()
            std = training_validation_dataset[:training_size].std()
            training_data = (training_validation_dataset - mean) / std
            testing_data = (test_closing_prices - mean) / std

        # prepare datasets for model
        train_x, train_y = prepare_dataset_single_model(training_data, 0, training_size)
        validation_x, validation_y = prepare_dataset_single_model(training_data, training_size, None)
        testing_x, testing_y = prepare_dataset_single_model(testing_data, 0, None)

        # create datasets for model
        training_dataset, validation_dataset = create_datasets(train_x, train_y, validation_x, validation_y, training_data.size)

        # try these LSTM unit configurations
        if "fixed" in data_set:
            units = [512]  # fixed normalization will be generated with 512 units
        elif data_set == 'validation_period':
            units = [20]  # unfixed validation period will be generated with 20 units
        else:
            units = [20, 512]  # 20 - won't be sufficient for testing period (512 will fix that)

        # create models with different units
        for unit in units:
            model_name = ref_stock + '_' + str(unit) + '_' + str(data_set)
            # create model parameter class
            model_params = ModelParameters(training_dataset, validation_dataset, True, model_name, 'models/exp-1-2/', [unit, 0])
            # create model
            model = Model(model_params)
            # predict 200 days
            predictions = model.predict(testing_x)
            # reshape to 1d array
            predictions = np.array(predictions.reshape(-1))
            # denormalize predictions
            if "fixed" in data_set:
                scaled_predictions = scaler.inverse_transform(predictions.reshape(-1, 1)).reshape(-1)
            else:
                scaled_predictions = predictions * std + mean

            # plot results
            t_mse = model.history['mse'][-1]
            v_mse = model.history['val_mse'][-1]
            # calculate test mse
            test_mse = mean_squared_error(testing_y, predictions)

            # store results to csv
            fixed_normalization = "fixed" in data_set
            with open('experiment-results/exp-1-2/experiment-1-2.csv', 'a') as f:
                for i in range(len(predictions)):
                    test_value = testing_y[i][0] if fixed_normalization else testing_y[i]
                    f.write(
                        f"{model_name},{predictions[i]},{scaled_predictions[i]},{original_prices[60 + i]},{test_value},{str(fixed_normalization)},{t_mse},{v_mse},{str(test_mse)}\n")

            if not model.stored():
                model.save('-' + str(v_mse))

            if "fixed" in data_set:
                del scaler

        # reset training_validation dataset
        training_validation_dataset = training_dataset_copy
