"""
Experiment 5. - find best learning rate and epochs for proposed architectures (train 72 models)
              - train 480 models, predict 200 days for each model from testing period, so they can be evaluated
"""
from src.constants import *
from src.stock import Stock
from src.model import *
from src.helper import *
import pandas_market_calendars as mcal
import copy


def plot_data(y1, y2, title, xlabel, ylabel, name, log, y1_label='', y2_label=''):
    """
    Custom plot method for plotting learning rate and epochs optimization
    :param y1: 
    :param y2: 
    :param title: 
    :param xlabel: 
    :param ylabel: 
    :param name: 
    :param log: 
    :param y1_label: 
    :param y2_label: 
    :return: 
    """
    # plot training and validation losses
    plt.figure(figsize=(18, 8))
    # epochs on x axis
    plot_epochs = np.arange(1, config['epochs'] + 1)
    if not log:
        # plot both y vars
        plt.plot(plot_epochs, y1, label=y1_label, lw=3)
        plt.plot(plot_epochs, y2, label=y2_label, lw=3)
    else:
        learning_rates = 1e-4 * (10 ** (np.arange(config['epochs']) / 100))
        plt.semilogx(learning_rates, y1, lw=3, label=y1_label)
    # set title and labels
    plt.title(title, size=28)
    plt.xlabel(xlabel, size=24, labelpad=15)
    plt.ylabel(ylabel, size=24, labelpad=15)
    # increase xticks and yticks size
    plt.xticks(fontsize=22)
    plt.yticks(fontsize=22)
    # print legend, save and show
    plt.legend(fontsize=24)
    plt.tight_layout(pad=1)
    plt.savefig('experiment-results/exp-5/' + name)
    plt.show()


# experiment parameters
training_tries = 10
units = 512
BATCH = 30
required_columns = ['Adj Close', 'Volume']
# if true will create 72 models for finding best learning rate and epochs
# else will create 480 models for testing
LR_TESTING = True
epochs_list = [65, 58, 52, 56] if not LR_TESTING else [100, 100, 100, 100]

# proposed architectures
configs = [
    {'model_type': '-single-', 'epochs': epochs_list[0], 'alpha': 0.0004, 'arch': [units, 0], 'name': 'single-65-0.0004-1-layer-'},

    {'model_type': '-multi-', 'epochs': epochs_list[1], 'alpha': 0.00037, 'arch': [units, 0], 'name': 'multi-58-0.00037-1-layer-'},

    {'model_type': '-single-', 'epochs': epochs_list[2], 'alpha': 0.00031, 'arch': [units, units], 'name': 'single-52-0.00031-2-layer-'},

    {'model_type': '-multi-', 'epochs': epochs_list[3], 'alpha': 0.00031, 'arch': [units, units], 'name': 'multi-56-0.00031-2-layer-'},
]

# data ranges 10, 20y datasets with and without early stopping
data_ranges = [
    {'start': datetime.datetime(2010, 1, 1), 'early_stop': False, 'name': '10y'},
    {'start': datetime.datetime(2000, 1, 1), 'early_stop': False, 'name': '20y'},
]

# add early stopping for testing mode
if not LR_TESTING:
    data_ranges.append({'start': datetime.datetime(2010, 1, 1), 'early_stop': True, 'name': '10y-early-stop'})
    data_ranges.append({'start': datetime.datetime(2000, 1, 1), 'early_stop': True, 'name': '20y-early-stop'})

# get trading days
nasdaq = mcal.get_calendar('NASDAQ')
trading_days = nasdaq.schedule(start_date='2022-11-01', end_date='2022-11-30')['market_open'].dt.date

