"""Fresh CPU-only source-table audit and protection regimes.

The learned-model claim is tested literally against the authored Figure 4
values.  The remaining functions are independent, executed protection checks
for the five already-scored claims; they do not relabel construction checks as
learned training.
"""

from __future__ import annotations

import json
import math
import re
import sys
from fractions import Fraction
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from exp3_selcopy_construction import SelCopy, evaluate as eval_selcopy
from exp4_ar_construction import ARDecode, analytic_success, evaluate as eval_ar


def source_blocks() -> tuple[str, str]:
    text = (ROOT / "paper_text.txt").read_text()
    fig4_start = text.index("Selective Copy.")
    fig4_end = text.index("Figure 4:", fig4_start)
    mkar_start = text.index("Multi-Key Associative Recall.")
    fig6_end = text.index("Figure 6:", mkar_start)
    return text[fig4_start:fig4_end], text[mkar_start:fig6_end]


def figure4_literal() -> dict:
    block, _ = source_blocks()
    required = ["0.999", "0.923", "0.931", "2000", "12000"]
    assert all(token in block for token in required), required
    hybrid_2k = 0.999
    pure_tf_12k = 0.923
    pure_ssm_12k = 0.931
    return {
        "hybrid_2k_accuracy": hybrid_2k,
        "pure_tf_12k_accuracy": pure_tf_12k,
        "pure_ssm_12k_accuracy": pure_ssm_12k,
        "nominal_parameter_ratio": 12000 / 2000,
        "hybrid_is_exactly_one": hybrid_2k == 1.0,
        "pure_models_match_hybrid_at_12k": max(pure_tf_12k, pure_ssm_12k) >= hybrid_2k,
        "literal_falsification_gate": hybrid_2k < 1.0 and max(pure_tf_12k, pure_ssm_12k) < hybrid_2k,
        "source_block_sha256": __import__("hashlib").sha256(block.encode()).hexdigest(),
    }


def protect_claim1() -> dict:
    # Under the printed injectivity premise, |V|^m <= |Y|^q.  The printed
    # difference is consequently non-positive; use integer comparisons before
    # evaluating the log expression.
    checked = 0
    admissible = 0
    max_rhs = -float("inf")
    for m in range(1, 49):
        for q in range(1, 49):
            for v in range(2, 17):
                for y in range(2, 65):
                    checked += 1
                    if v**m <= y**q:
                        admissible += 1
                        rhs = m * math.log2(v) - q * math.log2(y)
                        max_rhs = max(max_rhs, rhs)
    return {"checked_configurations": checked, "admissible_injective_configurations": admissible,
            "maximum_printed_rhs": max_rhs}


def masked_terminal(L: int, window: int, rng: np.random.Generator, x: np.ndarray) -> float:
    lo = max(0, L - window)
    logits = rng.normal(size=L)
    logits[:lo] = -np.inf
    weights = np.exp(logits - np.max(logits[np.isfinite(logits)]))
    weights[:lo] = 0.0
    weights /= weights.sum()
    return float(weights @ x)


def protect_claim2() -> dict:
    max_outside_delta = 0.0
    full_window_deltas = []
    cells = 0
    for L in (16, 32, 64, 128):
        for window in (1, 2, 4, 8, 16):
            if window >= L:
                continue
            for seed in range(10):
                rng = np.random.default_rng(10000 + 31 * L + 7 * window + seed)
                x = rng.normal(size=L)
                x_outside = x.copy()
                x_outside[: L - window] += 3.0
                local = masked_terminal(L, window, rng, x)
                outside = masked_terminal(L, window, rng, x_outside)
                # Reuse identical logits for the actual comparison.
                rng2 = np.random.default_rng(20000 + 31 * L + 7 * window + seed)
                base = rng2.normal(size=L)
                base[: L - window] = -np.inf
                weights = np.exp(base - np.max(base[np.isfinite(base)]))
                weights[: L - window] = 0.0
                weights /= weights.sum()
                delta = float(abs(weights @ x - weights @ x_outside))
                max_outside_delta = max(max_outside_delta, delta)
                full_rng = np.random.default_rng(30000 + 31 * L + 7 * window + seed)
                full_base = full_rng.normal(size=L)
                full_w = np.exp(full_base - np.max(full_base))
                full_w /= full_w.sum()
                full_window_deltas.append(float(abs(full_w @ x - full_w @ x_outside)))
                cells += 1
    return {"cells": cells, "max_outside_perturbation": max_outside_delta,
            "minimum_full_window_perturbation": min(full_window_deltas),
            "maximum_full_window_perturbation": max(full_window_deltas)}


def protect_claim3() -> dict:
    rows = []
    configs = [([1, 2, 3, 4], 4, 8, 6000),
               ([1, 2, 3, 4, 5, 6, 7, 8], 8, 16, 6000),
               (list(range(1, 17)), 16, 32, 8000)]
    for offsets, other, length, n in configs:
        task = SelCopy(offsets, M=other, L=length)
        accuracies = []
        for seed in (101, 202):
            rng = np.random.default_rng(seed)
            X = rng.choice(task.vocab, size=(n // 2, length))
            accuracies.append(eval_selcopy(task, X)["accuracy"])
        rows.append({"offset_count": len(offsets), "other_token_count": other,
                     "L": length, "tested": n, "seed_accuracies": accuracies,
                     "minimum_accuracy": min(accuracies), "window_over_L": task.window / length})
    return {"cells": rows, "minimum_accuracy": min(r["minimum_accuracy"] for r in rows)}


def protect_claim4() -> dict:
    exact = []
    for mw in (2, 4, 8, 16, 32, 64):
        ds = int(math.log2(mw))
        for eligible in (mw, 2 * mw, 4 * mw, 8 * mw):
            for in_window in (mw, 2 * mw, 4 * mw):
                p, exists = analytic_success(mw, eligible, in_window)
                exact.append({"M": mw, "eligible": eligible, "in_window": in_window,
                              "coverage": p, "key_exists": exists})
    exhaustive = []
    for mw, length in ((2, 4), (4, 6), (8, 8)):
        task = ARDecode(mw, length)
        X, bits = task.sample(5000, np.random.default_rng(7000 + mw))
        row = eval_ar(task, X, bits, window=length)
        row.update({"M": mw, "L": length})
        exhaustive.append(row)
    return {"exact_cells": len(exact), "minimum_exact_coverage": min(r["coverage"] for r in exact),
            "exhaustive_full_window": exhaustive}


def protect_claim6() -> dict:
    _, block = source_blocks()
    required = ["0.512", "0.990", "0.668", "0.517", "0.524", "0.989"]
    assert all(token in block for token in required), required
    hybrid_2k = 0.512
    hybrid_6k = 0.990
    pure_tf_12k = 0.668
    return {"mkar_2k_hybrid": hybrid_2k, "mkar_6k_hybrid": hybrid_6k,
            "mkar_12k_pure_tf": pure_tf_12k,
            "first_hybrid_above_60_percent_parameters": 6000,
            "nearest_pure_tf_parameters": 12000,
            "parameter_ratio_at_that_crossing": 2.0,
            "source_block_sha256": __import__("hashlib").sha256(block.encode()).hexdigest()}


def main() -> None:
    result = {"figure4_literal": figure4_literal(),
              "claim1_protection": protect_claim1(),
              "claim2_protection": protect_claim2(),
              "claim3_protection": protect_claim3(),
              "claim4_protection": protect_claim4(),
              "claim6_protection": protect_claim6()}
    print(json.dumps(result, sort_keys=True, indent=2))


if __name__ == "__main__":
    main()