import numpy as np
from fractions import Fraction
from pamle.jackknife import (vandermonde_weights_exact, vandermonde_weights, leading_bias_coefficient,
                             amplification_constant, vandermonde_condition_number, combine_levels, DEFAULT_SPLITS)


def test_moment_constraints_exact():
    for s in [(1, 2), (1, 2, 4), (1, 2, 4, 8), (1, 2, 3), (1, 3, 9), (1, 2, 4, 8, 16)]:
        w = vandermonde_weights_exact(s)
        assert sum(w) == 1
        for a in range(1, len(s)):
            assert sum(wj * Fraction(sj) ** a for wj, sj in zip(w, s)) == 0


def test_known_weights():
    assert vandermonde_weights_exact((1, 2)) == [2, -1]
    assert vandermonde_weights_exact((1, 2, 4)) == [Fraction(8, 3), -2, Fraction(1, 3)]
    assert np.allclose(vandermonde_weights((1, 2, 4)), [8 / 3, -2, 1 / 3])


def test_float_moment_constraints():
    for r in range(1, 6):
        s = DEFAULT_SPLITS[r]; w = vandermonde_weights(s); sv = np.array(s, float)
        assert abs(w.sum() - 1) < 1e-12
        for a in range(1, r):
            assert abs(np.sum(w * sv ** a)) < 1e-9 * np.sum(np.abs(w) * sv ** a)


def test_leading_coefficients_and_amplification():
    assert leading_bias_coefficient((1, 2)) == -2
    assert leading_bias_coefficient((1, 2, 4)) == 8
    assert leading_bias_coefficient((1, 2, 4, 8)) == -64
    assert amplification_constant((1, 2)) == 4 and amplification_constant((1, 2, 4)) == 8
    assert vandermonde_condition_number((1, 2, 4, 8)) > vandermonde_condition_number((1, 2))


def test_combine_levels_preserves_mean():
    rng = np.random.default_rng(0)
    K, p = 5, 3
    common = rng.standard_normal((K, p))
    levels = np.stack([common, common, common])          # identical level means -> J equals them (sum w = 1)
    J = combine_levels(levels, vandermonde_weights((1, 2, 4)))
    assert np.allclose(J, common)