#!/usr/bin/env python3
"""Answer-retention proxy for SOMA CoT compressors.

The competition score combines token reduction with tasks passed. Compression
ratio alone therefore says nothing about whether a compressor is any good. The
official check requires an LLM answering questions from the compressed context;
this is a cheaper LEXICAL PROXY and is explicitly not the official metric.

Method: for each reference answer, extract its content terms (alphanumeric
tokens >=4 chars, minus stopwords, plus any dotted/underscored identifiers and
quoted literals). Retention for one QA pair is the fraction of those terms that
still appear in the compressed context. A term that never appeared in the
UNCOMPRESSED context is excluded, since no compressor could have kept it -- this
makes the score measure compression damage rather than corpus coverage.
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import re
import statistics
from pathlib import Path

PLUGIN = Path("/var/lib/octave/sn114/external/SOMA-plugin")
COT = Path(
    "/var/lib/octave/sn114/repo/miner/plain_text_compression/sample_tasks/"
    "cot_compression_tasks.jsonl"
)
COMPRESSORS = {
    "performance": "performance_focused_compressor.py",
    "aggressive": "aggressive_hybrid_compressor.py",
    "code_focused": "code_focused_compressor.py",
    "hybrid": "hybrid_compressor.py",
    "adaptive": "adaptive_compressor.py",
    "thinking_strip": "thinking_strip_compressor.py",
    "dialogue": "dialogue_focused_compressor.py",
    "structural": "structural_cot_compressor.py",
    "structural_reasoning": "structural_reasoning_compressor.py",
    "adaptive_structural": "adaptive_structural_compressor.py",
    "structural_lowcomp": "structural_lowcomp_compressor.py",
    "structural_selective": "structural_selective_compressor.py",
    "structural_selective_call": "structural_selective_call_compressor.py",
    "structural_pathkeep": "structural_pathkeep_compressor.py",
    "structural_threshold": "structural_threshold_compressor.py",
    "structural_perspan": "structural_perspan_compressor.py",
    "structural_forward": "structural_forward_compressor.py",
    "structural_combined": "structural_combined_compressor.py",
    "structural_gist": "structural_gist_compressor.py",
}
STOP = {
    "that","this","with","from","were","have","been","which","when","what","the",
    "and","for","was","are","not","but","its","their","there","then","than","into",
    "onto","after","before","because","would","could","should","also","such","only",
    "user","wanted","issue","problem","code","function","file","line","fix","fixed",
    "correctly","incorrectly","caused","causing","result","resulting","string",
}


def load(name: str, filename: str):
    spec = importlib.util.spec_from_file_location(f"r_{name}", PLUGIN / filename)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod.compress_messages


def content_terms(answer: str) -> set[str]:
    terms: set[str] = set()
    # dotted / underscored identifiers and paths carry the most signal
    for ident in re.findall(r"[A-Za-z_][\w./]*[\w]", answer):
        if ("." in ident or "_" in ident or "/" in ident) and len(ident) >= 4:
            terms.add(ident.lower())
    for lit in re.findall(r"'([^']{2,40})'|\"([^\"]{2,40})\"", answer):
        val = (lit[0] or lit[1]).strip().lower()
        if len(val) >= 3:
            terms.add(val)
    for w in re.findall(r"[A-Za-z][A-Za-z0-9]{3,}", answer):
        lw = w.lower()
        if lw not in STOP:
            terms.add(lw)
    return terms


def retention(context: str, compressed: str, qa: list[dict]) -> dict:
    ctx = context.lower()
    comp = compressed.lower()
    per_q, kept_t, total_t = [], 0, 0
    for pair in qa:
        terms = content_terms(pair.get("answer", ""))
        # only terms the uncompressed context could have supplied
        avail = {t for t in terms if t in ctx}
        if not avail:
            continue
        kept = sum(1 for t in avail if t in comp)
        per_q.append(kept / len(avail))
        kept_t += kept
        total_t += len(avail)
    return {
        "questions_scored": len(per_q),
        "mean_per_question_retention": statistics.fmean(per_q) if per_q else 0.0,
        "min_per_question_retention": min(per_q) if per_q else 0.0,
        "micro_term_retention": kept_t / total_t if total_t else 0.0,
        "terms_available": total_t,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, required=True)
    args = ap.parse_args()

    tasks = [json.loads(l) for l in COT.read_text().splitlines() if l.strip()]
    results = {}

    for name, filename in COMPRESSORS.items():
        fn = load(name, filename)
        rows, ratios = [], []
        for task in tasks:
            src = task["source_text"]
            qa = task.get("qa") or []
            if not qa:
                continue
            out = fn([{"role": "user", "content": src}])
            comp = out[0]["content"] if out else ""
            rows.append(retention(src, comp, qa))
            ratios.append(len(comp) / len(src))

        comp_pct = (1 - statistics.fmean(ratios)) * 100
        mean_ret = statistics.fmean(r["mean_per_question_retention"] for r in rows)
        micro = statistics.fmean(r["micro_term_retention"] for r in rows)
        worst = min(r["min_per_question_retention"] for r in rows)
        results[name] = {
            "compression_percent": round(comp_pct, 3),
            "mean_answer_retention": round(mean_ret, 4),
            "micro_term_retention": round(micro, 4),
            "worst_question_retention": round(worst, 4),
            "tasks": len(rows),
            "questions_scored": sum(r["questions_scored"] for r in rows),
            # joint view: retention weighted by how much was removed
            "retention_x_compression": round(mean_ret * (comp_pct / 100), 4),
        }
        print(f"{name:<16} comp={comp_pct:6.2f}%  retention={mean_ret:.4f}  "
              f"worst_q={worst:.4f}")

    args.out.write_text(json.dumps(
        {"metric": "lexical answer-retention proxy (NOT the official LLM-judged score)",
         "results": results}, indent=2))
    print(f"\nwrote {args.out}")


if __name__ == "__main__":
    main()