"""
Helper functions that prepare data for training and testing and can create sample plots.
"""
from src.constants import *


def prepare_dataset_single_model(dataset, start, end):
    """
    Creates a dataset for a single model.
    :param dataset:
    :param start:
    :param end:
    :return:
    """
    inputs = []
    outputs = []

    # check if Volume column is present and drop it
    if isinstance(dataset, pd.DataFrame) and 'Volume' in dataset.columns:
        dataset = dataset.drop(columns=['Volume'])

    # if end is not specified -> take rest of dataset
    if end is None:
        end = len(dataset)

    # create windows of BATCH_SIZE days
    for priceIndex in range(start + BATCH_SIZE, end):
        indices = range(priceIndex - BATCH_SIZE, priceIndex)
        inputs.append(dataset.iloc[indices])
        outputs.append(dataset.iloc[priceIndex])

    # convert to numpy arrays
    return np.array(inputs), np.array(outputs)


def prepare_dataset_multi_model(dataset, start, end, packages=None):
    """
    Creates a dataset for a multi model.
    :param dataset:
    :param start:
    :param end:
    :param packages:
    :return:
    """
    inputs = []
    outputs = []

    # check if date is present and drop it
    if isinstance(dataset, pd.DataFrame) and 'Date' in dataset.columns:
        dataset = dataset.drop(columns=['Date'])

    # creating rolling window of WINDOW_SIZE
    window = dataset['Adj Close'].rolling(window=WINDOW_SIZE)

    # determine mode -> no packages -> use best features (std, min, max) + volume
    #                               -> use provided features as model input
    if packages is None:
        dataset['std'] = window.std()
        dataset['max'] = window.max()
        dataset['min'] = window.min()
    else:
        # add min max std
        if 'MMS' in packages:
            dataset['min'] = window.min()
            dataset['max'] = window.max()
            dataset['std'] = window.std()

        # check if volume should be present
        if 'VOL' not in packages:
            dataset = dataset.drop(columns=['Volume'])

        # add moving average
        if 'MA' in packages:
            dataset['ma'] = window.mean()

        # add exponential moving average
        if 'XMA' in packages:
            dataset['ema'] = dataset['Adj Close'].ewm(span=WINDOW_SIZE, adjust=False).mean()

    # if end is not specified -> take rest of dataset
    if end is None:
        end = len(dataset)

    # create windows of BATCH_SIZE days
    for priceIndex in range(start + BATCH_SIZE, end):
        start_index = priceIndex - BATCH_SIZE
        indices = range(start_index, priceIndex)
        inputs.append(dataset.iloc[indices].to_numpy())
        outputs.append(dataset.iloc[priceIndex]['Adj Close'])

    # first usable window is WINDOW_SIZE - 1, but that would yield 501 predictions, to keep consistent with single model
    # only use data from WINDOW_SIZE on
    return np.array(inputs[WINDOW_SIZE:]), np.array(outputs[WINDOW_SIZE:])


def create_datasets(train_x, train_y, validation_x, validation_y, dataset_len):
    """
    Creates training and validation datasets from provided data.
    :param train_x:
    :param train_y:
    :param validation_x:
    :param validation_y:
    :param dataset_len:
    :return:
    """
    # create training mini batches
    batch_size = 30
    buffer_size = dataset_len

    # prepare datasets for TF (shuffle, batches)
    training_dataset = Dataset \
        .from_tensor_slices((train_x, train_y)) \
        .shuffle(buffer_size) \
        .batch(batch_size, drop_remainder=True) \
        .repeat()

    validation_dataset = Dataset.from_tensor_slices(
        (validation_x, validation_y)).batch(BATCH_SIZE).repeat()

    return training_dataset, validation_dataset


def normalize_dataset(training_data, validation_data, testing_data, expected_columns=['Adj Close']):
    """
    Normalizes data using MaxAbsScaler.
    :param training_data:
    :param validation_data:
    :param testing_data:
    :param expected_columns:
    :return:
    """
    # init scaler
    scaler = MaxAbsScaler()
    # transform data
    if len(expected_columns) == 1:
        # reshape for experiment 1 and 2
        training_data = scaler.fit_transform(training_data.values.reshape(-1, 1))
        validation_data = scaler.transform(validation_data.values.reshape(-1, 1))
        testing_data = scaler.transform(testing_data.values.reshape(-1, 1))
    else:
        # transform for experiment 3, 4, 5
        training_data[expected_columns] = scaler.fit_transform(training_data[expected_columns])
        validation_data[expected_columns] = scaler.transform(validation_data[expected_columns])
        testing_data[expected_columns] = scaler.transform(testing_data[expected_columns])

    # join training and validation data
    training_data = np.concatenate((training_data, validation_data), axis=0)
    # create dataframes
    training_data = pd.DataFrame(training_data, columns=expected_columns)
    testing_data = pd.DataFrame(testing_data, columns=expected_columns)
    return scaler, training_data, testing_data


def plot_predictions(predicted_values, real_values, stock_name='BRK-A', save_name='test.png', save_dir='exp-3', ylabel='Cena ($)', x_vals=201, xlabel='Predikované dni'):
    """
    Plots predicted values against real values.
    :param predicted_values:
    :param real_values:
    :param stock_name:
    :param save_name:
    :param save_dir:
    :param ylabel:
    :param x_vals:
    :param xlabel:
    :return:
    """
    plt.figure()
    plt.plot(np.arange(1, x_vals), real_values, label='reálna krivka (%s)' % stock_name)
    plt.plot(np.arange(1, x_vals), predicted_values, label='krivka predikcií (%s)' % stock_name)
    plt.xlabel(xlabel, fontsize=14, labelpad=8)
    plt.ylabel(ylabel, fontsize=14, labelpad=8)
    plt.legend()
    plt.tight_layout(pad=0.5)
    plt.savefig('experiment-results/{}/{}'.format(save_dir, save_name))
    plt.show()
