import numpy as np
import matplotlib.pyplot as plt
from numpy.typing import NDArray
import argparse


def twodim_array_to_pgfplots_table_string(a: NDArray):
    return (
        " \\\\\n".join([" ".join([str(vali) for vali in val]) for val in a]) + "\\\\\n"
    )


def main():
    # Parse command line arguments

    parser = argparse.ArgumentParser()
    parser.add_argument("--correlation", "-c", type=np.float32, required=True)
    parser.add_argument("-N", type=np.int32, required=True)
    parser.add_argument("--plot", "-p", action="store_true")

    args = parser.parse_args()

    # Generate & plot data

    means = np.array([0, 0])
    cov = np.array([[1, args.correlation], [args.correlation, 1]])

    x = np.random.multivariate_normal(means, cov, size=args.N)

    print(twodim_array_to_pgfplots_table_string(x))

    if args.plot:
        plt.scatter(x[:, 0], x[:, 1])
        plt.show()


if __name__ == "__main__":
    main()