from numpy import array, corrcoef, inf, sqrt
from numpy.testing import assert_allclose, assert_almost_equal, assert_array_equal
from numpy.typing import NDArray
import os
from pytest import mark
from unittest import mock

from hestia_earth.utils.stats import (
    _calc_confidence_level,
    add_normal_distributions,
    avg_run_in_columnwise,
    avg_run_in_rowwise,
    calc_confidence_level_monte_carlo,
    calc_precision_monte_carlo,
    calc_required_iterations_monte_carlo,
    calc_z_critical,
    correlated_normal_2d,
    discrete_uniform_1d,
    discrete_uniform_2d,
    gen_seed,
    grouped_avg,
    lerp_normal_distributions,
    normal_1d,
    normal_2d,
    plus_minus_uncertainty_to_normal_1d,
    plus_minus_uncertainty_to_normal_2d,
    repeat_1d_array_as_columns,
    repeat_array_as_columns,
    repeat_array_as_rows,
    repeat_single,
    subtract_normal_distributions,
    triangular_1d,
    triangular_2d,
    truncated_normal_1d,
    truncated_normal_2d,
)

# confidence_level, n_sided, z_critical
CONFIDENCE_INTERVAL_PARAMS = [
    # 1 sided
    (0, 1, -inf),
    (50, 1, 0),
    (80, 1, 0.8416),
    (90, 1, 1.2816),
    (95, 1, 1.6449),
    (99, 1, 2.3263),
    (100, 1, inf),
    # 2 sided
    (0, 2, 0),
    (50, 2, 0.6745),
    (80, 2, 1.2816),
    (90, 2, 1.6449),
    (95, 2, 1.9600),
    (99, 2, 2.5758),
    (100, 2, inf),
]


@mark.parametrize(
    "confidence_level, n_sided, z_critical",
    CONFIDENCE_INTERVAL_PARAMS,
    ids=[f"z={z}, n={n}" for _, n, z in CONFIDENCE_INTERVAL_PARAMS],
)
def test_calc_confidence_level(confidence_level, n_sided, z_critical):
    result = _calc_confidence_level(z_critical, n_sided=n_sided)
    assert_almost_equal(result, confidence_level, decimal=2)


@mark.parametrize(
    "confidence_level, n_sided, z_critical",
    CONFIDENCE_INTERVAL_PARAMS,
    ids=[f"conf={conf}, n={n}" for conf, n, _ in CONFIDENCE_INTERVAL_PARAMS],
)
def test_calc_z_critical(confidence_level, n_sided, z_critical):
    result = calc_z_critical(confidence_level, n_sided=n_sided)
    assert_almost_equal(result, z_critical, decimal=4)


# confidence_level, n_iterations, precision, sd
MONTE_CARLO_PARAMS = [
    (95, 80767, 0.01, 1.45),
    (95, 1110, 0.01, 0.17),
    (99, 1917, 0.01, 0.17),
    (50, 102, 100.18, 1500),
]


@mark.parametrize(
    "confidence_level, n_iterations, precision, sd",
    MONTE_CARLO_PARAMS,
    ids=[f"n={n}, prec={prec}, sd={sd}" for _, n, prec, sd in MONTE_CARLO_PARAMS],
)
def test_calc_confidence_level_monte_carlo(
    confidence_level, n_iterations, precision, sd
):
    result = calc_confidence_level_monte_carlo(
        n_iterations,
        precision,
        sd,
    )
    assert_almost_equal(result, confidence_level, decimal=2)


@mark.parametrize(
    "confidence_level, n_iterations, precision, sd",
    MONTE_CARLO_PARAMS,
    ids=[
        f"conf={conf}, prec={prec}, sd={sd}" for conf, _, prec, sd in MONTE_CARLO_PARAMS
    ],
)
def test_calc_required_iterations_monte_carlo(
    confidence_level, n_iterations, precision, sd
):
    result = calc_required_iterations_monte_carlo(confidence_level, precision, sd)
    assert result == n_iterations


@mark.parametrize(
    "confidence_level, n_iterations, precision, sd",
    MONTE_CARLO_PARAMS,
    ids=[f"conf={conf}, n={n}, sd={sd}" for conf, n, _, sd in MONTE_CARLO_PARAMS],
)
def test_calc_precision_monte_carlo(confidence_level, n_iterations, precision, sd):
    result = calc_precision_monte_carlo(confidence_level, n_iterations, sd)
    assert_almost_equal(result, precision, decimal=2)


