import json
import numpy as np
from polar_slack import tetrahedron_design, slack_matrix, polar_attention_logits, row_softmax


def numerical_rank(A, tol=1e-9):
    s = np.linalg.svd(A, compute_uv=False)
    return int(np.sum(s > tol * max(s[0], 1.0))), s


def main():
    np.random.seed(7)
    X = tetrahedron_design()
    d, N = X.shape
    h = np.ones(N)
    U = -np.eye(d)
    c = 1.0 / 3.0
    A = slack_matrix(X, h, c, U)
    rank, singular = numerical_rank(A)
    B = X.T @ U.T @ X
    zero_tol = 1e-10

    # Exact tetrahedral self-polar realization checks.
    math = {
        "shape": list(A.shape),
        "min_A": float(A.min()),
        "max_abs_A_minus_expected": float(np.max(np.abs(A - (4.0 / 3.0) * np.eye(N)))),
        "rank_A": rank,
        "expected_rank_d_plus_1": d + 1,
        "rank_B": numerical_rank(B)[0],
        "B_nonzero_singular_values": [float(v) for v in np.linalg.svd(B, compute_uv=False) if v > 1e-9],
        "row_sums": [float(v) for v in A.sum(1)],
        "zero_entries_per_row": int(np.sum(np.abs(A[0]) <= zero_tol)),
    }

    # Scope check: arbitrary supports need not preserve the paper's guarantee.
    A_invalid = slack_matrix(X, np.array([0.2, 1.0, 1.0, 1.0]), c, U)
    invalid_check = {"min_A_with_arbitrary_h": float(A_invalid.min()),
                     "has_negative_entry": bool(np.any(A_invalid < -zero_tol))}

    # Controlled attention proxy: each position's geometric incidence target is itself.
    # Every method receives exactly the same random nuisance QK logits.
    trials, beta = 2000, 1.0
    noise = np.random.normal(size=(trials, N, N))
    qk = beta * noise
    target = np.broadcast_to(np.arange(N), (trials, N))
    dense_p = row_softmax(qk)
    top_idx = np.argmax(qk, axis=-1)
    top_acc = float(np.mean(top_idx == target))
    geo_logits = polar_attention_logits(qk, A, alpha=1.0, eps=1e-8)
    geo_p = row_softmax(geo_logits)
    geo_idx = np.argmax(geo_logits, axis=-1)
    geo_acc = float(np.mean(geo_idx == target))
    rows = np.arange(N)[None, :]
    dense_target = float(np.mean(dense_p[np.arange(trials)[:, None], rows, target]))
    geo_target = float(np.mean(geo_p[np.arange(trials)[:, None], rows, target]))
    experiment = {
        "trials": trials, "N": N, "beta": beta,
        "dense_target_probability": dense_target,
        "polar_target_probability": geo_target,
        "ordinary_top1_accuracy": top_acc,
        "polar_argmax_accuracy": geo_acc,
        "dense_edges_per_row": N,
        "ordinary_top1_edges_per_row": 1,
        "polar_positive_edges_per_row": int(np.sum(A > zero_tol)),
        "polar_effective_edges_per_row": int(np.sum(A[0] > zero_tol)),
        "polar_cross_entropy": float(-np.mean(np.log(np.maximum(geo_p[np.arange(trials)[:, None], rows, target], 1e-30)))),
        "dense_cross_entropy": float(-np.mean(np.log(np.maximum(dense_p[np.arange(trials)[:, None], rows, target], 1e-30)))),
    }
    print(json.dumps({"math": math, "invalid_relaxation": invalid_check, "experiment": experiment}, indent=2))


if __name__ == '__main__':
    main()