"""Exact Sudoku counts and immutable raw-argmax evaluation artifacts."""

from __future__ import annotations

import hashlib
import json
from datetime import UTC, datetime
from pathlib import Path


def ratio(numerator: int, denominator: int) -> dict:
    return {
        "numerator": numerator,
        "denominator": denominator,
        "value": numerator / denominator if denominator else None,
    }


def valid_board(board: list[int]) -> bool:
    """Structural validity only; this does not compare with a labeled solution."""
    if len(board) != 81 or any(v not in range(9) for v in board):
        return False
    units = [board[r * 9 : (r + 1) * 9] for r in range(9)]
    units += [board[c::9] for c in range(9)]
    units += [
        [board[(br + r) * 9 + bc + c] for r in range(3) for c in range(3)]
        for br in (0, 3, 6)
        for bc in (0, 3, 6)
    ]
    return all(set(unit) == set(range(9)) for unit in units)


def score_stage(predictions: list, labels: list, masks: list) -> dict:
    """Keep every input board, including incorrect or structurally invalid boards."""
    if not predictions or len(predictions) != len(labels) or len(labels) != len(masks):
        raise ValueError("Expected nonempty, aligned predictions, labels, and masks")
    cell_correct = clue_correct = blank_correct = clue_count = 0
    board_correct = csr_count = clue_consistent = 0
    for board, truth, mask in zip(predictions, labels, masks, strict=True):
        if len(board) != 81 or len(truth) != 81 or len(mask) != 81:
            raise ValueError("Expected flat 81-cell boards and masks")
        if any(v not in range(9) for v in truth) or any(v not in (0, 1) for v in mask):
            raise ValueError("Expected labels in 0..8 and binary masks")
        correct = [a == b for a, b in zip(board, truth, strict=True)]
        valid = valid_board(board)
        cell_correct += sum(correct)
        board_correct += all(correct)
        csr_count += valid
        clue_correct += sum(ok for ok, clue in zip(correct, mask, strict=True) if clue)
        blank_correct += sum(ok for ok, clue in zip(correct, mask, strict=True) if not clue)
        clue_count += sum(bool(clue) for clue in mask)
        clue_consistent += valid and all(ok for ok, clue in zip(correct, mask, strict=True) if clue)
    n = len(labels)
    return {
        "sym_acc": ratio(cell_correct, n * 81),
        "board_acc": ratio(board_correct, n),
        "clue_acc": ratio(clue_correct, clue_count),
        "blank_acc": ratio(blank_correct, n * 81 - clue_count),
        "csr": ratio(csr_count, n),
        "clue_consistent_valid_rate": ratio(clue_consistent, n),
    }


def file_sha256(path: Path) -> str:
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


DEFINITIONS = {
    "sym_acc": "Exact cell accuracy across all 81 cells.",
    "board_acc": "All 81 cells equal the labeled solution.",
    "clue_acc": "Exact cell accuracy on labeled clue cells.",
    "blank_acc": "Exact cell accuracy on labeled blank cells.",
    "csr": "Every row, column, and box contains 0..8 exactly once; no solution comparison.",
    "clue_consistent_valid_rate": "Structural validity AND agreement with all labeled clues.",
}


def render_raw_report(report: dict) -> str:
    """Console and Markdown share this rendering of the serialized result record."""
    lines = [
        "# Sudoku raw-argmax evaluation",
        "",
        "No refinement, greedy decoding, oracle labels, or solver completion/repair.",
        "Structural validity is distinct from exact solution agreement.",
        "All boards remain in the denominators; undefined rates are N/A.",
        "",
    ]
    evidence = report["provenance"].get("dataset_identity", {}).get("data_protocol", {})
    if evidence:
        lines.extend([evidence["evidence_label"], ""])
    lines.extend(["| Metric | Correct / total | Percent |", "|---|---:|---:|"])
    for name, metric in report["stages"]["raw_argmax"].items():
        value = metric["value"]
        percent = "N/A" if value is None else f"{100 * value:.4f}"
        lines.append(f"| {name} | {metric['numerator']} / {metric['denominator']} | {percent} |")
    lines.extend(["", *[f"- {key}: {value}" for key, value in report["definitions"].items()]])
    return "\n".join(lines) + "\n"


def write_raw_report(destination: Path, predictions, labels, masks, provenance: dict) -> dict:
    """Save every raw prediction; score the saved arrays and render their exact counts."""
    import numpy as np

    arrays = {
        "raw_argmax": predictions.detach().cpu().numpy(),
        "labels": labels.detach().cpu().numpy(),
        "puzzle_mask": masks.detach().cpu().numpy(),
        "board_index": np.arange(len(labels)),
    }
    metrics = score_stage(
        arrays["raw_argmax"].tolist(), arrays["labels"].tolist(), arrays["puzzle_mask"].tolist()
    )
    destination.mkdir(parents=True, exist_ok=False)
    prediction_path = destination / "predictions.npz"
    np.savez_compressed(prediction_path, **arrays)
    report = {
        "schema_version": 1,
        "created_utc": datetime.now(UTC).isoformat(),
        "decoder": "raw_argmax",
        "provenance": provenance,
        "predictions_sha256": file_sha256(prediction_path),
        "definitions": DEFINITIONS,
        "board_index": "Zero-based order in the selected HDF5 split; no boards filtered.",
        "stages": {"raw_argmax": metrics},
    }
    result_path = destination / "metrics.json"
    result_path.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
    saved = json.loads(result_path.read_text())
    rendered = render_raw_report(saved)
    (destination / "metrics.md").write_text(rendered)
    print(rendered, end="", flush=True)
    return saved