"""Split-granularity extrapolation weights (generalized jackknife / Richardson extrapolation).

For split multipliers s_1 < ... < s_r (s_j sub-blocks of size m/s_j per worker; s_1 = 1 is the full block),
the weights solve the Vandermonde system
    sum_j w_j = 1,   sum_j w_j s_j^a = 0,  a = 1, ..., r-1,
so that a bias expansion b_1/n + ... + b_{r-1}/n^{r-1} evaluated at n = m/s_j cancels through order m^{-(r-1)}.
"""
from __future__ import annotations
from fractions import Fraction
import numpy as np

DEFAULT_SPLITS = {1: (1,), 2: (1, 2), 3: (1, 2, 4), 4: (1, 2, 4, 8), 5: (1, 2, 4, 8, 16)}


def vandermonde_weights_exact(splits):
    """Exact rational weights (list of Fraction) for the split multipliers `splits`."""
    s = [Fraction(int(x)) for x in splits]
    r = len(s)
    # Solve by Gaussian elimination on exact rationals.
    A = [[sj ** a for sj in s] for a in range(r)]
    b = [Fraction(1)] + [Fraction(0)] * (r - 1)
    M = [row[:] + [bi] for row, bi in zip(A, b)]
    for col in range(r):
        piv = next(i for i in range(col, r) if M[i][col] != 0)
        M[col], M[piv] = M[piv], M[col]
        pv = M[col][col]
        M[col] = [x / pv for x in M[col]]
        for i in range(r):
            if i != col and M[i][col] != 0:
                f = M[i][col]
                M[i] = [x - f * y for x, y in zip(M[i], M[col])]
    return [M[i][r] for i in range(r)]


def vandermonde_weights(splits):
    """Float64 weights."""
    return np.array([float(w) for w in vandermonde_weights_exact(splits)])


def leading_bias_coefficient(splits):
    """beta_r = sum_j w_j s_j^r: multiplies b_r m^{-r} in the residual bias of J_r (exact rational)."""
    w = vandermonde_weights_exact(splits); r = len(splits)
    return sum(wj * Fraction(int(sj)) ** r for wj, sj in zip(w, splits))


def amplification_constant(splits):
    """A_r = sum_j |w_j| s_j: multiplies sqrt(C)/m in the L2 bound of the corrected remainder (Lemma 3)."""
    w = vandermonde_weights_exact(splits)
    return sum(abs(wj) * Fraction(int(sj)) for wj, sj in zip(w, splits))


def weight_l1(splits):
    return sum(abs(w) for w in vandermonde_weights_exact(splits))


def vandermonde_condition_number(splits):
    s = np.asarray(splits, dtype=float); r = len(s)
    A = np.vstack([s ** a for a in range(r)])
    return float(np.linalg.cond(A))


def combine_levels(level_means, weights):
    """J = sum_j w_j T_{s_j}.  level_means: array (r, ..., p) of per-level averages; weights: (r,)."""
    w = np.asarray(weights, dtype=float)
    return np.tensordot(w, np.asarray(level_means), axes=(0, 0))


def level_averages(estimates_by_level):
    """estimates_by_level[j]: array (K, s_j, p) of sub-block estimates at level j -> (r, K, p) level means."""
    return np.stack([np.mean(E, axis=1) for E in estimates_by_level], axis=0)