"""Own toy experiment: what is learned or lost by a wearing readout?

Not a model of human memory or a reproduction of a physical experiment.
Small exact enumeration, no dependencies, network, subprocesses or file reads.
Run with Python's standard library: python -B observe.py
"""
from collections import defaultdict
from fractions import Fraction as F
from itertools import product
import json


def walk(profile, destination):
    """A walker reduces each traversed patch by one, with no regrowth."""
    if not 0 <= destination <= len(profile):
        raise ValueError("destination outside path")
    return tuple(max(0, h - int(i < destination)) for i, h in enumerate(profile))


def history_profile(history, initial=(2, 2, 2)):
    profile = initial
    for destination in history:
        profile = walk(profile, destination)
    return profile


def history_classes(length):
    classes = defaultdict(list)
    for history in product((1, 2), repeat=length):
        classes[history_profile(history)].append(history)
    return classes


def passive(height):
    return height, height


def read_then_wear(height):
    return height, max(0, height - 1)


def wear_then_read(height):
    after = max(0, height - 1)
    return after, after


def observation_classes(reader):
    classes = defaultdict(list)
    for height in range(3):
        observation, after = reader(height)
        classes[observation].append(height)
    return dict(classes)


def pair_recovery(reader, keep_first_record=False):
    """Two sequential readers, uniform prior on initial heights 0, 1, 2.

    Optimal one-shot classification: each observation's most probable initial
    height wins. With a uniform prior, number of classes / number of hypotheses.
    The second reader can optionally access the first reader's separate record.
    """
    first, second = defaultdict(list), defaultdict(list)
    for initial in range(3):
        obs1, state = reader(initial)
        obs2, _ = reader(state)
        first[(obs1,)].append(initial)
        key = (obs1, obs2) if keep_first_record else (obs2,)
        second[key].append(initial)
    return F(len(first), 3), F(len(second), 3)


def two_readout_laws():
    """Same visible height 1: frozen or regrowing before a wearing probe.

    Frozen law: every probe removes one unit.
    Regrowing law: before each probe, a unit grows with probability (2-h)/2.
    Observation is the height AFTER that probe. Equal prior on the two laws.
    This investigates a transition law, not recovery of a historical sequence.
    """
    laws = {"frozen": {(1, ()): F(1)}, "regrowing": {(1, ()): F(1)}}
    rows = []
    for count in range(1, 7):
        next_laws = {}
        distributions = {}
        for label, paths in laws.items():
            nxt = defaultdict(F)
            for (h, observations), probability in paths.items():
                growth = F(0) if label == "frozen" else F(2-h, 2)
                for grows, chance in ((0, 1-growth), (1, growth)):
                    if chance:
                        after = max(0, h + grows - 1)
                        nxt[(after, observations + (after,))] += probability * chance
            assert sum(nxt.values()) == 1
            next_laws[label] = dict(nxt)
            obs_dist = defaultdict(F)
            for (_, observations), p in nxt.items():
                obs_dist[observations] += p
            distributions[label] = obs_dist
        union = set(distributions["frozen"]) | set(distributions["regrowing"])
        success = sum(max(distributions["frozen"].get(o, F(0)),
                          distributions["regrowing"].get(o, F(0))) for o in union) / 2
        rows.append({"probes": count, "best_law_guess_success": str(success),
                     "regrowing_probability_still_height_1":
                         str(sum(p for (h, _), p in next_laws["regrowing"].items() if h == 1))})
        laws = next_laws
    return rows


def adaptive_restore(max_waits):
    """Probe a transition law, then restore the original visible height.

    Wait/observe until growth to 2 is seen, or the finite horizon is reached.
    If 2 is seen, take one wearing step: growth at 2 is zero, so this restores 1.
    Keeping the transcript is explicit; no claim of cost-free physical erasure.
    """
    observations = {}
    terminal_heights = {}
    for law in ("frozen", "regrowing"):
        paths = {(1, (), False): F(1)}
        for _ in range(max_waits):
            nxt = defaultdict(F)
            for (h, transcript, done), probability in paths.items():
                if done:
                    nxt[(h, transcript, done)] += probability
                    continue
                growth = F(0) if law == "frozen" else F(2-h, 2)
                for grows, chance in ((0, 1-growth), (1, growth)):
                    if not chance:
                        continue
                    observed = h + grows
                    record = transcript + (("wait", observed),)
                    if observed == 2:
                        # Both laws grow by zero at the maximum height.
                        nxt[(1, record + (("wear", 1),), True)] += probability * chance
                    else:
                        nxt[(observed, record, False)] += probability * chance
            paths = dict(nxt)
        assert sum(paths.values()) == 1
        heights, transcripts = defaultdict(F), defaultdict(F)
        for (h, transcript, _), p in paths.items():
            heights[h] += p
            transcripts[transcript] += p
        assert dict(heights) == {1: F(1)}
        terminal_heights[law] = {str(h): str(p) for h, p in heights.items()}
        observations[law] = transcripts
    keys = set(observations["frozen"]) | set(observations["regrowing"])
    success = sum(max(observations["frozen"].get(k, F(0)),
                      observations["regrowing"].get(k, F(0))) for k in keys) / 2
    expected = 1 - F(1, 2) ** (max_waits + 1)
    assert success == expected
    return {"maximum_waits": max_waits, "law_guess_success_with_transcript": str(success),
            "terminal_visible_height_distributions": terminal_heights,
            "law_guess_success_from_terminal_height_alone": "1/2"}


