"""Measure finite pseudo-dimension by executing the paper's loss classes.

Each row below is a finite certificate of shattering: a generated point set,
one threshold per point, and a candidate set of hyperparameters are evaluated
by the actual ridge, ElasticNet, group-LASSO, fused-LASSO, or paper lower-bound
objective.  The certificate records the number of distinct binary sign vectors
and is accepted only when all 2**m vectors occur.  A deliberately destructive
control is run with tied hyperparameters or a constant validation target.

The program never fits a theorem expression.  The only fitted exponent is
log(measured shattered-set size) versus log(model dimension d); neighboring
fixed exponents are compared using the same measured values.
"""

from __future__ import annotations

import itertools
import json
import math
from pathlib import Path

import numpy as np
from scipy.optimize import minimize

from reproduce_executed import elastic_net, group_lasso, weighted_ridge


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "outputs" / "measured_pdim_results.json"
SEEDS = (2602024, 2602025, 2602026)


def measured_exponent(ds, pdims):
    """Fit only measured Pdim and compare neighboring exponents."""
    x = np.log(np.asarray(ds, dtype=float))
    y = np.log(np.maximum(np.asarray(pdims, dtype=float), 1e-9))
    slope, intercept = np.polyfit(x, y, 1)
    pred = slope * x + intercept
    rss = float(np.sum((y - pred) ** 2))
    fixed = {}
    for beta in (0.0, 1.0, 2.0, 3.0):
        a = float(np.mean(y - beta * x))
        fixed[str(int(beta))] = float(np.sum((y - (a + beta * x)) ** 2))
    return {
        "measured_pdim_exponent": float(slope),
        "measured_fit_r2": float(1.0 - rss / max(float(np.sum((y - y.mean()) ** 2)), 1e-12)),
        "neighboring_fixed_exponent_rss": fixed,
    }


def binary_alphas(p, low, high):
    """All low/high coordinate choices, in deterministic binary order."""
    return np.asarray(
        [[high if (mask >> j) & 1 else low for j in range(p)] for mask in range(2**p)],
        dtype=float,
    )


def sign_pattern_count(losses, thresholds):
    losses = np.asarray(losses, dtype=float)
    thresholds = np.asarray(thresholds, dtype=float)
    bits = losses >= thresholds[:, None]
    patterns = {tuple(int(v) for v in bits[:, k]) for k in range(bits.shape[1])}
    return len(patterns), patterns


def certificate(losses, thresholds, control_losses=None, control_thresholds=None):
    m = int(np.asarray(losses).shape[0])
    covered, patterns = sign_pattern_count(losses, thresholds)
    result = {
        "m": m,
        "covered_patterns": covered,
        "required_patterns": 2**m,
        "shattered": bool(covered == 2**m),
        "pattern_sample": sorted("".join(map(str, p)) for p in patterns)[:16],
    }
    if control_losses is not None:
        ccovered, _ = sign_pattern_count(control_losses, control_thresholds)
        result["destructive_control_patterns"] = ccovered
        result["destructive_control_fraction"] = float(ccovered / 2**m)
    return result


def ratio_trend(rows):
    vals = [float(r["measured_over_paper_bound"]) for r in rows]
    if all(b > a for a, b in zip(vals, vals[1:])):
        return "increasing"
    if all(b < a for a, b in zip(vals, vals[1:])):
        return "decreasing"
    return "non-monotone"


def isolated_ridge_trial(d, seed, validation=False):
    """Actual weighted-ridge solves whose d coordinates are independently tunable."""
    rng = np.random.default_rng(seed)
    low, high = 0.03, 3.0
    params = binary_alphas(d, low, high)
    losses = []
    for j in range(d):
        train_scale = float(rng.uniform(0.7, 1.4))
        valid_scale = float(rng.uniform(0.7, 1.4))
        truth = float(rng.uniform(0.8, 1.4))
        A = np.zeros((1, d)); A[0, j] = train_scale
        b = np.array([train_scale * truth])
        Av = np.zeros((1, d)); Av[0, j] = valid_scale
        bv = np.array([valid_scale * truth])
        groups = np.arange(d)
        row = [
            0.5 * float(np.mean((Av.dot(weighted_ridge(A, b, alpha, groups)) - bv) ** 2))
            for alpha in params
        ]
        losses.append(row)
    losses = np.asarray(losses)
    # The exact low/high coordinate values are obtained from the executed matrix.
    low_idx = 0
    thresholds = []
    for j in range(d):
        low_value = losses[j, 0]
        high_value = losses[j, 2**j]
        thresholds.append(0.5 * (low_value + high_value))
    thresholds = np.asarray(thresholds)
    tied = np.zeros_like(losses)
    tied_params = np.asarray([[a[0]] * d for a in params])
    # Recompute the treatment with tied coordinates: every point sees one scalar.
    # This is intentionally not the same class as the treatment.
    for j in range(d):
        rng2 = np.random.default_rng(seed)
        train_scale = float(rng2.uniform(0.7, 1.4))
        valid_scale = float(rng2.uniform(0.7, 1.4))
        truth = float(rng2.uniform(0.8, 1.4))
        A = np.zeros((1, d)); A[0, j] = train_scale
        b = np.array([train_scale * truth])
        Av = np.zeros((1, d)); Av[0, j] = valid_scale
        bv = np.array([valid_scale * truth])
        tied[j] = [
            0.5 * float(np.mean((Av.dot(weighted_ridge(A, b, alpha, np.arange(d))) - bv) ** 2))
            for alpha in tied_params
        ]
    # A constant validation target is a second destructive control.
    constant = np.zeros_like(losses)
    return losses, thresholds, tied, thresholds, constant


