"""Probe a learner for behaviours nobody programmed into it directly.

Each probe is a small experiment with a number that could have come out either way.
Runs the learner inside this process, so only use learner files you trust.

  python probe.py                                        # the seed
  python probe.py --learner runs/first/cands/00004/learner.py --out runs/first/probe.json
"""
from __future__ import annotations

import argparse
import copy
import json
import math
import random
import sysconfig
import time
from pathlib import Path

import numpy as np

from compsearch.classify import load_learner_factory

KB = 1024
PROBE_BYTES = 20 * KB


def read_streams() -> tuple[bytes, bytes]:
    """Python source (the held-out stream) and English prose (Python's changelog)."""
    python = Path("data/holdout.bin").read_bytes()
    english_path = Path("data/english.bin")
    if not english_path.is_file():
        news = Path(sysconfig.get_paths()["stdlib"]).parent / "NEWS.txt"
        english_path.write_bytes(news.read_bytes()[:1024 * KB])
    return python, english_path.read_bytes()


def feed(learner, text: bytes) -> list[tuple[float, float, bool]]:
    """Run text through the learner (it learns). Per byte: (bits, top probability, top was right)."""
    trace = []
    for byte in text:
        probs = np.asarray(learner.predict(), dtype=np.float64)
        probs = probs / probs.sum()
        top = int(probs.argmax())
        trace.append((-math.log2(max(float(probs[byte]), 1e-12)), float(probs[top]), top == byte))
        learner.update(byte)
    return trace


def bpb(trace: list[tuple[float, float, bool]]) -> float:
    return sum(item[0] for item in trace) / len(trace)


def frozen(learner, text: bytes) -> list[tuple[float, float, bool]]:
    """Measure on a copy, so the learner itself does not learn the text."""
    return feed(copy.deepcopy(learner), text)


def probe_learning_curve(factory, python: bytes) -> tuple[dict, object]:
    learner, windows = factory(), []
    for start in range(0, 400 * KB, 100 * KB):
        windows.append(round(bpb(feed(learner, python[start:start + 100 * KB])), 3))
    return {"bits_per_byte_by_100KB_window": windows,
            "improvement_first_to_last": round(windows[0] - windows[-1], 3)}, learner


def probe_forgetting(python_learner, python_probe: bytes, english: bytes) -> dict:
    before = bpb(frozen(python_learner, python_probe))
    feed(python_learner, english[:400 * KB])
    after = bpb(frozen(python_learner, python_probe))
    return {"python_cost_before_english": round(before, 3),
            "python_cost_after_400KB_english": round(after, 3),
            "change_percent": round(100 * (after - before) / before, 2)}


def probe_transfer(factory, python_learner, english: bytes) -> tuple[dict, object]:
    sample = english[:100 * KB]
    fresh = factory()
    from_scratch = bpb(feed(fresh, sample))
    pretrained = bpb(frozen(python_learner, sample))  # copy keeps learning as it reads
    return {"english_cost_from_scratch": round(from_scratch, 3),
            "english_cost_after_python_first": round(pretrained, 3),
            "gain_percent": round(100 * (from_scratch - pretrained) / from_scratch, 2)}, fresh


def probe_order(factory, python: bytes, english: bytes, python_probe: bytes, english_probe: bytes) -> dict:
    first, second = python[:150 * KB], english[:150 * KB]
    results = []
    for order in ((first, second), (second, first)):
        learner = factory()
        for chunk in order:
            feed(learner, chunk)
        results.append((bpb(frozen(learner, python_probe)), bpb(frozen(learner, english_probe))))
    (py_a, en_a), (py_b, en_b) = results
    return {"python_probe_python_first_vs_english_first": [round(py_a, 3), round(py_b, 3)],
            "english_probe_python_first_vs_english_first": [round(en_a, 3), round(en_b, 3)],
            "largest_difference_percent": round(100 * max(abs(py_a - py_b) / py_a, abs(en_a - en_b) / en_a), 2)}


def probe_one_shot(learner, filler: bytes, rng: random.Random) -> dict:
    learner = copy.deepcopy(learner)
    novel = bytes(rng.choice(b"abcdefghijklmnopqrstuvwxyz0123456789") for _ in range(48))
    first = bpb(feed(learner, novel))
    feed(learner, filler[:2 * KB])
    soon = bpb(frozen(learner, novel))
    feed(learner, filler[2 * KB:202 * KB])
    later = bpb(frozen(learner, novel))
    return {"random_string_first_sight": round(first, 2),
            "seen_once_then_2KB_later": round(soon, 2),
            "seen_once_then_200KB_later": round(later, 2)}