# mu_1, sigma_1, mu_2, sigma_2, rho, sum_mean, sum_sigma, diff_mean, diff_sigma
PARAMS_NORMAL_DIST = [
    # 2 standard normal distributions, perfectly negative correlation
    (0, 1, 0, 1, -1, 0, 0, 0, 2),
    # 2 standard normal distributions, negative correlation
    (0, 1, 0, 1, -0.5, 0, 1, 0, sqrt(3)),
    # 2 standard normal distributions, no correlation
    (0, 1, 0, 1, 0, 0, sqrt(2), 0, sqrt(2)),
    # 2 standard normal distributions, positive correlation
    (0, 1, 0, 1, 0.5, 0, sqrt(3), 0, 1),
    # 2 standard normal distributions, perfectly positive correlation
    (0, 1, 0, 1, 1, 0, 2, 0, 0),
    # different normal distributions, perfectly negative correlation
    (50000, 3000, 45000, 9000, -1, 95000, 6000, 5000, 12000),
    # different normal distributions, no correlation
    (50000, 3000, 45000, 9000, 0, 95000, sqrt(90000000), 5000, sqrt(90000000)),
    # different normal distributions, perfectly positive correlation
    (50000, 3000, 45000, 9000, 1, 95000, 12000, 5000, 6000),
]
IDS_ADD_NORMAL_DIST = [
    f"N({mu_1}, {sigma_1}^2) + N({mu_2}, {sigma_2}^2), rho: {rho}"
    for mu_1, sigma_1, mu_2, sigma_2, rho, *_ in PARAMS_NORMAL_DIST
]
IDS_SUBTRACT_DIST = [
    f"N({mu_1}, {sigma_1}^2) - N({mu_2}, {sigma_2}^2), rho: {rho}"
    for mu_1, sigma_1, mu_2, sigma_2, rho, *_ in PARAMS_NORMAL_DIST
]


@mark.parametrize(
    "mu_1, sigma_1, mu_2, sigma_2, rho, sum_mean, sum_sigma, _diff_mean, _diff_sigma",
    PARAMS_NORMAL_DIST,
    ids=IDS_ADD_NORMAL_DIST,
)
def test_add_normal_distributions(
    mu_1, sigma_1, mu_2, sigma_2, rho, sum_mean, sum_sigma, _diff_mean, _diff_sigma
):
    result = add_normal_distributions(mu_1, sigma_1, mu_2, sigma_2, rho)
    assert result == (sum_mean, sum_sigma)


@mark.parametrize(
    "mu_1, sigma_1, mu_2, sigma_2, rho, _sum_mean, _sum_sigma, diff_mean, diff_sigma",
    PARAMS_NORMAL_DIST,
    ids=IDS_SUBTRACT_DIST,
)
def test_subtract_normal_distributions(
    mu_1, sigma_1, mu_2, sigma_2, rho, _sum_mean, _sum_sigma, diff_mean, diff_sigma
):
    result = subtract_normal_distributions(mu_1, sigma_1, mu_2, sigma_2, rho)
    assert result == (diff_mean, diff_sigma)


