"""
Experiment 3. - find optimal model features based on packages
"""
from src.constants import *
from src.stock import Stock
from src.model import *
from src.helper import *

# overall 16 different feature selections will be evaluated
packages = [
    {'names': ['MMS']},
    {'names': ['MMS', 'MA']},
    {'names': ['MMS', 'MA', 'VOL']},
    {'names': ['MMS', 'VOL']},
    {'names': ['MMS', 'XMA']},
    {'names': ['MMS', 'XMA', 'VOL']},
    {'names': ['VOL', 'MA']},
    {'names': ['VOL', 'XMA']},
]

# standard training dataset of 10ys will be used
# proposed model features will be evaluated on testing period predicting 200 days
TESTING_START = datetime.datetime(2020, 1, 1)
# baseline for testing end is 12.1.2021 but multivariate model need extra 30 days for computation
TESTING_END = datetime.datetime(2021, 2, 26)
# columns to be kept from fetched data
required_columns = ['Adj Close', 'Volume']

for ref_stock in ['BRK-A', 'GOOG', 'MSFT']:
    # get stock price data
    stock = Stock(ref_stock, START, TESTING_END, FORCE_FETCH_STOCK)
    stock.fetch()

    for main_feature in ['price', 'returns']:
        if main_feature == 'returns':
            # keep volume, calculate returns and remove first now with NAN value
            whole_dataset = pd.DataFrame([stock.df['Adj Close'].pct_change()[1:], stock.df['Volume'][1:]], index=required_columns).T
        else:
            # remove first row so both datasets have same length
            whole_dataset = stock.df[required_columns][1:]

        # 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 datasets
        scaler, training_dataset, testing_data = normalize_dataset(training_data, validation_data, testing_data, required_columns)

        # save original prices for later
        original_prices = whole_dataset[whole_dataset.index >= TESTING_START].copy()['Adj Close']

        for package in packages:
            # prepare datasets for model
            train_x, train_y = prepare_dataset_multi_model(training_dataset, 0, training_size, package['names'])
            validation_x, validation_y = prepare_dataset_multi_model(training_dataset, training_size, None, package['names'])
            testing_x, testing_y = prepare_dataset_multi_model(testing_data, 0, None, package['names'])

            # create tf datasets
            tf_training_dataset, tf_validation_dataset = create_datasets(train_x, train_y, validation_x, validation_y, training_data.size)

            current_packages = '-'.join(package['names'])
            # run 3 models for each package
            for model_idx in range(1, 4):
                model_name = ref_stock.lower() + '-' + str(model_idx) + '-multi-' + main_feature + '-' + current_packages
                # create param class
                model_params = ModelParameters(tf_training_dataset, tf_validation_dataset, RETRAIN, model_name, 'models/exp-3/', [512, 0])
                # init model
                model = Model(model_params)
                # predict 200 days from testing period
                predictions = model.predict(testing_x)
                # prepare for rescale
                predictions = predictions
                dummy_volume = np.zeros(predictions.shape[0]).reshape(-1, 1)
                # denormalize predictions
                scaled_predictions = scaler.inverse_transform(np.concatenate((predictions, dummy_volume), axis=1))[:, 0]
                # plot results
                price_offset = BATCH_SIZE + WINDOW_SIZE
                y_label = 'Cena ($)' if main_feature == 'price' else 'Medzidňová návratnosť'
                plot_predictions(scaled_predictions, original_prices[price_offset:], ref_stock, 'experiment-3-{}.png'.format(main_feature), 'exp-3', y_label)
                # store model and it's results if not already stored
                if not model.stored():
                    t_mse = model.history['mse'][-1]
                    v_mse = model.history['val_mse'][-1]
                    model_name += f"-{t_mse}-{v_mse}"
                    # store results
                    with open('experiment-results/exp-3/experiment-3.csv', 'a') as f:
                        for i in range(len(predictions)):
                            f.write(
                                f"{model_name},{predictions[i]},{scaled_predictions[i]},{original_prices[price_offset + i]},{testing_y[i]}\n")
                    model.save(str(v_mse))