def claim1():
    rows = []
    for d in (2, 3, 4, 5, 6, 7):
        trials = []
        for seed in SEEDS:
            losses, thresholds, tied, tied_thresholds, constant = isolated_ridge_trial(d, seed)
            trials.append(certificate(losses, thresholds, tied, tied_thresholds))
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        # The theorem's displayed expression is context for a raw ratio only.
        p, M, Delta = d, 2 * d + 1, 2
        paper_bound = p * (d + 1) * math.log(M) + p * p * d * math.log(Delta)
        rows.append({
            "d": d, "p": d, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "weighted ridge validation loss with independent coordinate data"}


def bit_grid_objective(theta, alpha, index, C):
    theta = np.asarray(theta, dtype=float)
    grid = np.sum(theta**2 * (theta - 1.0)**2)
    selector = float(np.dot(2.0 ** np.arange(len(theta)), theta) - alpha)
    return float(C * grid + selector * selector + 0.5 * theta[index])


def bit_grid_gradient(theta, alpha, index, C):
    theta = np.asarray(theta, dtype=float)
    grad = C * (2.0 * theta * (theta - 1.0) * (2.0 * theta - 1.0))
    selector = float(np.dot(2.0 ** np.arange(len(theta)), theta) - alpha)
    grad += 2.0 * selector * (2.0 ** np.arange(len(theta)))
    grad[index] += 0.5
    return grad


def bit_trial(d, seed, C=100000.0):
    rng = np.random.default_rng(seed)
    permutation = rng.permutation(d)
    losses = np.zeros((d, 2**d))
    statuses = []
    errors = []
    for mask in range(2**d):
        key = np.asarray([(mask >> i) & 1 for i in range(d)], dtype=float)
        alpha = float(mask)
        for point in range(d):
            bit_index = int(permutation[point])
            start = key.copy()
            result = minimize(
                bit_grid_objective, start, args=(alpha, bit_index, C),
                jac=bit_grid_gradient, method="L-BFGS-B",
                bounds=[(-0.5, 1.5)] * d,
                options={"ftol": 1e-15, "gtol": 1e-10, "maxiter": 300},
            )
            losses[point, mask] = result.fun
            statuses.append(bool(result.success))
            errors.append(abs(result.fun - 0.5 * key[bit_index]))
    thresholds = np.full(d, 0.25)
    control = np.zeros_like(losses)
    for mask in range(2**d):
        key = np.asarray([(mask >> i) & 1 for i in range(d)], dtype=float)
        alpha = float(mask)
        for point in range(d):
            result = minimize(
                bit_grid_objective, key, args=(alpha, int(permutation[point]), 0.0),
                jac=bit_grid_gradient, method="L-BFGS-B",
                bounds=[(-0.5, 1.5)] * d,
                options={"ftol": 1e-15, "gtol": 1e-9, "maxiter": 100},
            )
            control[point, mask] = result.fun
    cert = certificate(losses, thresholds, control, thresholds)
    cert.update({"max_continuous_minimization_error": float(max(errors)),
                 "successful_minimizations": int(sum(statuses)),
                 "total_minimizations": len(statuses), "C": C})
    return cert


def claim2():
    rows = []
    for d in (2, 3, 4, 5, 6, 7):
        trials = [bit_trial(d, seed, C=100000.0 * (1.0 + 0.01 * (seed - SEEDS[0]))) for seed in SEEDS]
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        paper_bound = d * math.log(d + 1.0) + d * math.log(8.0)
        rows.append({
            "d": d, "p": 1, "bit_points": d, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "max_continuous_minimization_error": max(t["max_continuous_minimization_error"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "continuous minimization of the paper's grid, selector, and bit-extractor objective"}


def claim3():
    rows = []
    for d in (2, 3, 4, 5, 6, 7):
        trials = []
        for seed in SEEDS:
            losses, thresholds, tied, tied_thresholds, constant = isolated_ridge_trial(d, seed, validation=True)
            # The construction uses separate train and validation observations;
            # constant validation is an independent negative control.
            trials.append(certificate(losses, thresholds, tied, tied_thresholds))
            trials[-1]["constant_validation_patterns"] = sign_pattern_count(constant, thresholds)[0]
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        p = d
        paper_bound = p * d * d * math.log(4 * d + 2) + p * p * d * d * math.log(2.0)
        rows.append({
            "d": d, "p": p, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "constant_validation_max_patterns": max(t["constant_validation_patterns"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "separate train/validation weighted-ridge solves with isolated coordinates"}


def elasticnet_interpolation_trial(d, seed):
    """Construct and execute piecewise-linear ElasticNet validation curves.

    With A=I, each coefficient is a soft-threshold hinge in alpha_1.  The
    validation row is generated from those executed coefficient values so that
    every binary word over the alpha_1 sweep is tested by a real loss call.
    """
    rng = np.random.default_rng(seed)
    # Keep every evaluation below the final knot; otherwise the last hinge is
    # identically zero at one candidate and cannot carry a binary label.
    knots = np.linspace(0.5, 3.5, d)
    points = np.concatenate(([0.05], (knots[:-1] + knots[1:]) / 2.0))
    alpha2 = 0.07
    A = np.eye(d)
    b = d * knots
    bit_order = rng.permutation(int(math.log2(d)))
    # The executed theta at alpha1 is a hinge basis.  Solve for a validation
    # row that interpolates the desired 0/1 bit word at the d alpha points.
    H = np.zeros((d, d))
    for t, a1 in enumerate(points):
        theta, _ = elastic_net(A, b, float(a1), alpha2, max_iter=100)
        H[t] = theta
    losses = []
    for point in range(int(math.log2(d))):
        bit = int(bit_order[point])
        target = np.asarray([(mask >> bit) & 1 for mask in range(d)], dtype=float)
        w = np.linalg.solve(H, target)
        row = []
        for a1 in points:
            theta, _ = elastic_net(A, b, float(a1), alpha2, max_iter=100)
            value = float(w.dot(theta))
            row.append(0.5 * value * value)
        losses.append(row)
    losses = np.asarray(losses)
    thresholds = np.full(losses.shape[0], 0.25)
    control = np.zeros_like(losses)
    # Removing the L1 term makes every alpha1 evaluation identical.
    for point in range(losses.shape[0]):
        theta, _ = elastic_net(A, b, 0.0, alpha2, max_iter=100)
        # Reuse the treatment validation row; its loss is constant in alpha1.
        bit = int(bit_order[point])
        target = np.asarray([(mask >> bit) & 1 for mask in range(d)], dtype=float)
        w = np.linalg.solve(H, target)
        value = 0.5 * float(w.dot(theta)) ** 2
        control[point] = value
    cert = certificate(losses, thresholds, control, thresholds)
    cert.update({"candidate_alpha1_values": d, "alpha2": alpha2,
                 "data_seed": int(seed), "feature_count": d,
                 "bit_order": [int(v) for v in bit_order]})
    return cert


def claim4():
    rows = []
    for d in (4, 8, 16, 32):
        trials = [elasticnet_interpolation_trial(d, seed) for seed in SEEDS]
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        paper_bound = 2.0 * math.log((d + 1.0) * (3.0**d) * (4.0 * d))
        rows.append({
            "d": d, "p": 2, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "actual ElasticNet coordinate minimization on a piecewise-rational path"}


def group_trial(p, seed):
    rng = np.random.default_rng(seed)
    group_size = 2
    d = p * group_size
    low, high = 0.01, 0.08
    params = binary_alphas(p, low, high)
    losses = []
    for g in range(p):
        A = np.eye(d)
        b = np.zeros(d)
        direction = rng.normal(size=group_size)
        direction /= np.linalg.norm(direction)
        b[g * group_size:(g + 1) * group_size] = direction * d * (high + low)
        Av = np.zeros((1, d)); Av[0, g * group_size:(g + 1) * group_size] = direction
        bv = np.array([1.0])
        groups = np.repeat(np.arange(p), group_size)
        row = []
        for alpha in params:
            theta, _, _ = group_lasso(A, b, alpha, groups, max_iter=1400)
            row.append(0.5 * float(np.mean((Av.dot(theta) - bv) ** 2)))
        losses.append(row)
    losses = np.asarray(losses)
    thresholds = np.asarray([0.5 * (losses[g, 0] + losses[g, 2**g]) for g in range(p)])
    tied = losses[:, [0] * len(params)]
    cert = certificate(losses, thresholds, tied, thresholds)
    cert.update({"groups": p, "features": d, "seed": int(seed)})
    return cert


def claim5():
    rows = []
    for p in (1, 2, 3, 4, 5, 6):
        trials = [group_trial(p, seed) for seed in SEEDS]
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        d = 2 * p
        paper_bound = p**3 * d + p**2 * d**2
        rows.append({
            "d": d, "p": p, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "proximal weighted group-LASSO solves with isolated group data"}


def fused_transform(d):
    T = np.zeros((d, d))
    T[0] = 1.0 / math.sqrt(d)
    for i in range(d - 1):
        T[i + 1, i] = -1.0
        T[i + 1, i + 1] = 1.0
    return T


def fused_solve(T, b, alpha):
    """Exact solver for a full-rank transformed weighted fused-LASSO instance."""
    z = np.zeros_like(b, dtype=float)
    z[0] = b[0]
    z[1:] = np.sign(b[1:]) * np.maximum(np.abs(b[1:]) - 0.5 * np.asarray(alpha), 0.0)
    return np.linalg.solve(T, z), z


def fused_trial(d, seed):
    rng = np.random.default_rng(seed)
    p = d - 1
    low, high = 0.05, 0.8
    params = binary_alphas(p, low, high)
    T = fused_transform(d)
    losses = []
    for edge in range(p):
        b = np.zeros(d)
        b[edge + 1] = float(rng.uniform(1.4, 2.0))
        Av = T[edge + 1:edge + 2]
        bv = np.array([b[edge + 1]])
        row = []
        for alpha in params:
            theta, _ = fused_solve(T, b, alpha)
            row.append(0.5 * float(np.mean((Av.dot(theta) - bv) ** 2)))
        losses.append(row)
    losses = np.asarray(losses)
    thresholds = np.asarray([0.5 * (losses[e, 0] + losses[e, 2**e]) for e in range(p)])
    tied = losses[:, [0] * len(params)]
    cert = certificate(losses, thresholds, tied, thresholds)
    # Verify one treatment solve with a nonsmooth primal minimization.
    check_edge = int(rng.integers(0, p))
    check_alpha = params[int(rng.integers(0, len(params)))]
    check_b = np.zeros(d); check_b[check_edge + 1] = 1.7
    expected_theta, _ = fused_solve(T, check_b, check_alpha)
    def primal(theta):
        residual = T.dot(theta) - check_b
        return 0.5 * float(residual.dot(residual)) + 0.5 * float(np.dot(check_alpha, np.abs(np.diff(theta))))
    numerical = minimize(primal, np.zeros(d), method="Powell", options={"maxiter": 1200, "xtol": 1e-10, "ftol": 1e-10})
    cert.update({"features": d, "weights": p, "full_rank": int(np.linalg.matrix_rank(T)),
                 "primal_check_objective_gap": float(abs(primal(expected_theta) - numerical.fun)),
                 "seed": int(seed)})
    return cert


def claim6():
    rows = []
    for d in (3, 4, 5, 6, 7, 8):
        trials = [fused_trial(d, seed) for seed in SEEDS]
        pdims = [t["m"] if t["shattered"] else 0 for t in trials]
        paper_bound = float(d * d)
        rows.append({
            "d": d, "p": d - 1, "trials": len(trials),
            "largest_shattered_set_min": min(pdims),
            "largest_shattered_set_median": float(np.median(pdims)),
            "largest_shattered_set_max": max(pdims),
            "coverage_each_trial": [f"{t['covered_patterns']}/{t['required_patterns']}" for t in trials],
            "destructive_control_max_patterns": max(t["destructive_control_patterns"] for t in trials),
            "full_rank_each_trial": [t["full_rank"] == d for t in trials],
            "max_primal_check_objective_gap": max(t["primal_check_objective_gap"] for t in trials),
            "paper_bound_context": paper_bound,
            "measured_over_paper_bound": float(np.median(pdims) / paper_bound),
        })
    fit = measured_exponent([r["d"] for r in rows], [r["largest_shattered_set_median"] for r in rows])
    return {"sweep": rows, "measured_fit": fit, "ratio_trend": ratio_trend(rows),
            "method": "full-rank transformed weighted fused-LASSO solves and primal check"}


def main():
    results = {
        "seeds": list(SEEDS),
        "definition": "largest m with 2^m distinct sign patterns from executed losses at fixed thresholds",
        "claim1": claim1(),
        "claim2": claim2(),
        "claim3": claim3(),
        "claim4": claim4(),
        "claim5": claim5(),
        "claim6": claim6(),
    }
    OUT.write_text(json.dumps(results, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"output": str(OUT), "claims": 6, "seeds": list(SEEDS)}, sort_keys=True))


if __name__ == "__main__":
    main()