# mu_1, sigma_1, mu_2, sigma_2, alpha, rho, Z_mean, Z_sigma
PARAMS_LERP_NORMAL_DIST = [
    # 2 standard normal distributions, perfectly negative correlation
    (0, 1, 0, 1, 0, -1, 0, 1),
    (0, 1, 0, 1, 0.5, -1, 0, 0),
    (0, 1, 0, 1, 1, -1, 0, 1),
    # 2 standard normal distributions, no correlation
    (0, 1, 0, 1, 0, 0, 0, 1),
    (0, 1, 0, 1, 0.5, 0, 0, sqrt(0.5)),
    (0, 1, 0, 1, 1, 0, 0, 1),
    # 2 standard normal distributions, perfectly positive correlation
    (0, 1, 0, 1, 0, 1, 0, 1),
    (0, 1, 0, 1, 0.5, 1, 0, 1),
    (0, 1, 0, 1, 1, 1, 0, 1),
    # different normal distributions, perfectly negative correlation
    (10000, 3000, 5000, 2500, -0.5, -1, 12500, 5750),
    (10000, 3000, 5000, 2500, 0, -1, 10000, 3000),
    (10000, 3000, 5000, 2500, 0.5, -1, 7500, 250),
    (10000, 3000, 5000, 2500, 1, -1, 5000, 2500),
    (10000, 3000, 5000, 2500, 1.5, -1, 2500, 5250),
    # different normal distributions, no correlation
    (10000, 3000, 5000, 2500, -0.5, 0, 12500, sqrt(21812500)),
    (10000, 3000, 5000, 2500, 0, 0, 10000, 3000),
    (10000, 3000, 5000, 2500, 0.5, 0, 7500, sqrt(3812500)),
    (10000, 3000, 5000, 2500, 1, 0, 5000, 2500),
    (10000, 3000, 5000, 2500, 1.5, 0, 2500, sqrt(16312500)),
    # different normal distributions, perfectly positive correlation
    (10000, 3000, 5000, 2500, -0.5, 1, 12500, 3250),
    (10000, 3000, 5000, 2500, 0, 1, 10000, 3000),
    (10000, 3000, 5000, 2500, 0.5, 1, 7500, 2750.0),
    (10000, 3000, 5000, 2500, 1, 1, 5000, 2500),
    (10000, 3000, 5000, 2500, 1.5, 1, 2500, 2250),
]
IDS_LERP_NORMAL_DIST = [
    f"N({mu_1}, {sigma_1}^2) - N({mu_2}, {sigma_2}^2), alpha: {alpha}, rho: {rho}"
    for mu_1, sigma_1, mu_2, sigma_2, alpha, rho, *_ in PARAMS_LERP_NORMAL_DIST
]


@mark.parametrize(
    "mu_1, sigma_1, mu_2, sigma_2, alpha, rho, Z_mean, Z_sigma",
    PARAMS_LERP_NORMAL_DIST,
    ids=IDS_LERP_NORMAL_DIST,
)
def test_lerp_normal_distributions(
    mu_1, sigma_1, mu_2, sigma_2, alpha, rho, Z_mean, Z_sigma
):
    result = lerp_normal_distributions(mu_1, sigma_1, mu_2, sigma_2, alpha, rho)
    assert result == (Z_mean, Z_sigma)


CYCLE = {
    "@id": "test-cycle",
    "name": "Product - Region, Country - 2026 - Description",
    "originalId": "original-cycle-id",
}

SITE = {
    "@id": "test-site",
    "name": "siteType - Organisation - Region, Country - Description",
    "originalId": "original-site-id",
}


@mark.parametrize(
    "node, args, expected",
    [
        (CYCLE, (), 2428136795),
        (CYCLE, ("arg1", "arg2"), 92189471),
        (SITE, (), 4270147077),
        (SITE, ("arg1", "arg2"), 3134241318),
    ],
    ids=["cycle", "cycle w/ args", "site", "site w/ args"],
)
def test_gen_seed(node, args, expected):
    result = gen_seed(node, *args)
    assert result == expected


@mark.parametrize(
    "node, args, MODELS_SEED_KEY, expected",
    [
        (CYCLE, (), "name", 583600750),
        (CYCLE, ("arg1", "arg2"), "name", 3970152803),
        (CYCLE, (), "originalId", 3994783498),
        (CYCLE, ("arg1", "arg2"), "originalId", 1040701350),
        (CYCLE, (), "name,@id", 4259822156),
        (CYCLE, ("arg1", "arg2"), "name,@id", 878902438),
        (
            CYCLE,
            (),
            "invalid key",
            3313293952,
        ),  # hash value of `"None"`, still works but won't be unique
        (SITE, (), "name", 539681188),
        (SITE, ("arg1", "arg2"), "name", 659108581),
        (SITE, (), "originalId", 3008679139),
        (SITE, ("arg1", "arg2"), "originalId", 446523220),
        (SITE, (), "name,@id", 2950540721),
        (SITE, ("arg1", "arg2"), "name,@id", 3266714212),
        (
            SITE,
            (),
            "invalid key",
            3313293952,
        ),  # hash value of `str(None)`, still works but won't be unique
    ],
    ids=[
        "cycle w/ name",
        "cycle w/ args & name",
        "cycle w/ originalId",
        "cycle w/ args & originalId",
        "cycle w/ name,@id",
        "cycle w/ args & name,@id",
        "cycle w/ missing key",
        "site w/ name",
        "site w/ args & name",
        "site w/ originalId",
        "site w/ args & originalId",
        "site w/ name,@id",
        "site w/ args & name,@id",
        "site w/ missing key",
    ],
)
def test_gen_seed_with_env_var_override(node, args, MODELS_SEED_KEY, expected):
    with mock.patch.dict(os.environ, {"MODELS_SEED_KEY": MODELS_SEED_KEY}):
        result = gen_seed(node, *args)

    assert result == expected


