"""Tests for the IsoFLOPs / scaling-law workbench.

    pytest test_lab.py -v                       # your lab.py
    LAB_MODULE=solution pytest test_lab.py -v   # the reference (must be green)
"""

import importlib
import math
import os

import pytest

lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))

TRUTH = (1.69, 406.4, 0.34, 410.7, 0.28)
BUDGETS = [1e19, 1e20, 1e21, 1e22]
SIZES = [[N * m for N in (1e7, 3e7, 1e8, 3e8, 1e9)] for m in (1, 3, 10, 30)]


def _ladder(**kw):
    return lab.synthetic_ladder(BUDGETS, SIZES, **kw)


def _points(ladder):
    return [(N, D, L) for _, N, D, L in ladder]


# ======================================================================================
# 1. Numerical spine
# ======================================================================================

def test_fit_line_recovers_exact_line():
    xs = [1.0, 2.0, 3.0, 4.0]
    ys = [3.0 * x + 7.0 for x in xs]
    slope, intercept = lab.fit_line(xs, ys)
    assert slope == pytest.approx(3.0)
    assert intercept == pytest.approx(7.0)


def test_fit_line_rejects_degenerate_input():
    with pytest.raises(ValueError):
        lab.fit_line([1.0], [1.0])                     # too few points
    with pytest.raises(ValueError):
        lab.fit_line([2.0, 2.0, 2.0], [1.0, 2.0, 3.0])  # zero x-variance
    with pytest.raises(ValueError):
        lab.fit_line([1.0, 2.0], [1.0])                # mismatched lengths


def test_fit_power_law_recovers_exponent():
    """A power law is a straight line in log-log. The slope IS the exponent."""
    xs = [1e18, 1e19, 1e20, 1e21, 1e22]
    ys = [0.6 * x ** 0.5 for x in xs]
    c, k = lab.fit_power_law(xs, ys)
    assert k == pytest.approx(0.5, abs=1e-9)
    assert c == pytest.approx(0.6, rel=1e-6)


def test_fit_power_law_rejects_nonpositive_data():
    with pytest.raises(ValueError):
        lab.fit_power_law([1.0, 2.0, 3.0], [1.0, -2.0, 3.0])
    with pytest.raises(ValueError):
        lab.fit_power_law([0.0, 2.0, 3.0], [1.0, 2.0, 3.0])


def test_fit_quadratic_recovers_exact_parabola():
    xs = [-2.0, -1.0, 0.0, 1.0, 2.0]
    ys = [3.0 * x * x - 4.0 * x + 5.0 for x in xs]
    a, b, c = lab.fit_quadratic(xs, ys)
    assert a == pytest.approx(3.0, abs=1e-6)
    assert b == pytest.approx(-4.0, abs=1e-6)
    assert c == pytest.approx(5.0, abs=1e-6)


def test_fit_quadratic_rejects_too_few_points():
    with pytest.raises(ValueError):
        lab.fit_quadratic([1.0, 2.0], [1.0, 2.0])


def test_fit_quadratic_rejects_collinear_design():
    """BOUNDARY: identical x values make the normal equations singular. Must raise a
    ValueError, not divide by zero."""
    with pytest.raises(ValueError):
        lab.fit_quadratic([1.0, 1.0, 1.0], [1.0, 2.0, 3.0])


def test_parabola_vertex_finds_the_minimum():
    # y = 2x^2 - 8x + 1 has its minimum at x = 2.
    assert lab.parabola_vertex(2.0, -8.0) == pytest.approx(2.0)


def test_parabola_vertex_rejects_downward_parabola():
    """BOUNDARY: a downward-opening parabola has a MAXIMUM, not a minimum. If your
    IsoFLOPs fit produces one, the sweep is wrong — fail loudly."""
    with pytest.raises(ValueError):
        lab.parabola_vertex(-1.0, 4.0)
    with pytest.raises(ValueError):
        lab.parabola_vertex(0.0, 4.0)


