#!/usr/bin/env python3
"""Fit and challenge a simple power-law learning curve with no dependencies."""

from __future__ import annotations

import math
import random
from dataclasses import dataclass


@dataclass(frozen=True)
class Fit:
    floor: float
    scale: float
    exponent: float
    rmse: float

    def predict(self, compute: float) -> float:
        return self.floor + self.scale * compute ** (-self.exponent)


def linear_fit(xs: list[float], ys: list[float]) -> tuple[float, float]:
    xbar, ybar = sum(xs) / len(xs), sum(ys) / len(ys)
    denom = sum((x - xbar) ** 2 for x in xs)
    slope = sum((x - xbar) * (y - ybar) for x, y in zip(xs, ys)) / denom
    return ybar - slope * xbar, slope


def fit_power_law(compute: list[float], loss: list[float]) -> Fit:
    """Grid-search the irreducible floor, then fit in log space."""
    best: Fit | None = None
    upper = min(loss) * 0.995
    for step in range(240):
        floor = upper * step / 240
        residual = [value - floor for value in loss]
        if min(residual) <= 0:
            continue
        intercept, slope = linear_fit(
            [math.log(value) for value in compute],
            [math.log(value) for value in residual],
        )
        candidate = Fit(floor, math.exp(intercept), -slope, 0.0)
        rmse = math.sqrt(
            sum((candidate.predict(x) - y) ** 2 for x, y in zip(compute, loss)) / len(loss)
        )
        candidate = Fit(candidate.floor, candidate.scale, candidate.exponent, rmse)
        if best is None or candidate.rmse < best.rmse:
            best = candidate
    if best is None:
        raise ValueError("No admissible fit")
    return best


def bootstrap_extrapolation(
    compute: list[float], loss: list[float], target: float, draws: int = 500, seed: int = 7
) -> tuple[float, float, float]:
    rng = random.Random(seed)
    estimates = []
    for _ in range(draws):
        indices = [rng.randrange(len(compute)) for _ in compute]
        if len({compute[i] for i in indices}) < 2:
            continue
        pairs = sorted((compute[i], loss[i]) for i in indices)
        try:
            fit = fit_power_law([x for x, _ in pairs], [y for _, y in pairs])
            estimates.append(fit.predict(target))
        except (ValueError, ZeroDivisionError):
            pass
    estimates.sort()
    if len(estimates) < 20:
        raise RuntimeError("Too few valid bootstrap fits")
    pick = lambda q: estimates[min(len(estimates) - 1, int(q * len(estimates)))]
    return pick(0.05), pick(0.50), pick(0.95)


def main() -> None:
    compute = [1, 2, 4, 8, 16, 32, 64, 128]
    loss = [1.48, 1.17, 0.96, 0.82, 0.72, 0.66, 0.62, 0.595]
    fit = fit_power_law(compute, loss)
    target = 1024
    interval = bootstrap_extrapolation(compute, loss, target)
    print(f"fit: floor={fit.floor:.3f}, scale={fit.scale:.3f}, exponent={fit.exponent:.3f}")
    print(f"in-range RMSE: {fit.rmse:.4f}")
    print(f"loss at {target} compute: {fit.predict(target):.3f}")
    print(f"bootstrap 90% interval: [{interval[0]:.3f}, {interval[2]:.3f}]")

    # Negative control: a late regime change should make distant extrapolation untrustworthy.
    shifted = loss[:-2] + [0.67, 0.67]
    shifted_fit = fit_power_law(compute, shifted)
    print(f"negative-control forecast: {shifted_fit.predict(target):.3f}")
    print("Compare the forecasts; a narrow curve fit is not evidence that the regime persists.")


if __name__ == "__main__":
    main()