"""
Experiment 4. - take closer look at month predictions
              - create 60 models and compare their performance
              - alternative evaluation method
"""
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
import pandas_market_calendars as mcal

# training results that will be stored 
test_stats = []
model_stats = []
# experiment parameters
model_path = 'models/exp-4/'
data_offset = BATCH_SIZE + WINDOW_SIZE
models_to_train = 10

# setup testing period of 499 days
TESTING_START = datetime.datetime(2020, 8, 23)
TESTING_END = datetime.datetime(2022, 12, 25)

# get all 2022 trading days + 60 days for first batch + 30 days for multivariate model
nasdaq = mcal.get_calendar('NASDAQ')
trading_days = nasdaq.schedule(start_date='2020-08-21', end_date='2022-12-25')['market_open'].dt.date
# expected columns from stock data
required_columns = ['Adj Close', 'Volume']

for loop_stock in ['BRK-A', 'MSFT', 'GOOG']:
    # get stock prices and create a dataframe (standard 10y dataset)
    stock = Stock(loop_stock, START, END, FORCE_FETCH_STOCK)
    stock.fetch()
    training_validation_dataset = stock.df[required_columns]
    # get testing data
    stock_test = Stock(loop_stock, TESTING_START, TESTING_END, FORCE_FETCH_STOCK)
    stock_test.fetch()
    testing_dataset = stock_test.df[required_columns]

    # split training validation
    training_size = int(len(training_validation_dataset) * TRAIN_VALIDATION_SPLIT)
    training_data = training_validation_dataset[:training_size]
    validation_data = training_validation_dataset[training_size:]

    # normalize dataset
    scaler, training_validation_dataset, testing_dataset = normalize_dataset(training_data, validation_data, testing_dataset, required_columns)

    for model_type in ['single_variable', 'multi_variable']:
        # create windows of BATCH_SIZE days
        if model_type == 'single_variable':
            train_x, train_y = prepare_dataset_single_model(training_validation_dataset, 0, training_size)
            validation_x, validation_y = prepare_dataset_single_model(training_validation_dataset, training_size, None)
            testing_x, testing_y = prepare_dataset_single_model(testing_dataset, 0, None)
            testing_x, testing_y = testing_x[WINDOW_SIZE:], testing_y[WINDOW_SIZE:]  # remove first 30 days so length matches multi variable model
        else:
            train_x, train_y = prepare_dataset_multi_model(training_validation_dataset, 0, training_size, ['MMS', 'VOL'])
            validation_x, validation_y = prepare_dataset_multi_model(training_validation_dataset, training_size, None, ['MMS', 'VOL'])
            testing_x, testing_y = prepare_dataset_multi_model(testing_dataset, 0, None, ['MMS', 'VOL'])

        # create tf datasets
        training_dataset, validation_dataset = create_datasets(train_x, train_y, validation_x, validation_y, training_validation_dataset.size)

        mtype = '-single-' if model_type == 'single_variable' else '-multi-'

        # create 10 models of chosen type
        for i in range(1, models_to_train + 1):
            model_name = str(i) + '-' + loop_stock.lower() + mtype + str(512) + '-'
            # model params
            model_params = ModelParameters(training_dataset, validation_dataset, False, model_name, model_path, [512, 0])
            # init model
            model = Model(model_params)
            # predict testing period
            predictions = model.predict(testing_x)
            # reshape to 1d array
            predictions = np.array(predictions.reshape(-1))

            # store model stats
            for prediction_id, prediction in enumerate(predictions):
                if prediction_id == 0:
                    continue

                # denormalize
                denormalized_prediction = scaler.inverse_transform([[prediction, 0]])[0][0]
                # get previous day price from testing_y and transform it into input that scaler expects
                previous_day_price = np.array([[testing_y[prediction_id - 1], 0]], dtype=object)
                previous_day_price = scaler.inverse_transform(previous_day_price)[0][0]
                # get today price from testing_y and transform it into input that scaler expects
                today_price = np.array([[testing_y[prediction_id], 0]], dtype=object)
                today_price = scaler.inverse_transform(today_price)[0][0]
                # check validity of prediction -> directional change should be preserved
                validity = 0
                real_change = today_price - previous_day_price
                predicted_change = denormalized_prediction - previous_day_price

                if (real_change >= 0) and (predicted_change >= 0):
                    validity = 1
                elif (real_change < 0) and (predicted_change < 0):
                    validity = 1

                test_stats.append({
                    'id': prediction_id,
                    'model': model_name,
                    'stock': loop_stock,
                    'date': trading_days[prediction_id + data_offset + 1],
                    'predicted_value': prediction,
                    'denormalized_prediction': denormalized_prediction,
                    'previous_day_price': previous_day_price,
                    'today_price': today_price,
                    'validity': validity
                })

            #  normalized values
            testing_y_1d = testing_y.reshape(-1)
            # calculate mse
            mse = mean_squared_error(testing_y_1d, predictions)
            # plot predictions
            plot_predictions(predictions, testing_y_1d, loop_stock, model_name + '.png', 'exp-4', 'Normalizovaná cena', len(predictions) + 1)
            # store model stats
            model_stats.append({
                'model': model_name,
                'test_mse': mse,
            })

            # store model
            if not model.stored():
                model.save(str(model.history['val_mse'][-1]))

# store all model stats with the model name into csv
with open('experiment-results/exp-4/experiment-4-model-stats.csv', 'a') as f:
    for stat in model_stats:
        f.write(f"{stat['model']},{stat['test_mse']}\n")

# test stats to df
test_stats_df = pd.DataFrame(test_stats)
test_stats_df = test_stats_df[test_stats_df['validity'] == 1].groupby(['model', 'validity']).size().reset_index()
test_stats_df = test_stats_df.rename(columns={0: 'success_count'}).sort_values(by=['success_count'], ascending=False)

# store all test stats with the model name into csv
with open('experiment-results/exp-4/experiment-4-predictions.csv', 'a') as f:
    for stat in test_stats:
        f.write(
            f"{stat['id']},{stat['model']},{stat['stock']},{stat['date']},{stat['predicted_value']},{stat['denormalized_prediction']},{stat['previous_day_price']},{stat['today_price']},{stat['validity']}\n")