# ======================================================================================
# 2. The loss model
# ======================================================================================

def test_parametric_loss_is_monotone_decreasing_in_both_axes():
    base = lab.parametric_loss(1e9, 1e11, *TRUTH)
    assert lab.parametric_loss(1e10, 1e11, *TRUTH) < base     # more params
    assert lab.parametric_loss(1e9, 1e12, *TRUTH) < base      # more tokens


def test_parametric_loss_approaches_the_irreducible_floor():
    """As N and D go to infinity, loss -> E and never below it."""
    huge = lab.parametric_loss(1e30, 1e30, *TRUTH)
    assert huge > TRUTH[0]
    assert huge == pytest.approx(TRUTH[0], abs=1e-3)


@pytest.mark.parametrize("n,d", [(0, 1e11), (1e9, 0), (-1e9, 1e11)])
def test_parametric_loss_rejects_nonpositive(n, d):
    with pytest.raises(ValueError):
        lab.parametric_loss(n, d, *TRUTH)


def test_loss_decomposition_fractions_sum_to_one():
    d = lab.loss_decomposition(70e9, 1.4e12)
    total = d["irreducible_frac"] + d["capacity_frac"] + d["data_frac"]
    assert total == pytest.approx(1.0)
    assert d["irreducible"] + d["capacity"] + d["data"] == pytest.approx(d["total"])


def test_loss_decomposition_shows_most_loss_is_irreducible():
    """The fact that reframes the whole industry: ~87% of the number is a floor."""
    d = lab.loss_decomposition(70e9, 1.4e12)
    assert d["irreducible_frac"] > 0.8


# ======================================================================================
# 3. The measurement bias — the Kaplan lesson
# ======================================================================================

def test_schedule_bias_is_zero_at_the_end_of_the_schedule():
    """BOUNDARY: reading the loss exactly at the scheduled horizon is unbiased."""
    assert lab.schedule_bias(1e11, 1e11, 0.15) == pytest.approx(0.0)


def test_schedule_bias_grows_as_you_read_earlier():
    """The further you are from the end of the schedule, the more the undecayed
    learning rate inflates the measured loss."""
    b_late = lab.schedule_bias(9e10, 1e11, 0.15)
    b_early = lab.schedule_bias(1e10, 1e11, 0.15)
    assert b_early > b_late > 0


def test_schedule_bias_never_negative():
    assert lab.schedule_bias(2e11, 1e11, 0.15) == pytest.approx(0.0)


def test_schedule_bias_rejects_bad_input():
    with pytest.raises(ValueError):
        lab.schedule_bias(1e11, 1e11, -0.1)
    with pytest.raises(ValueError):
        lab.schedule_bias(1e11, 0, 0.15)


def test_uniform_offset_is_absorbed_into_E_and_changes_nothing():
    """THE CONTROL for the money test below.

    A CONSTANT shift in every measured loss is absorbed into the fitted E. It cannot
    change the exponents, and therefore cannot change the recommendation. This is
    exactly why a naive 'the losses were all a bit high' story does NOT explain the
    Kaplan/Chinchilla discrepancy — the bias has to be non-uniform.
    """
    clean = _ladder(bias_mode="none")
    # Keep the ORIGINAL C values: recomputing 6*N*D in floating point would split a
    # budget into two groups and is itself a nice illustration of why you carry the
    # budget label around rather than re-deriving it.
    shifted = [(C, N, D, L + 0.05) for C, N, D, L in clean]

    exp_clean = lab.fit_scaling_exponents(lab.isoflops_ladder(clean))
    exp_shift = lab.fit_scaling_exponents(lab.isoflops_ladder(shifted))
    assert exp_shift["n_exponent"] == pytest.approx(exp_clean["n_exponent"], abs=1e-6)


