"""评测专题：Python 3.9+，标准库；自编有限实验，不调用外部模型或服务。"""
from argparse import ArgumentParser
from collections import Counter, defaultdict
from fractions import Fraction
from hashlib import sha256
from itertools import permutations, product
from math import exp, log, log2
from pathlib import Path
import json
import platform
import sqlite3

ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
TARGET = "ABCDEF"
MODES = ("base", "repeat", "all_orders", "alphabet")
POLICIES = ("correct", "says_done", "cancel_all", "wrong_person")


def fingerprint(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def plan():
    return {
        "source_sha256": fingerprint(__file__),
        "alphabet": ALPHABET, "target": TARGET, "smoothing": 1,
        "base": "each ordered pair of A..Z once, with BOS and EOS",
        "extras": {"base": "none", "repeat": "ABCDEF repeated 50 times",
                   "all_orders": "every permutation of ABCDEF once",
                   "alphabet": "AB, BC, ..., YZ each repeated 20 times"},
        "permutation": "all 720 orders; include ties and canonical order",
        "mixture": "fixed illustrative per-group rates, no fitted model",
        "environment": "three fixed tasks, four handwritten policies; in-memory SQLite",
        "compression": "ABAB; known length; binary iid probabilities, no EOS coding",
    }


def corpus(mode):
    data = Counter("".join(pair) for pair in product(ALPHABET, repeat=2))
    if mode == "repeat":
        data[TARGET] += 50
    elif mode == "all_orders":
        data.update("".join(p) for p in permutations(TARGET))
    elif mode == "alphabet":
        data.update({ALPHABET[i:i+2]: 20 for i in range(25)})
    elif mode != "base":
        raise ValueError(mode)
    return data


def fit(data):
    counts = defaultdict(Counter)
    for sequence, repeats in data.items():
        tokens = ["<BOS>"] + list(sequence) + ["<EOS>"]
        for left, right in zip(tokens, tokens[1:]):
            counts[left][right] += repeats
    return counts


def probability(counts, sequence):
    value = Fraction(1)
    tokens = ["<BOS>"] + list(sequence) + ["<EOS>"]
    # 26 个字母加 EOS；BOS 只作条件，不是下一符号候选。
    for left, right in zip(tokens, tokens[1:]):
        row = counts[left]
        value *= Fraction(row[right] + 1, sum(row.values()) + 27)
    return value


def permutation_experiment():
    result = {}
    orders = ["".join(p) for p in permutations(TARGET)]
    for mode in MODES:
        data = corpus(mode)
        counts = fit(data)
        probs = [probability(counts, s) for s in orders]
        canonical = probs[orders.index(TARGET)]
        at_least = sum(p >= canonical for p in probs)
        rows = [{"sequence": s, "numerator": p.numerator,
                 "denominator": p.denominator, "log_probability": log(p)}
                for s, p in zip(orders, probs)]
        result[mode] = {
            "training_sequences": dict(sorted(data.items())),
            "transition_counts": {s: dict(sorted(row.items()))
                                  for s, row in sorted(counts.items())},
            "training_size": sum(data.values()),
            "canonical_training_copies": data[TARGET],
            "canonical_log_probability": log(canonical),
            "at_least_canonical": at_least, "orders": len(orders),
            "p_value": at_least / len(orders), "scores": rows,
        }
    return result


def mixtures():
    rates = {"A": {"easy": .9, "hard": .4},
             "B": {"easy": .8, "hard": .6}}
    return {"group_rates": rates, "comparisons": [
        {"name": name, "easy_weight": weight,
         "scores": {m: weight*r["easy"] + (1-weight)*r["hard"]
                    for m, r in rates.items()}}
        for name, weight in [("old", .8), ("new", .2), ("equal", .5)]]}


def tasks():
    return [
        {"user": user, "target": offset+1,
         "before": [[offset+1, user, "active"],
                    [offset+2, other, "active"],
                    [offset+3, user, "active"],
                    [offset+4, user, "cancelled"]]}
        for offset, user, other in [(100, "alice", "bob"),
                                   (200, "bob", "chen"),
                                   (300, "chen", "alice")]
    ]


def act(connection, task, policy):
    # 故意保留三种错误策略，供评分器比较。所有任务均为本人已确认的取消。
    if policy == "correct":
        connection.execute("UPDATE bookings SET status='cancelled' WHERE id=? AND owner=?",
                           (task["target"], task["user"]))
    elif policy == "cancel_all":
        connection.execute("UPDATE bookings SET status='cancelled' WHERE owner=?",
                           (task["user"],))
    elif policy == "wrong_person":
        connection.execute("UPDATE bookings SET status='cancelled' WHERE id=?",
                           (task["target"]+1,))
    elif policy != "says_done":
        raise ValueError(policy)
    return "已取消指定预约。"


def grade(task, response, after):
    # 独立从任务契约构造期望状态；不读取策略名称或策略内部 SQL。
    expected = [[key, owner, "cancelled" if key == task["target"]
                 and owner == task["user"] else state]
                for key, owner, state in task["before"]]
    return {
        "text_only": "已取消" in response,
        "target_only": any(key == task["target"] and state == "cancelled"
                           for key, _, state in after),
        "full_state": sorted(after) == sorted(expected),
    }


def state_experiment():
    records = []
    for task in tasks():
        for policy in POLICIES:
            with sqlite3.connect(":memory:") as connection:
                connection.execute("CREATE TABLE bookings (id INTEGER PRIMARY KEY, owner TEXT, status TEXT)")
                connection.executemany("INSERT INTO bookings VALUES (?, ?, ?)", task["before"])
                response = act(connection, task, policy)
                after = [list(r) for r in connection.execute("SELECT * FROM bookings ORDER BY id")]
            connection.close()
            records.append({"task": task, "policy": policy, "response": response,
                            "after": after, "grades": grade(task, response, after)})
    return {"records": records, "passes": {
        key: sum(r["grades"][key] for r in records)
        for key in ("text_only", "target_only", "full_state")}}


def compression():
    text = "ABAB"
    result = []
    for units, probabilities in [(list(text), [.5]*4), (["AB", "AB"], [.25]*2)]:
        bits = -sum(log2(p) for p in probabilities)
        ppl = exp(-sum(log(p) for p in probabilities)/len(units))
        result.append({"units": units, "probabilities": probabilities,
                       "bits": bits, "perplexity": ppl,
                       "utf8_bytes": len(text.encode("utf-8")),
                       "rate_percent": 100*bits/(8*len(text.encode("utf-8")))})
    return result


def run():
    return {"source_sha256": fingerprint(__file__),
            "environment": {"python": platform.python_version(), "sqlite": sqlite3.sqlite_version},
            "permutation": permutation_experiment(), "mixtures": mixtures(),
            "states": state_experiment(), "compression": compression()}


def show(record):
    print("题目混合：")
    for row in record["mixtures"]["comparisons"]:
        print(f'{row["name"]}: A={row["scores"]["A"]:.1%}, B={row["scores"]["B"]:.1%}')
    print("全部 720 种排列（包含并列）：")
    for mode, row in record["permutation"].items():
        print(f'{mode}: 完整原序列出现 {row["canonical_training_copies"]} 次, '
              f'p={row["at_least_canonical"]}/{row["orders"]}={row["p_value"]:.6f}')
    print("实际状态评分（12 次执行）：")
    for key, count in record["states"]["passes"].items():
        print(f"{key}: {count}/12")
    for row in record["compression"]:
        print(f'{len(row["units"])} 个词元: PPL={row["perplexity"]:.1f}, '
              f'理想码长={row["bits"]:.1f} bit, 占 UTF-8={row["rate_percent"]:.1f}%')


def save(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2)+"\n", encoding="utf-8")


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--plan", type=Path)
    parser.add_argument("--record", type=Path)
    args = parser.parse_args()
    if args.plan:
        save(args.plan, plan())
        return
    record = run()
    if args.record:
        save(args.record, record)
    show(record)


if __name__ == "__main__":
    main()