SEED = 0
N_ITERATIONS = 1000
SHAPE = (1000, 1000)


def assert_rows_identical(arr: NDArray):
    """
    Covert array to a set to remove repeated rows and check that number remaining rows is 1.
    """
    assert len(set(map(tuple, arr))) == 1


def assert_rows_unique(arr: NDArray):
    """
    Covert array to a set to remove repeated rows and check that number remaining rows is the same as the number of
    original rows.
    """
    assert len(set(map(tuple, arr))) == len(arr)


def assert_elements_between(arr: NDArray, min: float, max: float):
    assert ((min <= arr) & (arr <= max)).all()


PARAMS_REPEAT_SINGLE = [
    (3.14159, None, 3.14159),
    (3.14159, bool, True),
    (True, None, True),
    (True, float, 1),
]

IDS_REPEAT_SINGLE = [
    f"{type(value).__name__}{f' -> {dtype.__name__}' if dtype else ''}"
    for value, dtype, _ in PARAMS_REPEAT_SINGLE
]


@mark.parametrize(
    "value, dtype, expected_element",
    [
        (3.14159, None, 3.14159),
        (3.14159, bool, True),
        (True, None, True),
        (True, float, 1),
    ],
    ids=IDS_REPEAT_SINGLE,
)
def test_repeat_single(value, dtype, expected_element):
    SHAPE = (3, 3)
    EXPECTED = array(
        [
            [expected_element, expected_element, expected_element],
            [expected_element, expected_element, expected_element],
            [expected_element, expected_element, expected_element],
        ]
    )
    result = repeat_single(SHAPE, value, dtype=dtype)
    assert_array_equal(result, EXPECTED)


def test_repeat_array_as_columns():
    INPUT = array([[1, 2, 3], [4, 5, 6]])
    EXPECTED = array([[1, 2, 3, 1, 2, 3], [4, 5, 6, 4, 5, 6]])
    result = repeat_array_as_columns(2, INPUT)
    assert_array_equal(result, EXPECTED)


def test_repeat_array_as_rows():
    INPUT = array([[1, 2, 3], [4, 5, 6]])
    EXPECTED = array([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]])
    result = repeat_array_as_rows(2, INPUT)
    assert_array_equal(result, EXPECTED)


def test_repeat_1d_array_as_columns():
    INPUT = array([1, 2, 3])
    EXPECTED = array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    result = repeat_1d_array_as_columns(3, INPUT)
    assert_array_equal(result, EXPECTED)


def test_discrete_uniform_1d():
    MIN, MAX = -100, 100
    result = discrete_uniform_1d(SHAPE, MIN, MAX, seed=SEED)
    assert_rows_identical(result)
    assert_elements_between(result, MIN, MAX)
    assert result.shape == SHAPE


def test_discrete_uniform_2d():
    MIN, MAX = -100, 100
    result = discrete_uniform_2d(SHAPE, MIN, MAX, seed=SEED)
    assert_rows_unique(result)
    assert_elements_between(result, MIN, MAX)
    assert result.shape == SHAPE


def test_discrete_triangular_1d():
    LOW, HIGH = -100, 100
    MODE = -50
    result = triangular_1d(SHAPE, LOW, HIGH, MODE, seed=SEED)
    assert_rows_identical(result)
    assert_elements_between(result, LOW, HIGH)
    assert result.shape == SHAPE


def test_discrete_triangular_2d():
    LOW, HIGH = -100, 100
    MODE = 50
    result = triangular_2d(SHAPE, LOW, HIGH, MODE, seed=SEED)
    assert_rows_unique(result)
    assert_elements_between(result, LOW, HIGH)
    assert result.shape == SHAPE


def test_normal_1d():
    MEAN = 0
    SD = 50
    result = normal_1d(SHAPE, MEAN, SD, seed=SEED)
    assert_rows_identical(result)
    assert result.shape == SHAPE


def test_normal_2d():
    MEAN = 0
    SD = 50
    result = normal_2d(SHAPE, MEAN, SD, seed=SEED)
    assert_rows_unique(result)
    assert result.shape == SHAPE


def test_truncated_normal_1d():
    MEAN = 0
    SD = 50
    LOW, HIGH = -50, 50
    result = truncated_normal_1d(SHAPE, MEAN, SD, LOW, HIGH, seed=SEED)
    assert_rows_identical(result)
    assert_elements_between(result, LOW, HIGH)
    assert result.shape == SHAPE