def test_nonuniform_bias_tilts_the_fitted_exponent():
    """THE MONEY TEST — Kaplan -> Chinchilla, reproduced from scratch.

    A schedule-mismatch bias hits different ladder points by different amounts. That
    NON-UNIFORMITY tilts the fitted exponent and therefore moves the flagship
    recommendation — even though the underlying ground truth is identical.
    """
    clean = _ladder(bias_mode="none")
    biased = _ladder(lr_decay_bias=0.15, bias_mode="fixed_schedule")

    exp_clean = lab.fit_scaling_exponents(lab.isoflops_ladder(clean))
    exp_biased = lab.fit_scaling_exponents(lab.isoflops_ladder(biased))

    assert abs(exp_biased["n_exponent"] - exp_clean["n_exponent"]) > 0.02, (
        "a non-uniform measurement bias MUST move the fitted exponent")


def test_synthetic_ladder_is_deterministic_under_a_seed():
    """DETERMINISM: same seed -> byte-identical ladder."""
    a = _ladder(noise_std=0.01, seed=42)
    b = _ladder(noise_std=0.01, seed=42)
    c = _ladder(noise_std=0.01, seed=43)
    assert a == b
    assert a != c


def test_synthetic_ladder_validates():
    with pytest.raises(ValueError):
        _ladder(noise_std=-1)
    with pytest.raises(ValueError):
        _ladder(bias_mode="wishful_thinking")
    with pytest.raises(ValueError):
        lab.synthetic_ladder([1e19, 1e20], [[1e8]])       # mismatched lengths


def test_synthetic_ladder_respects_the_budget_constraint():
    """Every point on an IsoFLOPs curve must cost the same: C = 6ND."""
    for C, N, D, _ in _ladder(bias_mode="none"):
        assert 6 * N * D == pytest.approx(C, rel=1e-9)


# ======================================================================================
# 4. The IsoFLOPs method
# ======================================================================================

def test_isoflop_optimum_finds_an_interior_minimum():
    C = 1e21
    pts = [(N, C / (6 * N), lab.parametric_loss(N, C / (6 * N), *TRUTH))
           for N in (1e8, 3e8, 1e9, 3e9, 1e10, 3e10)]
    N_opt, (a, b, c) = lab.isoflop_optimum(pts)
    assert a > 0                              # opens upward: it IS a minimum
    assert min(N for N, _, _ in pts) < N_opt < max(N for N, _, _ in pts)


def test_isoflop_optimum_needs_at_least_three_sizes():
    with pytest.raises(ValueError):
        lab.isoflop_optimum([(1e9, 1e11, 2.0), (1e10, 1e10, 2.1)])


def test_isoflops_ladder_returns_one_optimum_per_budget_sorted():
    res = lab.isoflops_ladder(_ladder(bias_mode="none"))
    assert len(res) == len(BUDGETS)
    assert [C for C, _, _ in res] == sorted(BUDGETS)
    for C, N, D in res:
        assert 6 * N * D == pytest.approx(C, rel=1e-9)


def test_optimal_size_grows_with_budget():
    res = lab.isoflops_ladder(_ladder(bias_mode="none"))
    sizes = [N for _, N, _ in res]
    assert sizes == sorted(sizes)


def test_exponents_sum_to_one():
    """THE FREE BUG DETECTOR. a + b == 1 follows directly from C = 6ND, so any fit
    violating it has an error in the FLOP accounting or in the fit itself."""
    exps = lab.fit_scaling_exponents(lab.isoflops_ladder(_ladder(bias_mode="none")))
    assert exps["exponent_sum"] == pytest.approx(1.0, abs=1e-6)
    assert exps["consistent"] is True


def test_fitted_exponent_is_near_one_half():
    """Chinchilla's headline: N and D scale at roughly the same rate."""
    exps = lab.fit_scaling_exponents(lab.isoflops_ladder(_ladder(bias_mode="none")))
    assert 0.35 < exps["n_exponent"] < 0.65


