import json
import math
from pathlib import Path
import numpy as np


def tail_coefficients(eta):
    eta = np.asarray(eta, dtype=float)
    tau = len(eta)
    c = np.empty(tau)
    for t in range(tau - 1):
        c[t] = eta[t] ** 2 / (2.0 * np.sum(eta[t + 1:]))
    c[-1] = eta[-1] / 2.0
    return c


def continuous_batches(c, s, budget):
    r = np.sqrt(np.maximum(c * s, 1e-30))
    return budget * r / r.sum()


def integer_batches(weights, budget, bmin=1, bmax=None):
    """Allocate integer examples, preserving budget and bounds as closely as possible."""
    n = len(weights)
    if bmax is None:
        bmax = budget
    if budget < n * bmin or budget > n * bmax:
        raise ValueError("infeasible bounds")
    x = budget * np.asarray(weights, dtype=float) / np.sum(weights)
    x = np.clip(x, bmin, bmax)
    # Iterative largest-remainder allocation with bounds.
    b = np.floor(x).astype(int)
    b = np.maximum(b, bmin)
    b = np.minimum(b, bmax)
    while b.sum() < budget:
        candidates = np.where(b < bmax)[0]
        score = x[candidates] - b[candidates]
        b[candidates[np.argmax(score)]] += 1
    while b.sum() > budget:
        candidates = np.where(b > bmin)[0]
        score = b[candidates] - x[candidates]
        b[candidates[np.argmax(score)]] -= 1
    return b


def objective(b, c, s):
    return float(np.sum(np.asarray(c) * np.asarray(s) / np.asarray(b)))


def cosine_lr(tau, eta_max=0.12, eta_min=0.012):
    x = np.arange(tau) / max(1, tau - 1)
    return eta_min + 0.5 * (eta_max - eta_min) * (1 + np.cos(np.pi * x))


def verify_math(seed=7):
    rng = np.random.default_rng(seed)
    tau = 32
    eta = cosine_lr(tau)
    c = tail_coefficients(eta)
    s = np.exp(rng.normal(0, 0.8, tau))
    C = 32 * 16
    b_star = continuous_batches(c, s, C)
    b_int = integer_batches(np.sqrt(c * s), C, 1, 64)
    uniform = np.full(tau, C / tau)
    # Predictions: (P1) ratios match sqrt(c*s); (P2) continuous optimum has the
    # Cauchy-Schwarz value (sum sqrt(cs))^2/C; (P3) multiplying all noise by k
    # scales batches by sqrt(k), while relative allocations do not change.
    ratio_err = np.max(np.abs((b_star / b_star[0]) / (np.sqrt(c*s) / np.sqrt(c[0]*s[0])) - 1))
    predicted = np.sum(np.sqrt(c * s)) ** 2 / C
    opt_gap = objective(b_star, c, s) / predicted - 1
    scales = np.array([0.25, 1.0, 4.0, 16.0])
    scale_rows = []
    for k in scales:
        bk = continuous_batches(c, k*s, C)
        # Relative allocation is predicted invariant; because C is fixed, absolute
        # batches are also invariant. The meaningful scaling prediction is objective.
        scale_rows.append({"noise_multiplier": float(k),
                           "objective_ratio_observed": objective(bk, c, k*s) / objective(b_star, c, s),
                           "objective_ratio_predicted": float(k),
                           "allocation_relative_max_error": float(np.max(np.abs(bk/bk.sum()-b_star/b_star.sum())))})
    # Tail sweep: for constant noise, compare first/last coefficient under cosine
    # horizons. The reported quantity is directly predicted by c_t formula.
    tail_rows = []
    for T in [8, 16, 32, 64, 128]:
        et = cosine_lr(T)
        ct = tail_coefficients(et)
        bt = continuous_batches(ct, np.ones(T), T * 16)
        tail_rows.append({"horizon": T,
                          "c_first_over_c_last": float(ct[0]/ct[-1]),
                          "batch_first_over_last": float(bt[0]/bt[-1]),
                          "prediction_error": float(abs(bt[0]/bt[-1] - math.sqrt(ct[0]/ct[-1])))})
    return {"P1_ratio_max_abs_error": float(ratio_err),
            "P2_optimality_relative_error": float(opt_gap),
            "P3_noise_scale_sweep": scale_rows,
            "tail_weight_sweep": tail_rows,
            "integer_objective_over_continuous": objective(b_int,c,s)/objective(b_star,c,s),
            "eta": eta.tolist(), "c": c.tolist(), "s": s.tolist(),
            "continuous_batches": b_star.tolist(), "integer_batches": b_int.tolist()}


def run_sgd(seed, schedule, eta, s_profile, dim=8):
    rng = np.random.default_rng(seed)
    # Strongly convex quadratic; gradient noise has E||noise||^2 approximately s_t/B.
    lam = 0.5
    w = rng.normal(0, 1, dim)
    target = np.zeros(dim)
    losses = []
    for t, b in enumerate(schedule):
        true_g = lam * (w - target)
        noise = rng.normal(0, math.sqrt(s_profile[t] / b / dim), dim)
        w = w - eta[t] * (true_g + noise)
        losses.append(0.5 * lam * float(np.dot(w, w)))
    return np.asarray(losses)


def mini_experiment(seed=19):
    tau, per_step = 64, 16
    C = tau * per_step
    eta = cosine_lr(tau, 0.16, 0.016)
    c = tail_coefficients(eta)
    # Noise falls during optimization but remains deliberately heterogeneous.
    s = 1.0 + 8.0 * np.exp(-np.arange(tau) / 18.0)
    idea = integer_batches(np.sqrt(c*s), C, bmin=4, bmax=48)
    static = np.full(tau, per_step, dtype=int)
    # Common hand-designed comparator: linear growth, same total examples.
    linear_weights = np.linspace(0.5, 1.5, tau)
    linear = integer_batches(linear_weights, C, bmin=4, bmax=48)
    all_results = {}
    for name, sched in [("static", static), ("linear_growth", linear), ("tail_weighted", idea)]:
        curves = np.array([run_sgd(seed+i, sched, eta, s) for i in range(40)])
        all_results[name] = {"final_loss_mean": float(curves[:, -1].mean()),
                             "final_loss_se": float(curves[:, -1].std(ddof=1)/math.sqrt(len(curves))),
                             "loss_at_step_32": float(curves[:,31].mean()),
                             "batch_min": int(sched.min()), "batch_max": int(sched.max()),
                             "batch_first": int(sched[0]), "batch_last": int(sched[-1])}
    all_results["theoretical_batch_objectives"] = {
        "static": objective(static,c,s), "linear_growth": objective(linear,c,s), "tail_weighted": objective(idea,c,s)}
    all_results["eta"] = eta.tolist()
    all_results["c"] = c.tolist()
    all_results["s_profile"] = s.tolist()
    return all_results


if __name__ == "__main__":
    out = {"verification": verify_math(), "mini_experiment": mini_experiment()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"verification": out["verification"], "mini_experiment": out["mini_experiment"]}, indent=2))