#!/usr/bin/env python3
"""Reproduce three finite observations on the decimal orbit of 196.

The windows and estimators are part of the certificate.  Nothing here makes a
claim about the orbit after step 4,000.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from collections import Counter
from pathlib import Path


def reverse_add(n: int) -> int:
    text = str(n)
    return n + int(text[::-1])


def carry_in_field(text: str) -> list[int]:
    digits = [int(char) for char in reversed(text)]
    carry = 0
    field: list[int] = []
    for index, digit in enumerate(digits):
        field.append(carry)
        carry = (digit + digits[-1 - index] + carry) // 10
    return field


def mutual_information(pair_counts: Counter[tuple[int, int]]) -> float:
    total = sum(pair_counts.values())
    left = Counter()
    right = Counter()
    for (a, b), count in pair_counts.items():
        left[a] += count
        right[b] += count
    return sum(
        count / total * math.log2(count * total / (left[a] * right[b]))
        for (a, b), count in pair_counts.items()
    )


def subword_counts(rows: list[str], maximum: int) -> list[int]:
    return [
        len(
            {
                row[start : start + width]
                for row in rows
                for start in range(len(row) - width + 1)
            }
        )
        for width in range(1, maximum + 1)
    ]


def build_certificate() -> dict[str, object]:
    states = ["196"]
    value = 196
    for _ in range(4_000):
        value = reverse_add(value)
        states.append(str(value))

    state_hash = hashlib.sha256(
        "".join(f"{step}:{state}\n" for step, state in enumerate(states)).encode()
    ).hexdigest()

    reflection_rows = states[1:4_001]
    reflection = [
        (step, len(row), sum(a != b for a, b in zip(row, reversed(row))))
        for step, row in enumerate(reflection_rows, start=1)
    ]
    minimum_distance = min(distance for _, _, distance in reflection)
    minimum_ratio = min(
        reflection, key=lambda item: item[2] / item[1]
    )
    long_ratios = [distance / length for _, length, distance in reflection if length > 100]

    shift_sums = {
        growth: {shift: 0.0 for shift in range(-2, 3)}
        for growth in (False, True)
    }
    shift_counts = Counter()
    for step in range(1, 3_001):
        before = states[step - 1]
        row = states[step]
        growth = len(row) > len(before)
        shift_counts[growth] += 1
        length = len(row)
        for shift in range(-2, 3):
            comparisons = 0
            mismatches = 0
            for index, digit in enumerate(row):
                partner = length - 1 - index + shift
                if 0 <= partner < length:
                    comparisons += 1
                    mismatches += digit != row[partner]
            shift_sums[growth][shift] += mismatches / comparisons

    raw_rows = states[:4_000]
    digit_counts = Counter("".join(raw_rows))
    carry_rows = [carry_in_field(row) for row in raw_rows]
    carry_counts = Counter(bit for row in carry_rows for bit in row)
    digit_pairs: Counter[tuple[int, int]] = Counter()
    carry_pairs: Counter[tuple[int, int]] = Counter()
    for row, carries in zip(raw_rows, carry_rows):
        length = len(row)
        for index in range(length):
            digit_pairs[(int(row[index]), int(row[-1 - index]))] += 1
            carry_pairs[(carries[index], carries[-1 - index])] += 1

    digit_total = sum(digit_counts.values())
    carry_total = sum(carry_counts.values())
    carry_strings = ["".join(map(str, row)) for row in carry_rows]

    return {
        "protocol": "fprd.reverse-add-196-orbit-audit.v1",
        "scope_boundary": "finite decimal orbit only; no assertion after T^4000(196)",
        "state_sequence": {
            "range": "T^k(196), 0 <= k <= 4000",
            "sha256_step_colon_state_newline": state_hash,
            "final_digit_length": len(states[4_000]),
        },
        "reflection_distance": {
            "window": "outputs T^k(196), 1 <= k <= 4000",
            "definition": "number of ordered positions differing from reflection",
            "minimum": minimum_distance,
            "minimum_ratio": {
                "step": minimum_ratio[0],
                "distance": minimum_ratio[2],
                "length": minimum_ratio[1],
                "value": round(minimum_ratio[2] / minimum_ratio[1], 12),
            },
            "mean_ratio_when_length_gt_100": round(
                sum(long_ratios) / len(long_ratios), 12
            ),
            "rows_when_length_gt_100": len(long_ratios),
        },
        "shifted_axis": {
            "window": "transitions T^(k-1)(196) -> T^k(196), 1 <= k <= 3000",
            "estimator": "mean of per-output normalized mismatch fractions over valid index pairs",
            "step_counts": {
                "no_growth": shift_counts[False],
                "growth": shift_counts[True],
            },
            "mean_normalized_distance": {
                str(shift): {
                    "no_growth": round(
                        shift_sums[False][shift] / shift_counts[False], 12
                    ),
                    "growth": round(
                        shift_sums[True][shift] / shift_counts[True], 12
                    ),
                }
                for shift in range(-2, 3)
            },
        },
        "raw_complexity": {
            "window": "states T^k(196), 0 <= k < 4000",
            "subwords": "union of contiguous within-state words; row boundaries excluded",
            "digit_subword_counts_k_1_to_4": subword_counts(raw_rows, 4),
            "carry_in_subword_counts_k_1_to_12": subword_counts(carry_strings, 12),
            "digit_count": digit_total,
            "maximum_digit_frequency_deviation_from_one_tenth": round(
                max(abs(digit_counts[str(d)] / digit_total - 0.1) for d in range(10)),
                12,
            ),
            "carry_in_count": carry_total,
            "carry_in_density": round(carry_counts[1] / carry_total, 12),
            "mirrored_digit_mutual_information_bits": round(
                mutual_information(digit_pairs), 12
            ),
            "mirrored_carry_in_mutual_information_bits": round(
                mutual_information(carry_pairs), 12
            ),
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path)
    args = parser.parse_args()
    payload = json.dumps(build_certificate(), indent=2, sort_keys=True) + "\n"
    if args.output:
        args.output.write_text(payload, encoding="utf-8")
    else:
        print(payload, end="")


if __name__ == "__main__":
    main()