import numpy as np


def generate_ground_truth_data(length):
    """
    Generate and save test data using the parameters from Ding et al. 2006
    section 5.1 Example 1 (p. 18).

    Parameters
    ----------
    length : int
        Length of generated time series data

    References
    ----------
    Ding, M., Chen, Y., & Bressler, S. L. (2006). Granger Causality: Basic
    Theory and Application to Neuroscience. arXiv preprint q-bio/0608035.

    """
    np.random.seed(123)
    order = 2
    time_series = np.zeros((2, length + order))

    weights_1 = np.array([[0.9, 0], [0.16, 0.8]])
    weights_2 = np.array([[-0.5, 0], [-0.2, -0.5]])

    weights = np.stack((weights_1, weights_2))

    noise_covariance = np.array([[1., 0.4], [0.4, 0.7]])

    for i in range(length):
        for lag in range(order):
            time_series[:, i + order] += np.dot(weights[lag],
                                                time_series[:, i + 1 - lag])
        rnd_var = np.random.multivariate_normal([0, 0],
                                                noise_covariance)
        time_series[:, i + order] += rnd_var

    time_series = time_series[:, 2:]

    # Save the generated data
    data_folder = '../data/'
    if length <= 30000:
        time_series_filename = 'time_series_small.npy'
    else:
        time_series_filename = 'time_series.npy'
    np.save(data_folder + time_series_filename, time_series)
    np.save(data_folder + 'weights.npy', weights)
    np.save(data_folder + 'noise_covariance.npy', noise_covariance)


if __name__ == '__main__':
    generate_ground_truth_data(30000)
    generate_ground_truth_data(300000)