#!/usr/bin/env python
from one_pass_stats import ParallelCovariance, ParallelDescriptiveStats


def _op_stats_to_str(stats: ParallelDescriptiveStats) -> str:
    return f"Mean {stats.mean} | Count {stats.count} | Stddev {stats.stddev()} | Min {stats.min} | Max {stats.max}"


def _make_stats(dist: list[float]) -> ParallelDescriptiveStats:
    stats = ParallelDescriptiveStats()
    for value in dist:
        stats.add_value(value)
    return stats


def _make_covar(args: tuple[list[float], list[float]]) -> ParallelCovariance:
    dist_x, dist_y = args
    covar = ParallelCovariance()
    for x, y in zip(dist_x, dist_y):
        covar.add_pair(x, y)
    print(f"make cov: {covar.covariance()} pears: {covar.pearson()}")
    return covar


def _naive_covariance(data1: list[float], data2: list[float], sample: bool = False) -> float:
    n = len(data1)
    sum12 = sum(data1[i] * data2[i] for i in range(n))
    sum1 = sum(data1)
    sum2 = sum(data2)
    return (sum12 - sum1 * sum2 / n) / (n - (1 if sample else 0))


def test_single_set(norm_dist: list[float]) -> ParallelDescriptiveStats:
    import numpy as np

    print("Calculating descriptive statistics using the one-pass algorithm...")
    onepass_stats = _make_stats(norm_dist)

    print("The following OnePass and numpy-calculated means should match: ")
    print(f"one-pass: {onepass_stats.mean}")
    print(f"numpy:    {np.mean(norm_dist)}")

    print("The following OnePass and numpy-calculated stddevs should match: ")
    print(f"one-pass: {onepass_stats.stddev()}")
    print(f"numpy:    {np.std(norm_dist, ddof=1)}")

    print(f"all one-pass stats: {_op_stats_to_str(onepass_stats)}")
    return onepass_stats


def test_merge(original_dist: list[float], parallel_nb: int) -> None:
    import numpy as np
    from multiprocessing import Pool

    dists = np.split(original_dist, parallel_nb)
    with Pool(processes=parallel_nb) as pool:
        stats_list = pool.map(_make_stats, dists)

    merged_stats = ParallelDescriptiveStats()
    for stats in stats_list:
        merged_stats = merged_stats.merge(stats)

    print(f"split ({parallel_nb}) & merged: {_op_stats_to_str(merged_stats)}")


def test_covar(dist_x: list[float], dist_y: list[float]) -> None:
    import numpy as np
    from multiprocessing import Pool

    print(f"distX = {dist_x}")
    print(f"distY = {dist_y}")

    covar = ParallelCovariance()
    for x, y in zip(dist_x, dist_y):
        covar.add_pair(x, y)
    print(f"OnePass Covariance: {covar.covariance()}; Pearson: {covar.pearson()}")

    ox = _make_stats(dist_x)
    oy = _make_stats(dist_y)
    print(f"correct stats: {_op_stats_to_str(ox)} {_op_stats_to_str(oy)}")

    naive_cov = _naive_covariance(dist_x, dist_y)
    print(f"Naive, two-pass Covariance: {naive_cov}; Pearson: {naive_cov / (ox.stddev() * oy.stddev())}")

    split_xs = np.split(np.array(dist_x), 4)
    split_ys = np.split(np.array(dist_y), 4)

    print("Testing covariance merge algorithm...")
    with Pool(processes=4) as pool:
        covars_list = pool.map(_make_covar, zip(split_xs, split_ys))

    merged_covar = ParallelCovariance()
    for covar in covars_list:
        merged_covar = merged_covar.merge(covar)

    print(merged_covar.covariance())
    print(merged_covar.pearson())


if __name__ == "__main__":
    import sys
    import time

    import numpy as np

    dist_size = 1000000
    expected_mean = 42.3
    expected_sigma = 67.3
    parallel_nb = 10

    if len(sys.argv) > 1:
        dist_size = int(sys.argv[1])
    if len(sys.argv) > 2:
        expected_mean = float(sys.argv[2])
    if len(sys.argv) > 3:
        expected_sigma = float(sys.argv[3])
    if len(sys.argv) > 4:
        parallel_nb = int(sys.argv[4])

    print(
        f"Testing normal distribution of size {dist_size} with "
        f"expected mean {expected_mean} and expected stddev {expected_sigma}."
    )
    norm_dist = np.random.normal(expected_mean, expected_sigma, dist_size)

    t0 = time.time()
    test_single_set(norm_dist)
    single = time.time() - t0
    print(f"Single threaded test run in {single} seconds")

    print("Now testing a map-reduce version using the same algorithm, but parallelized for speedup...")
    t0 = time.time()
    test_merge(norm_dist, parallel_nb)
    multi = time.time() - t0
    print(f"Multi-threaded test run in {multi} seconds, for a speedup of {single / multi}")

    x = [1, 3, 5, 7, 9, 11, 13, 15]
    y = [4, 17, 24, 33, 77, 119, 200, 270]

    test_covar(x, y)

    blah = ParallelCovariance()
    for i in range(len(x)):
        blah.add_pair(x[i], y[i])
    print(blah.covariance())

    blah1 = _make_covar((x[0:4], y[0:4]))
    blah2 = _make_covar((x[4:8], y[4:8]))
    blah = blah1.merge(blah2)
    print(blah.covariance(sample=True))
    print(_naive_covariance(x, y, sample=True))
