#!/usr/bin/env python3
"""Test whether EVA tokens predict illustrated sections without scribe/style leakage.

This is a reproducible first-stage decipherment probe, not a translation system.
Pass a local copy of Ched-ai/voynich-eva/eva_lines.jsonl with --input. The
transcription is research-use material and is deliberately not bundled here.
"""

import argparse
from collections import Counter, defaultdict
import hashlib
import json
import math
from pathlib import Path
import random
import re


SOURCE = "https://huggingface.co/datasets/Ched-ai/voynich-eva"
CONTRASTS = [
    ("hand1_currierA_herbal_vs_pharmaceutical", "1", "A", None, "herbal", "pharmaceutical"),
    ("hand2_currierB_herbal_vs_biological", "2", "B", None, "herbal", "biological"),
    ("hand1_currierA_quireO_herbal_vs_pharmaceutical", "1", "A", "O", "herbal", "pharmaceutical"),
]


def folio_id(page_id):
    match = re.match(r"^(f\d+)", page_id)
    return match.group(1) if match else page_id


def load_pages(path):
    pages = {}
    with path.open(encoding="utf-8") as stream:
        for raw in stream:
            row = json.loads(raw)
            if row.get("line_type") != "paragraph":
                continue
            page_id = row["page_id"]
            if page_id not in pages:
                pages[page_id] = {
                    "page_id": page_id,
                    "folio": folio_id(page_id),
                    "section": row["section"],
                    "hand": str(row["hand"]),
                    "currier": str(row["currier_language"]),
                    "quire": str(row["quire"]),
                    "words": Counter(),
                }
            page = pages[page_id]
            if (page["section"], page["hand"], page["currier"]) != (
                row["section"], str(row["hand"]), str(row["currier_language"])
            ):
                raise ValueError(f"Inconsistent metadata for {page_id}")
            # EVA uses Latin letters as glyph labels, not sounds or plaintext.
            page["words"].update(
                word for word in re.findall(r"[a-z]+", row["text_clean"].lower())
                if 2 <= len(word) <= 20
            )
    return [page for page in pages.values() if page["words"]]


def fold_map(pages, labels, seed):
    by_label = defaultdict(set)
    for page in pages:
        by_label[page["section"]].add(page["folio"])
    assignments = {}
    rng = random.Random(seed)
    fold_count = min(5, *(len(by_label[label]) for label in labels))
    if fold_count < 2:
        raise ValueError("Need at least two folios per label")
    for label in labels:
        folios = sorted(by_label[label])
        rng.shuffle(folios)
        for index, folio in enumerate(folios):
            if folio in assignments:
                raise ValueError(f"Folio {folio} has both contrast labels")
            assignments[folio] = index % fold_count
    return assignments, fold_count


def classify(train, test, labels):
    counts = {label: Counter() for label in labels}
    for page in train:
        counts[page["section"]].update(page["words"])
    global_counts = counts[labels[0]] + counts[labels[1]]
    total_global = sum(global_counts.values())
    alpha = 100.0
    totals = {label: sum(counts[label].values()) for label in labels}
    results = []
    for page in test:
        score = 0.0
        for word, frequency in page["words"].items():
            if word not in global_counts:
                continue
            prior = global_counts[word] / total_global
            probabilities = [
                (counts[label][word] + alpha * prior) / (totals[label] + alpha)
                for label in labels
            ]
            score += frequency * math.log(probabilities[0] / probabilities[1])
        predicted = labels[0] if score >= 0 else labels[1]
        results.append((page["section"], predicted, page["page_id"]))
    return results


def cross_validate(pages, labels, seed):
    folds, fold_count = fold_map(pages, labels, seed)
    results = []
    for fold in range(fold_count):
        train = [p for p in pages if folds[p["folio"]] != fold]
        test = [p for p in pages if folds[p["folio"]] == fold]
        results.extend(classify(train, test, labels))
    recalls = {
        label: sum(actual == label and predicted == label for actual, predicted, _ in results)
        / sum(actual == label for actual, _, _ in results)
        for label in labels
    }
    return round(sum(recalls.values()) / 2, 4), recalls, results


def permutation_scores(pages, labels, count, seed):
    folio_label = {page["folio"]: page["section"] for page in pages}
    folios = sorted(folio_label)
    original_labels = [folio_label[folio] for folio in folios]
    scores = []
    for iteration in range(count):
        rng = random.Random(seed + iteration + 1)
        shuffled = original_labels[:]
        rng.shuffle(shuffled)
        replacement = dict(zip(folios, shuffled))
        null_pages = [dict(page, section=replacement[page["folio"]]) for page in pages]
        score, _, _ = cross_validate(null_pages, labels, seed)
        scores.append(score)
    return scores


def candidate_tokens(pages, labels):
    counts = {label: Counter() for label in labels}
    for page in pages:
        counts[page["section"]].update(page["words"])
    totals = {label: sum(counts[label].values()) for label in labels}
    candidates = {}
    for label, other in (labels, labels[::-1]):
        ranked = []
        for word, frequency in counts[label].items():
            if frequency < 10:
                continue
            rate = (frequency + 1) / (totals[label] + 1)
            other_rate = (counts[other][word] + 1) / (totals[other] + 1)
            ranked.append((math.log(rate / other_rate), word, frequency, counts[other][word]))
        candidates[label] = [
            {"eva_token": word, "count_here": here, "count_other": there,
             "log_rate_ratio": round(score, 3)}
            for score, word, here, there in sorted(ranked, reverse=True)[:10]
        ]
    return candidates


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--permutations", type=int, default=100)
    args = parser.parse_args()
    pages = load_pages(args.input)
    report = {
        "source": SOURCE,
        "input_sha256": hashlib.sha256(args.input.read_bytes()).hexdigest(),
        "method": "up-to-five-fold held-out-folio token-frequency classifier with pooled prior",
        "warning": "Section prediction is not decipherment or translation. EVA tokens are glyph labels.",
        "contrasts": [],
    }
    for name, hand, currier, quire, first, second in CONTRASTS:
        labels = (first, second)
        subset = [
            page for page in pages
            if page["section"] in labels and page["hand"] == hand
            and page["currier"] == currier and (quire is None or page["quire"] == quire)
        ]
        score, recall, results = cross_validate(subset, labels, 408)
        null = permutation_scores(subset, labels, args.permutations, 408)
        report["contrasts"].append({
            "name": name,
            "pages": {label: sum(p["section"] == label for p in subset) for label in labels},
            "folios": {label: len({p["folio"] for p in subset if p["section"] == label}) for label in labels},
            "balanced_accuracy": score,
            "recall": recall,
            "permutation_count": args.permutations,
            "null_mean": round(sum(null) / len(null), 4) if null else None,
            "null_max": max(null) if null else None,
            "permutation_p_value": round((1 + sum(x >= score for x in null)) / (1 + len(null)), 4) if null else None,
            "candidate_tokens": candidate_tokens(subset, labels),
            "predictions": [
                {"page": page, "actual": actual, "predicted": predicted}
                for actual, predicted, page in results
            ],
        })
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    for contrast in report["contrasts"]:
        print(contrast["name"], contrast["pages"],
              "balanced_accuracy", contrast["balanced_accuracy"],
              "null_mean", contrast["null_mean"],
              "p", contrast["permutation_p_value"])
    print(args.output)


if __name__ == "__main__":
    main()