def test_fit_scaling_exponents_needs_two_budgets():
    with pytest.raises(ValueError):
        lab.fit_scaling_exponents([(1e20, 1e9, 1.6e10)])


# ======================================================================================
# 5. Fitting the surface
# ======================================================================================

def test_residuals_vanish_at_the_true_parameters():
    pts = _points(_ladder(bias_mode="none"))
    for space in ("log", "linear"):
        rs = lab.residuals(pts, TRUTH, space)
        assert max(abs(r) for r in rs) < 1e-9


def test_residuals_rejects_unknown_space():
    with pytest.raises(ValueError):
        lab.residuals(_points(_ladder(bias_mode="none")), TRUTH, space="quaternion")


def test_huber_is_quadratic_near_zero_and_linear_in_the_tail():
    delta = 1e-3
    small = 1e-4
    assert lab.huber_loss([small], delta) == pytest.approx(0.5 * small ** 2)
    big = 1.0
    assert lab.huber_loss([big], delta) == pytest.approx(delta * (big - 0.5 * delta))


def test_huber_rejects_nonpositive_delta():
    with pytest.raises(ValueError):
        lab.huber_loss([0.1], delta=0.0)


def test_huber_suppresses_a_diverged_run_that_least_squares_chases():
    """One diverged run costs least squares orders of magnitude more than Huber.
    This is why Chinchilla used a robust loss, and it is not incidental: real ladders
    contain diverged runs and bad data shards."""
    pts = _points(_ladder(bias_mode="none"))
    dirty = pts[:-1] + [(pts[-1][0], pts[-1][1], pts[-1][2] + 1.5)]
    rs = lab.residuals(dirty, TRUTH, "log")
    assert lab.squared_loss(rs) > 100 * lab.huber_loss(rs)


def test_fit_parametric_recovers_truth_on_clean_data():
    pts = _points(_ladder(bias_mode="none"))
    params, score = lab.fit_parametric(pts)
    # The exponents are what drive the recommendation; check those tightly.
    assert params[2] == pytest.approx(TRUTH[2], abs=0.08)   # alpha
    assert params[4] == pytest.approx(TRUTH[4], abs=0.08)   # beta
    assert score < 1e-4


def test_fit_parametric_is_deterministic():
    """DETERMINISM: no library optimizer, no randomness. Same input, same bytes.
    A fit that moves between runs cannot support a nine-figure decision."""
    pts = _points(_ladder(bias_mode="none"))
    assert lab.fit_parametric(pts) == lab.fit_parametric(pts)


def test_fit_parametric_validates():
    pts = _points(_ladder(bias_mode="none"))
    with pytest.raises(ValueError):
        lab.fit_parametric(pts, objective="wishful")
    with pytest.raises(ValueError):
        lab.fit_parametric(pts, space="quaternion")
    with pytest.raises(ValueError):
        lab.fit_parametric(pts[:2])


def test_estimator_choice_changes_the_recommendation():
    """Feinberg's open problem, made concrete: 'Least squares vs MLE ... imply
    different scaling recommendations! Formalize.'"""
    pts = _points(_ladder(noise_std=0.02, seed=7, bias_mode="none"))
    recs = []
    for obj, space in (("squared", "linear"), ("squared", "log"), ("huber", "log")):
        p, _ = lab.fit_parametric(pts, objective=obj, space=space)
        recs.append(lab.analytic_optimum(1e24, *p)[0])
    assert max(recs) != min(recs), "different estimators should not all agree exactly"


# ======================================================================================
# 6. The analytic optimum
# ======================================================================================

def test_analytic_optimum_satisfies_the_budget():
    N, D = lab.analytic_optimum(1e24, *TRUTH)
    assert 6 * N * D == pytest.approx(1e24, rel=1e-9)