def finite_restoration(height, initial, max_waits, repair_steps):
    """Control: away from the ceiling, a wearing step may itself regrow.

    Wait until one growth event, then attempt at most repair_steps wearing
    operations. Stop repair as soon as the original height is reached.
    Only the regrowing law is needed to check this restoration limitation.
    """
    assert 0 < initial < height and max_waits >= 1 and repair_steps >= 0
    paths = {(initial, False): F(1)}
    for _ in range(max_waits):
        nxt = defaultdict(F)
        for (h, grown), p in paths.items():
            growth = F(0) if grown else F(height-h, height)
            for event, chance in ((0, 1-growth), (1, growth)):
                if chance:
                    nxt[(h+event, grown or bool(event))] += p*chance
        paths = dict(nxt)
    for _ in range(repair_steps):
        nxt = defaultdict(F)
        for (h, grown), p in paths.items():
            if not grown or h == initial:
                nxt[(h, grown)] += p
                continue
            growth = F(height-h, height)
            for event, chance in ((0, 1-growth), (1, growth)):
                if chance:
                    nxt[(h+event-1, grown)] += p*chance
        paths = dict(nxt)
    assert sum(paths.values()) == 1
    failure = sum(p for (h, _), p in paths.items() if h != initial)
    formula = (1-F(initial, height)**max_waits) * F(height-initial-1, height)**repair_steps
    assert failure == formula
    return {"maximum_height": height, "initial_height": initial,
            "maximum_waits": max_waits, "repair_steps": repair_steps,
            "probability_not_restored_under_regrowing_law": str(failure)}


def certain_two_step_probe():
    """Wear once; if the result is 0, wait once. Always identify the law.

    Under the frozen law this necessarily leaves 0, not the original 1.
    Under regrowth both possible first outcomes end at 1 after the protocol.
    """
    laws = {}
    for law in ("frozen", "regrowing"):
        observations = defaultdict(F)
        first_growth = F(0) if law == "frozen" else F(1, 2)
        for grows, p in ((0, 1-first_growth), (1, first_growth)):
            if not p:
                continue
            after_wear = grows  # initial 1 + grows - 1
            if after_wear == 1:
                observations[(1, (("wear", 1),))] += p
            else:
                after_wait = 0 if law == "frozen" else 1
                observations[(after_wait, (("wear", 0), ("wait", after_wait)))] += p
        laws[law] = dict(observations)
        assert sum(observations.values()) == 1
    transcript_sets = [{trace for (_, trace) in paths} for paths in laws.values()]
    assert transcript_sets[0].isdisjoint(transcript_sets[1])
    assert {h for (h, _) in laws["frozen"]} == {0}
    assert {h for (h, _) in laws["regrowing"]} == {1}
    return {"maximum_operations": 2, "law_identification_success": "1",
            "frozen_final_height": 0, "regrowing_final_height": 1,
            "original_height": 1,
            "scope": "Certainty for these two laws and available operations, not guaranteed state restoration."}


def main():
    histories = []
    for length in range(1, 9):
        classes = history_classes(length)
        assert sum(map(len, classes.values())) == 2**length
        histories.append({"visits": length, "histories": 2**length,
                          "different_profiles": len(classes),
                          "largest_indistinguishable_class": max(map(len, classes.values()))})
    collisions = history_classes(4)
    witness = next((profile, histories) for profile, histories in collisions.items()
                   if any(len(set(h)) == 1 for h in histories) and
                   any(len(set(h)) == 2 for h in histories))
    profile, paths = witness
    # Equal complete states remain equal under every identical finite probe word.
    same = [h for h in paths if len(set(h)) == 1][0]
    mixed = [h for h in paths if len(set(h)) == 2][0]
    checks = 0
    for n in range(7):
        for probes in product((1, 2), repeat=n):
            a, b = history_profile(same), history_profile(mixed)
            for probe in probes:
                a, b = walk(a, probe), walk(b, probe)
                assert a == b
            checks += 1
    readers = {}
    for name, reader in [("passive", passive), ("read_then_wear", read_then_wear),
                         ("wear_then_read", wear_then_read)]:
        first, second = pair_recovery(reader)
        _, shared = pair_recovery(reader, keep_first_record=True)
        readers[name] = {"observation_classes": observation_classes(reader),
                         "first_reader_success": str(first),
                         "second_reader_success_without_record": str(second),
                         "second_reader_success_with_first_record": str(shared)}
    assert readers["passive"]["second_reader_success_without_record"] == "1"
    assert readers["read_then_wear"]["first_reader_success"] == "1"
    assert readers["read_then_wear"]["second_reader_success_without_record"] == "2/3"
    assert readers["read_then_wear"]["second_reader_success_with_first_record"] == "1"
    assert readers["wear_then_read"]["first_reader_success"] == "2/3"
    assert readers["wear_then_read"]["second_reader_success_with_first_record"] == "2/3"
    laws = two_readout_laws()
    print(json.dumps({"history_enumeration": histories,
                      "same_profile_witness": {"profile": profile, "one_bench": same,
                                               "both_benches": mixed},
                      "future_probe_words_checked": checks,
                      "reader_order": readers, "transition_law_probes": laws,
                      "certain_two_step_probe": certain_two_step_probe(),
                      "adaptive_restore": [adaptive_restore(n) for n in range(1, 7)],
                      "finite_restoration_controls": [finite_restoration(H, h, n, k)
                          for H in range(2, 7) for h in range(1, H)
                          for n in (1, 2, 4) for k in (0, 1, 2, 4)]}, indent=2))


if __name__ == "__main__":
    main()