#!/usr/bin/env python3
"""PIST 8×8 strand adjacency matrix builder.

This is a raw-feature shim: it carries no admissibility logic. It tokenizes
input text, assigns each token to a strand by vocab_index % 8, and counts
bigram adjacencies projected onto strands.

Output schema matches `rrc_pist_predictions_250_v1.json` (matrix-only).
"""

from __future__ import annotations

import argparse
import hashlib
import json
import re
from typing import Any


def tokenize(text: str) -> list[str]:
    """Split text into normalized tokens."""
    text = text.lower()
    # Keep letters, digits, and a small set of math symbols
    tokens = re.findall(r"[a-z0-9]+|[+\-*/=^(){}\[\]]", text)
    return tokens


def strand(token: str, vocab: list[str]) -> int:
    return vocab.index(token) % 8


def build_matrix(tokens: list[str]) -> list[list[int]]:
    """Build 8×8 strand adjacency matrix from token bigrams."""
    vocab = sorted(set(tokens))
    matrix = [[0 for _ in range(8)] for _ in range(8)]
    for t_i, t_j in zip(tokens, tokens[1:]):
        matrix[strand(t_i, vocab)][strand(t_j, vocab)] += 1
    return matrix


def canonical_json_matrix(matrix: list[list[int]]) -> str:
    """Canonical row-major JSON with no whitespace for hashing."""
    return json.dumps(matrix, separators=(",", ":"))


def build_record(text: str, equation_id: str, name: str) -> dict[str, Any]:
    tokens = tokenize(text)
    vocab = sorted(set(tokens))
    matrix = build_matrix(tokens)
    matrix_json = canonical_json_matrix(matrix)
    return {
        "equation_id": equation_id,
        "name": name,
        "schema": "rrc_pist_predictions_250_v1",
        "claim_boundary": "matrix-only;no-classifier;no-lean-spectral",
        "matrix_schema": "token_strand_adjacency_8x8_v1",
        "matrix_8x8": matrix,
        "matrix_hash": hashlib.sha256(matrix_json.encode("utf-8")).hexdigest(),
        "global_vocab_hash": hashlib.sha256(
            json.dumps(vocab, separators=(",", ":")).encode("utf-8")
        ).hexdigest(),
        "proxy_pred": None,
        "exact_pred": None,
        "source_records": [{"equation_record_id": equation_id, "name": name}],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Build PIST 8×8 adjacency matrix")
    parser.add_argument("--text", required=True, help="input equation text")
    parser.add_argument("--equation-id", default="rrc_eq_unknown")
    parser.add_argument("--name", default="unknown")
    parser.add_argument("--output", help="write JSON record to file")
    args = parser.parse_args()

    record = build_record(args.text, args.equation_id, args.name)
    out = json.dumps(record, indent=2)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(out + "\n")
    else:
        print(out)


if __name__ == "__main__":
    main()