def probe_word_boundaries(trace: list, text: bytes) -> dict:
    starts, insides = [], []
    for index in range(1, len(text)):
        if chr(text[index]).isalpha():
            (insides if chr(text[index - 1]).isalpha() else starts).append(trace[index][0])
    start_bits, inside_bits = sum(starts) / len(starts), sum(insides) / len(insides)
    return {"surprise_at_first_letter_of_word": round(start_bits, 2),
            "surprise_inside_word": round(inside_bits, 2),
            "ratio": round(start_bits / inside_bits, 1)}


def probe_typos(learner, text: bytes, rng: random.Random) -> dict:
    letters = [i for i in range(1, len(text) - 1) if chr(text[i]).islower()]
    typos = sorted(rng.sample(letters, 40))
    corrupted = bytearray(text)
    for index in typos:
        corrupted[index] = rng.choice([c for c in b"abcdefghijklmnopqrstuvwxyz" if c != text[index]])
    trace = frozen(learner, bytes(corrupted))
    flagged_count = len(text) // 50  # flag the 2% most surprising bytes
    flagged = set(sorted(range(len(text)), key=lambda i: -trace[i][0])[:flagged_count])
    caught = sum(any(i in flagged for i in (index, index + 1)) for index in typos)
    return {"typos_planted": len(typos), "caught_in_top_2_percent": caught,
            "caught_percent": round(100 * caught / len(typos), 1), "chance_percent": 4.0}


def probe_calibration(trace: list) -> dict:
    edges = [(0.0, 0.5), (0.5, 0.8), (0.8, 0.95), (0.95, 1.01)]
    rows, error = [], 0.0
    for low, high in edges:
        bucket = [item for item in trace if low <= item[1] < high]
        if not bucket:
            continue
        claimed = sum(item[1] for item in bucket) / len(bucket)
        right = sum(item[2] for item in bucket) / len(bucket)
        error += abs(claimed - right) * len(bucket) / len(trace)
        rows.append({"confidence": f"{low:.2f}-{min(high, 1.0):.2f}", "share_of_bytes": round(len(bucket) / len(trace), 3),
                     "claimed": round(claimed, 3), "actually_right": round(right, 3)})
    return {"buckets": rows, "average_gap": round(error, 4)}


def probe_generation(learner, prompt: bytes, rng: np.random.Generator) -> dict:
    learner = copy.deepcopy(learner)
    feed(learner, prompt)
    produced = bytearray()
    for _ in range(240):
        probs = np.asarray(learner.predict(), dtype=np.float64) ** 1.3  # slightly sharpened
        byte = int(rng.choice(256, p=probs / probs.sum()))
        produced.append(byte)
        learner.update(byte)
    text = produced.decode("latin-1")
    words = [w.strip(".,:;()'\"`").lower() for w in text.split()]
    return {"prompt": prompt.decode("latin-1"), "continuation": text, "word_count": len(words)}


def memory_rows(learner) -> int:
    """Rough size of what the learner holds: entries across its dicts, lists and arrays."""
    total = 0
    for value in vars(learner).values():
        items = value if isinstance(value, (list, tuple)) else [value]
        for item in items:
            if isinstance(item, (dict, bytearray, list)):
                total += len(item)
            elif isinstance(item, np.ndarray):
                total += item.size
    return total


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--learner", default="seeds/sponge.py")
    parser.add_argument("--out", help="also write the results as JSON here")
    args = parser.parse_args()

    factory = load_learner_factory(args.learner)
    python, english = read_streams()
    python_probe = python[600 * KB:600 * KB + PROBE_BYTES]
    english_probe = english[600 * KB:600 * KB + PROBE_BYTES]
    started = time.monotonic()
    report = {"learner": args.learner}

    report["keeps_improving"], python_learner = probe_learning_curve(factory, python)
    report["memory_rows_after_400KB"] = memory_rows(python_learner)
    report["transfer_python_to_english"], english_learner = probe_transfer(factory, python_learner, english)
    report["forgetting"] = probe_forgetting(python_learner, python_probe, english)
    feed(english_learner, english[100 * KB:300 * KB])
    clean_trace = frozen(english_learner, english_probe)
    report["word_boundaries"] = probe_word_boundaries(clean_trace, english_probe)
    report["calibration"] = probe_calibration(clean_trace)
    report["typo_spotting"] = probe_typos(english_learner, english_probe, random.Random(1))
    report["one_shot_memory"] = probe_one_shot(english_learner, english[300 * KB:], random.Random(2))
    report["generation"] = probe_generation(english_learner, b"The interpreter now ", np.random.default_rng(3))
    report["order_of_learning"] = probe_order(factory, python, english, python_probe, english_probe)
    report["minutes"] = round((time.monotonic() - started) / 60, 1)

    print(json.dumps(report, indent=1))
    if args.out:
        Path(args.out).write_text(json.dumps(report, indent=1), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())