def test_truncated_normal_2d():
    MEAN = 0
    SD = 50
    LOW, HIGH = -50, 50
    result = truncated_normal_2d(SHAPE, MEAN, SD, LOW, HIGH, seed=SEED)
    assert_rows_unique(result)
    assert_elements_between(result, LOW, HIGH)
    assert result.shape == SHAPE


def test_plus_minus_uncertainty_to_normal_1d():
    MEAN = 10
    UNCERTAINTY = 10
    CONFIDENCE_INTERVAL = 95
    result = plus_minus_uncertainty_to_normal_1d(
        SHAPE, MEAN, UNCERTAINTY, CONFIDENCE_INTERVAL
    )
    assert_rows_identical(result)
    assert result.shape == SHAPE


def test_plus_minus_uncertainty_to_normal_2d():
    MEAN = 10
    UNCERTAINTY = 10
    CONFIDENCE_INTERVAL = 95
    result = plus_minus_uncertainty_to_normal_2d(
        SHAPE, MEAN, UNCERTAINTY, CONFIDENCE_INTERVAL
    )
    assert_rows_unique(result)
    assert result.shape == SHAPE


def test_grouped_avg():
    INPUT = array(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15], [16, 17, 18]]
    )
    EXPECTED = array([[4, 5, 6], [13, 14, 15]])
    result = grouped_avg(INPUT, n=3)
    assert_array_equal(result, EXPECTED)


def test_avg_run_in_columnwise():
    INPUT = array(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15], [16, 17, 18]]
    )
    EXPECTED = array([[4, 5, 6], [10, 11, 12], [13, 14, 15], [16, 17, 18]])
    result = avg_run_in_columnwise(INPUT, n=3)
    assert_array_equal(result, EXPECTED)


def test_avg_run_in_rowwise():
    INPUT = array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
    EXPECTED = array([[2, 4, 5], [7, 9, 10], [12, 14, 15]])
    result = avg_run_in_rowwise(INPUT, n=3)
    assert_array_equal(result, EXPECTED)


# means, sds, correlation_matrix, tol_kwargs
PARAMS_CORRELATED_NORMAL_2D = [
    (
        array([0, 0, 0, 0, 0]),
        array([1, 1, 1, 1, 1]),
        array(
            [
                [1.0, 0.5, 0.25, 0.125, 0.0625],
                [0.5, 1.0, 0.5, 0.25, 0.125],
                [0.25, 0.5, 1.0, 0.5, 0.25],
                [0.125, 0.25, 0.5, 1.0, 0.5],
                [0.0625, 0.125, 0.25, 0.5, 1.0],
            ]
        ),
        {"atol": 0.1},
    ),
    (
        array([40000, 42000, 43000, 44333.333]),
        array([4000, 4200, 4300, 4433.333]),
        array(
            [
                [1.0, 0.965867, 0.932987, 0.901227],
                [0.965867, 1.0, 0.965959, 0.933076],
                [0.932987, 0.965959, 1.0, 0.965959],
                [0.901227, 0.933076, 0.965959, 1.0],
            ]
        ),
        {"rtol": 0.05},
    ),
    (
        array([40000, 42000, 43000, 44333.333]),
        array([4000, 4200, 4300, 4433.333]),
        array([[1.0, 0, 0, 0], [0, 1.0, 0, 0], [0, 0, 1.0, 0], [0, 0, 0, 1.0]]),
        {"rtol": 0.05},
    ),
]
IDS_CORRELATED_NORMAL_2D = [
    "standard normal distributions",
    "custom normal distributions",
    "no correlation",
]


@mark.parametrize(
    "means, sds, correlation_matrix, tol_kwargs",
    PARAMS_CORRELATED_NORMAL_2D,
    ids=IDS_CORRELATED_NORMAL_2D,
)
def test_correlated_normal_2d_standard_normal_dist(
    means: NDArray, sds: NDArray, correlation_matrix: NDArray, tol_kwargs: dict
):
    result = correlated_normal_2d(
        N_ITERATIONS, means, sds, correlation_matrix, seed=SEED
    )
    empirical_correlation_matrix = corrcoef(result)

    assert_allclose(result.mean(axis=1), means, **tol_kwargs)
    assert_allclose(result.std(axis=1), sds, **tol_kwargs)
    assert_allclose(empirical_correlation_matrix, correlation_matrix, atol=0.1)