def test_analytic_optimum_beats_its_neighbours():
    """INVARIANT: the closed form must actually be the minimum along the IsoFLOPs curve."""
    C = 1e22
    N, D = lab.analytic_optimum(C, *TRUTH)
    best = lab.parametric_loss(N, D, *TRUTH)
    for factor in (0.5, 0.8, 1.25, 2.0):
        N2 = N * factor
        assert lab.parametric_loss(N2, C / (6 * N2), *TRUTH) >= best


def test_analytic_optimum_agrees_with_the_numerical_sweep():
    """The two methods in this phase must give the same answer."""
    C = 1e21
    N_analytic, _ = lab.analytic_optimum(C, *TRUTH)
    pts = [(N, C / (6 * N), lab.parametric_loss(N, C / (6 * N), *TRUTH))
           for N in (1e8, 3e8, 1e9, 3e9, 1e10)]
    N_sweep, _ = lab.isoflop_optimum(pts)
    # A 5-point parabola fit is a coarse approximation; agree within a factor of 3.
    assert 1 / 3 < N_sweep / N_analytic < 3


def test_optimum_exponent_is_half_when_alpha_equals_beta():
    """Chinchilla's ~0.5 corresponds exactly to alpha == beta."""
    assert lab.optimum_exponent(0.34, 0.34) == pytest.approx(0.5)
    assert lab.optimum_exponent(0.34, 0.28) < 0.5


def test_analytic_optimum_validates():
    with pytest.raises(ValueError):
        lab.analytic_optimum(0, *TRUTH)
    with pytest.raises(ValueError):
        lab.optimum_exponent(0.0, 0.3)


# ======================================================================================
# 7. Uncertainty
# ======================================================================================

def test_bootstrap_is_deterministic_under_a_seed():
    pts = _points(_ladder(noise_std=0.01, seed=1, bias_mode="none"))
    a = lab.bootstrap_forecast(pts, 1e24, n_boot=30, seed=5)
    b = lab.bootstrap_forecast(pts, 1e24, n_boot=30, seed=5)
    assert a == b


def test_bootstrap_interval_brackets_its_median():
    pts = _points(_ladder(noise_std=0.01, seed=1, bias_mode="none"))
    r = lab.bootstrap_forecast(pts, 1e24, n_boot=40, seed=5)
    assert r["ci_low"] <= r["median"] <= r["ci_high"]
    assert r["width"] == pytest.approx(r["ci_high"] - r["ci_low"])
    assert r["n_successful"] > 0


def test_noisier_ladder_gives_a_wider_interval():
    quiet = _points(_ladder(noise_std=0.002, seed=1, bias_mode="none"))
    loud = _points(_ladder(noise_std=0.05, seed=1, bias_mode="none"))
    w_quiet = lab.bootstrap_forecast(quiet, 1e24, n_boot=40, seed=5)["width"]
    w_loud = lab.bootstrap_forecast(loud, 1e24, n_boot=40, seed=5)["width"]
    assert w_loud > w_quiet


def test_bootstrap_validates():
    pts = _points(_ladder(bias_mode="none"))
    with pytest.raises(ValueError):
        lab.bootstrap_forecast(pts, 1e24, n_boot=2)
    with pytest.raises(ValueError):
        lab.bootstrap_forecast(pts, 1e24, ci=1.5)


def test_loss_delta_to_compute_multiplier():
    """0.01 nats is roughly a third more compute. This is the scale intuition that makes
    a confidence interval interpretable."""
    m = lab.compute_multiplier_for_loss_delta(0.01)
    assert 0.2 < m < 0.5
    # Linear in the delta.
    assert (lab.compute_multiplier_for_loss_delta(0.02)
            == pytest.approx(2 * m))


def test_compute_multiplier_rejects_loss_below_floor():
    with pytest.raises(ValueError):
        lab.compute_multiplier_for_loss_delta(0.01, loss=1.5, irreducible=1.69)


# ======================================================================================
# 8. Experimental design
# ======================================================================================

