import numpy as np
from pamle.likelihoods import Exponential, Logistic, gather_blocks
from pamle.algorithms import (local_split_fits, worker_corrected_estimates, pilot_estimator, global_newton,
                              aee_lin_xi, csl, analytic_bias_correction, estimate_all, quantize_relative)


def _exp_design(seed=0, N=2 ** 16, K=2 ** 8):
    rng = np.random.default_rng(seed); ex = Exponential()
    d = ex.simulate(rng, N, 1.0)
    return ex, d, gather_blocks(d, np.arange(N).reshape(K, N // K)), N, K


def test_newton_identity_on_data():
    ex, d, db, N, K = _exp_design()
    th_c = 1 / d['x'].mean()
    for x in (0.7, 1.0, 1.4):
        path, diag = global_newton(ex, db, np.array([x]), q=1)
        assert abs(path[1][0] - (2 * x - x * x * d['x'].mean())) < 1e-12
        assert abs((path[1][0] - th_c) + (x - th_c) ** 2 / th_c) < 1e-12


def test_pilot_bias_orders_exponential():
    """Average over many samples: raw bias ~ theta0/m, jackknife bias ~ -2 theta0/m^2 (exact formulas)."""
    rng = np.random.default_rng(1); ex = Exponential(); m = 32; K = 2000; reps = 200
    raw, j2, j3 = [], [], []
    for _ in range(reps):
        d = ex.simulate(rng, m * K, 1.0); db = gather_blocks(d, np.arange(m * K).reshape(K, m))
        fits = local_split_fits(ex, db, (1, 2, 4))
        raw.append(pilot_estimator(fits[1][0][:, 0]))
        j2.append(pilot_estimator(worker_corrected_estimates(fits, (1, 2))[0]))
        j3.append(pilot_estimator(worker_corrected_estimates(fits, (1, 2, 4))[0]))
    raw, j2, j3 = np.array(raw).ravel(), np.array(j2).ravel(), np.array(j3).ravel()
    se = 1 / np.sqrt(m * K * reps)
    assert abs(raw.mean() - m / (m - 1)) < 4 * se
    assert abs(j2.mean() - (1 - 2 / ((m - 1) * (m - 2)))) < 4 * se
    assert abs(j3.mean() - (1 + 8 / ((m - 1) * (m - 2) * (m - 4)))) < 4 * se


def test_estimate_all_agrees_with_central_when_predicted():
    rng = np.random.default_rng(2); m = Logistic(3); N = 2 ** 15; K = 64
    d = m.simulate(rng, N, [0.5, 1.0, -1.0]); db = gather_blocks(d, np.arange(N).reshape(K, N // K))
    out, diag = estimate_all(m, db, splits_max=3, q_max=2)
    c = out['central']
    gap = {k: np.sqrt(N) * np.linalg.norm(v - c) for k, v in out.items()}
    assert gap['r1_q2'] < gap['r1_q1'] < gap['r1_q0']
    assert gap['r2_q1'] < gap['r1_q1']
    assert gap['r2_q2'] < 1e-6 and gap['r3_q2'] < 1e-5
    assert gap['csl2'] < gap['csl1']
    assert diag['central_status'] == 0


def test_aee_and_csl_basic():
    ex, d, db, N, K = _exp_design(seed=3)
    fits = local_split_fits(ex, db, (1,))
    th = fits[1][0][:, 0]; H = fits[1][2][:, 0]
    a = aee_lin_xi(th, H)
    assert np.isfinite(a).all()
    path = csl(ex, db, th.mean(axis=0), rounds=1)
    # exponential model: L_0 and L_N differ by a linear term, so the full CSL surrogate maximizer IS the centralized MLE
    assert abs(path[1][0] - 1 / d['x'].mean()) < 1e-10
    # the quadratic (one-step) CSL variant with machine 0's Hessian equals the exact Newton step here (data-free Hessian)
    single, _ = global_newton(ex, db, th.mean(axis=0), q=1, hessian_mode="single")
    exact, _ = global_newton(ex, db, th.mean(axis=0), q=1)
    assert abs(single[1][0] - exact[1][0]) < 1e-12


def test_analytic_bias_coefficient_exponential():
    ex, d, db, N, K = _exp_design(seed=4)
    th = local_split_fits(ex, db, (1,))[1][0][:, 0]
    thc, b = analytic_bias_correction(ex, db, th)
    assert abs(b.mean() - 1.0) < 0.05          # b_1 = theta0 = 1


def test_quantize_relative():
    A = np.array([[1.234567, -0.001234], [3.0, 1e-9]])
    Q = quantize_relative(A, 8)
    assert np.all(np.abs(Q - A) <= np.abs(A) * 2.0 ** -8 + 1e-300)
    assert np.allclose(quantize_relative(A, 52), A)