for config in configs:
    model_type = config['model_type']
    ARCH = config['arch']

    for data_range in data_ranges:
        START = data_range['start']
        REAL_START = data_range['start']
        EARLY_STOPPING = data_range['early_stop']

        if model_type == '-single-':
            TESTING_START = datetime.datetime(2020, 10, 6)
        else:
            TESTING_START = datetime.datetime(2020, 8, 24)

        TESTING_END = datetime.datetime(2022, 12, 24)

        for stock_name in ['BRK-A', 'GOOG', 'MSFT']:
            # let's use last 500 trading days (testing) and rest for training (80/20 split)
            # get stock prices and create a dataframe for both training and testing
            stock = Stock(stock_name, START, TESTING_END, FORCE_FETCH_STOCK)
            stock.fetch()

            whole_dataset = stock.df[required_columns]

            # split training and testing
            training_dataset = whole_dataset[whole_dataset.index < TESTING_START]
            # split training validation
            training_size = int(len(training_dataset) * TRAIN_VALIDATION_SPLIT)
            validation_data = training_dataset[training_size:]
            training_data = training_dataset[:training_size]
            testing_data = whole_dataset[whole_dataset.index >= TESTING_START]

            # normalize data
            scaler, training_dataset, testing_data = normalize_dataset(training_data, validation_data, testing_data, required_columns)

            # prepare data
            if model_type == '-single-':
                train_x, train_y = prepare_dataset_single_model(training_dataset, 0, training_size)
                validation_x, validation_y = prepare_dataset_single_model(training_dataset, training_size, None)
                testing_x, testing_y = prepare_dataset_single_model(testing_data, 0, None)
            else:
                train_x, train_y = prepare_dataset_multi_model(training_dataset, 0, training_size, ['MMS', 'VOL'])
                validation_x, validation_y = prepare_dataset_multi_model(training_dataset, training_size, False, ['MMS', 'VOL'])
                testing_x, testing_y = prepare_dataset_multi_model(testing_data, 0, None, ['MMS', 'VOL'])

            # create tf datasets
            training_dataset, validation_dataset = create_datasets(train_x, train_y, validation_x, validation_y, whole_dataset.size)
            # store created models (train each config. 10 models)
            models = []
            for i in range(1, training_tries + 1):
                # create model
                model_name = str(i) + '-' + stock_name.lower() + model_type + str(units) + '-'
                # model params
                model_params = ModelParameters(training_dataset, validation_dataset, RETRAIN, model_name, 'models/exp-5/',
                                               ARCH, True, config['epochs'], config['alpha'], data_range['early_stop'],
                                               True, LR_TESTING)
                # init model
                model = Model(model_params)
                # store model copy
                models.append(copy.copy(model))
                # save model if not loaded
                if not model.loaded:
                    model.save(str(model.history['val_mse'][-1]))

            # store model statistics
            for idx, model in enumerate(models):
                # some models will produce too big values so graph will be unreadable -> cap them on certain value
                cap_val = 0.03 if stock_name == 'MSFT' else 0.001
                val_loss = [cap_val if x > cap_val else x for x in model.history['val_loss']]
                # find at which epoch is val_loss minimum
                min_epoch = np.argmin(val_loss)
                print('Minimum validation loss: ' + str(val_loss[min_epoch]))
                print('Minimum validation loss epoch: ' + str(min_epoch + 1))

                if LR_TESTING:
                    print('Learning rate at min epoch: ' + str(model.history['lr'][min_epoch]))
                    # store results into learning rate csv
                    with open('experiment-results/exp-5/experiment-5-full-lr-test.csv', 'a') as f:
                        f.write(
                            f"{stock_name},{config['name'] + str(idx + 1)},{val_loss[min_epoch]},{min_epoch + 1},{model.history['lr'][min_epoch]},{model.history['val_mse'][-1]}\n")

                    # cap loss at cap_val
                    loss = [cap_val if x > cap_val else x for x in model.history['loss']]
                    graph_name = stock_name + '-' + config['name'] + str(idx + 1) + '-val-vs-loss.png'
                    plot_data(loss, val_loss, 'Porovnanie trénovacej a validačnej odchylky', 'Epocha', 'Odchylka (MSE)', graph_name, False, 'Trénovacia odchylka',
                              'Validačná odchylka')
                    # plot learning rate vs. validation loss
                    graph_name = stock_name + '-' + config['name'] + str(idx + 1) + '-lr.png'
                    plot_data(val_loss, [], 'Porovnanie miery učenia a validačnej odchylky', 'Miera učenia (log)', 'Validačná odchylka (MSE)', graph_name, True,
                              'Validačná odchylka')

                else:
                    # store only partial learning rate results
                    with open('experiment-results/exp-5/experiment-5-lr-log.csv', 'a') as f:
                        f.write(
                            f"{stock_name},{config['name'] + str(idx + 1)},{val_loss[min_epoch]},{min_epoch + 1},{model.history['val_mse'][-1]}\n")

                    # store all model metrics
                    with open('experiment-results/exp-5/experiment-5.csv', 'a') as f:
                        t_rms = model.history['root_mean_squared_error'][-1]
                        t_mae = model.history['mae'][-1]
                        t_mse = model.best_loss if EARLY_STOPPING else model.history['mse'][-1]
                        v_rms = model.history['val_root_mean_squared_error'][-1]
                        v_mae = model.history['val_mae'][-1]
                        v_mse = model.best_val_loss if EARLY_STOPPING else model.history['val_mse'][-1]
                        epoch = model.best_epoch + 1 if EARLY_STOPPING else config['epochs']

                        f.write(
                            f"{t_rms},{t_mae},{t_mse},{v_rms},{v_mae},{v_mse},{epoch},{config['name'] + data_range['name'] + '-' + stock_name}\n")