def test_spread_beats_density_for_extrapolation():
    """13x lower variance from the SAME number of runs, purely from placement."""
    target = 25.0
    clustered = lab.extrapolation_variance([19.0, 19.2, 19.4, 19.6], target)
    spread = lab.extrapolation_variance([18.0, 19.0, 20.0, 21.0], target)
    assert spread < clustered / 10


def test_two_well_placed_runs_beat_four_clustered_ones():
    """Placement dominates count. Do not fill in the grid; go further out."""
    target = 25.0
    clustered4 = lab.extrapolation_variance([19.0, 19.2, 19.4, 19.6], target)
    extremes2 = lab.extrapolation_variance([18.0, 21.0], target)
    assert extremes2 < clustered4


def test_extrapolation_variance_validates():
    with pytest.raises(ValueError):
        lab.extrapolation_variance([19.0], 25.0)
    with pytest.raises(ValueError):
        lab.extrapolation_variance([19.0, 19.0, 19.0], 25.0)


def test_design_ladder_stays_within_budget_and_maximizes_spread():
    flagship = 1e25
    budgets = lab.design_ladder(flagship, ladder_budget_fraction=0.03, n_budgets=5)
    assert len(budgets) == 5
    assert sum(budgets) < flagship * 0.03
    assert budgets == sorted(budgets)
    decades = math.log10(budgets[-1] / budgets[0])
    assert decades > 3


def test_design_ladder_validates():
    with pytest.raises(ValueError):
        lab.design_ladder(1e25, ladder_budget_fraction=1.5)
    with pytest.raises(ValueError):
        lab.design_ladder(1e25, n_budgets=1)
    with pytest.raises(ValueError):
        lab.design_ladder(1e19, ladder_budget_fraction=0.001, min_flops=1e18)


def test_extrapolation_risk_grows_with_distance_and_uncertainty():
    near = lab.extrapolation_risk(1e24, 1e25, 0.005)
    far = lab.extrapolation_risk(1e20, 1e25, 0.005)
    wide = lab.extrapolation_risk(1e24, 1e25, 0.05)
    assert far["score"] > near["score"]
    assert wide["score"] > near["score"]


def test_interpolation_is_flagged_as_such():
    """BOUNDARY: if the target is inside the ladder, you are not extrapolating."""
    r = lab.extrapolation_risk(1e24, 1e23, 0.01)
    assert r["band"] == "INTERPOLATION"


# ======================================================================================
# 9. The decision
# ======================================================================================

def test_compare_recipes_finds_the_crossover():
    """A candidate that learns faster but has a higher floor MUST cross the baseline.
    The crossover, not the winner, is the deliverable — 'better' assumes a scale."""
    baseline = TRUTH
    candidate = (1.85, 300.0, 0.40, 300.0, 0.33)
    small = lab.compare_recipes(baseline, candidate, 1e20)
    large = lab.compare_recipes(baseline, candidate, 1e26)
    assert small["candidate_wins"] is True
    assert large["candidate_wins"] is False
    assert small["crossover_flops"] is not None


def test_compare_recipes_reports_a_signed_delta():
    r = lab.compare_recipes(TRUTH, TRUTH, 1e24)
    assert r["delta"] == pytest.approx(0.0)
    assert r["baseline_loss"] == pytest.approx(r["candidate_loss"])


def test_compare_recipes_validates_scan_range():
    with pytest.raises(ValueError):
        lab.compare_recipes(TRUTH, TRUTH, 1e24, scan_lo=25.0, scan_hi=20.0)


def test_a_delta_smaller_than_the_interval_is_not_a_result():
    """The single most useful sentence a scaling engineer can say in a review:
    'that difference is inside our error bars.'"""
    assert lab.decision_is_supported(0.05, 0.01) is True
    assert lab.decision_is_supported(0.005, 0.01) is False


def test_loss_at_budget_improves_monotonically_with_compute():
    losses = [lab.loss_at_budget(TRUTH, C) for C in (1e20, 1e22, 1e24, 1e26)]
    assert losses == sorted(losses, reverse=True)