"""Recompute the published model comparisons from text-free evaluation records.

Python 3.11+, standard library only. Run: python metrics.py --directory .
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from statistics import mean

CUTOFFS = (1, 3, 5, 10, 20, 50)
FACETS = ("trap", "decision", "constraint", "mechanism", "procedure")


def classification_metrics(labels: list[int], scores: list[float]) -> dict[str, float]:
    """Threshold-grouped AP and best precision at or above each recall target."""
    if len(labels) != len(scores) or not labels or set(labels) - {0, 1}:
        raise ValueError("Classification labels and scores must be aligned binary rows")
    if not all(math.isfinite(s) for s in scores) or sum(labels) == 0:
        raise ValueError("Classification scores must be finite with positive support")
    order = sorted(range(len(labels)), key=lambda i: -scores[i])
    positives = sum(labels)
    true_positives = 0
    previous_recall = 0.0
    ap = 0.0
    points = []
    for rank, index in enumerate(order, 1):
        true_positives += labels[index]
        if rank < len(order) and scores[order[rank]] == scores[index]:
            continue
        precision = true_positives / rank
        recall = true_positives / positives
        ap += (recall - previous_recall) * precision
        previous_recall = recall
        points.append((precision, recall))
    return {
        "average_precision": ap,
        "precision_at_recall_90": max(p for p, r in points if r >= 0.9),
        "precision_at_recall_95": max(p for p, r in points if r >= 0.95),
        "prevalence": positives / len(labels),
    }


def dcg(grades: list[int], k: int) -> float:
    return sum((2**g - 1) / math.log2(i + 2) for i, g in enumerate(grades[:k]))


def ranking_metrics(grades: list[int], scores: list[float], k: int) -> dict[str, float]:
    if len(grades) != len(scores) or len(grades) < k or k < 1:
        raise ValueError("Ranking inputs must be aligned and cover the cutoff")
    if any(type(g) is not int or g not in range(4) for g in grades):
        raise ValueError("Ranking requires fully judged integer grades 0 through 3")
    if not all(math.isfinite(s) for s in scores):
        raise ValueError("Ranking scores must be finite")
    order = sorted(range(len(scores)), key=lambda i: (-scores[i], i))
    ranked = [grades[i] for i in order]
    useful = sum(g >= 2 for g in ranked[:k])
    relevant = sum(g >= 2 for g in grades)
    ideal = dcg(sorted(grades, reverse=True), k)
    result = {"hit": float(useful > 0), "precision": useful / k, "useful": float(useful)}
    if relevant:
        result["recall"] = useful / relevant
        result["ndcg"] = dcg(ranked, k) / ideal
    return result


def random_ranking_metrics(grades: list[int], k: int) -> dict[str, float]:
    """Exact expectation under a uniform permutation of one fixed judged pool."""
    n = len(grades)
    if k < 1 or k > n or any(type(g) is not int or g not in range(4) for g in grades):
        raise ValueError("Random reference requires judged grades and a valid cutoff")
    relevant = sum(g >= 2 for g in grades)
    misses = math.comb(n - relevant, k) if n - relevant >= k else 0
    result = {"hit": 1 - misses / math.comb(n, k), "precision": relevant / n, "useful": k * relevant / n}
    if relevant:
        expected_dcg = mean(2**g - 1 for g in grades) * sum(1 / math.log2(i + 2) for i in range(k))
        result["recall"] = k / n
        result["ndcg"] = expected_dcg / dcg(sorted(grades, reverse=True), k)
    return result


def summarize_classifier(data: dict) -> dict:
    rows = data["rows"]
    result = {"rows": len(rows), "families": len({r["group"] for r in rows}), "models": {}}
    for model in data["model_order"]:
        by_facet = {}
        for facet in FACETS:
            resolved = [r for r in rows if r["labels"][facet] is not None]
            if any(r["labels"][facet] not in (0, 1) for r in resolved):
                raise ValueError("Invalid resolved classifier label")
            by_facet[facet] = {"rows": len(resolved), **classification_metrics(
                [r["labels"][facet] for r in resolved], [r["scores"][model][facet] for r in resolved]
            )}
        macro = {k: mean(v[k] for v in by_facet.values()) for k in ("average_precision", "precision_at_recall_90", "precision_at_recall_95")}
        result["models"][model] = {"facets": by_facet, "macro": macro}
    return result


def summarize_reranker(data: dict) -> dict:
    rows = data["rows"]
    answerable = [r for r in rows if any(g >= 2 for g in r["grades"])]
    output = {"queries": len(rows), "answerable": len(answerable), "models": {}}
    for model in ["random", *data["model_order"]]:
        scopes = {}
        for scope, subset in [("answerable", answerable), ("all", rows)]:
            cuts = {}
            for k in CUTOFFS:
                values = [random_ranking_metrics(r["grades"], k) if model == "random" else ranking_metrics(r["grades"], r["scores"][model], k) for r in subset]
                fields = ("hit", "precision", "useful", "recall", "ndcg") if scope == "answerable" else ("hit", "precision", "useful")
                cuts[str(k)] = {field: mean(v[field] for v in values) for field in fields}
            scopes[scope] = cuts
        output["models"][model] = scopes
    return output


def summarize_retriever(data: dict) -> dict:
    """Dense rankings, with both full-panel and known-answerable summaries."""
    if type(data['corpus_passages']) is not int or data['corpus_passages'] <= 0:
        raise ValueError('Dense comparison requires a positive corpus size')
    return {**_summarize_ranked_grades(data, include_selected=False),
            'answerable': _summarize_ranked_grades(data, include_selected=False, answerable_only=True),
            'corpus_passages': data['corpus_passages']}


def summarize_dimensions(data: dict) -> dict:
    """Average the measured public nDCG values for each embedding width."""
    rows, widths = data['rows'], data['widths']
    if not rows or len({row['id'] for row in rows}) != len(rows):
        raise ValueError('Width comparison requires distinct query identities')
    if not widths or len(set(widths)) != len(widths) or any(type(w) is not int or w <= 0 for w in widths):
        raise ValueError('Embedding widths must be distinct positive integers')
    for row in rows:
        if set(row['ndcg@10']) != {str(w) for w in widths} or any(
            isinstance(v, bool) or not isinstance(v, int | float) or not math.isfinite(v) or not 0 <= v <= 1
            for v in row['ndcg@10'].values()
        ):
            raise ValueError('Every query must have a finite nDCG value for every width')
    panels = {}
    for name in sorted({row['dataset'] for row in rows}):
        panel = [row for row in rows if row['dataset'] == name]
        panels[name] = {'queries': len(panel), 'ndcg@10': {
            str(w): mean(row['ndcg@10'][str(w)] for row in panel) for w in widths
        }}
    result = {'queries': len(rows), 'vector_bytes_fp32': {str(w): w * 4 for w in widths}, 'public': panels}
    if 'private' in data:
        private = data['private']
        if private['model_order'] != [str(w) for w in widths]:
            raise ValueError('Private width rankings must cover the same ordered widths')
        result['private'] = summarize_record('retriever', private)
    return result


def _summarize_ranked_grades(data: dict, *, include_selected: bool, answerable_only: bool = False) -> dict:
    """Reviewed exclusions inside each original dense or hybrid prefix.

    A null is permitted only for an explicitly excluded judging abstention.
    Cut first, remove exclusions second, and never backfill from a deeper rank.
    """
    rows = data['rows']
    if not rows or len({row['id'] for row in rows}) != len(rows):
        raise ValueError('Ranked rows require unique nonempty query identities')
    # Published promotion and serving v1 records contain twenty positions.
    # Extended dense/serving records declare their measured depth explicitly.
    depth_limit = data.get('ranking_depth', 20)
    allowed_depths = (20, 50) if include_selected else (10, 20, 50)
    if type(depth_limit) is not int or depth_limit not in allowed_depths:
        raise ValueError('Unsupported measured ranking depth')
    # Validate before excluding fully abstained prefixes. An omitted query
    # must never conceal a malformed grade, model map or selected depth.
    for row in rows:
        counts = row['reference_grade_counts']
        if set(counts) != {'0', '1', '2', '3'} or any(type(n) is not int or n < 0 for n in counts.values()):
            raise ValueError('Reference grade counts must cover grades zero through three')
        for model in data['model_order']:
            ranked, excluded = row['ranked_grades'][model], row['excluded'][model]
            if (len(ranked) != depth_limit or len(excluded) != depth_limit
                    or any(type(x) is not bool for x in excluded)):
                raise ValueError('Ranked rows must cover exactly the declared ranking depth')
            if include_selected:
                depth = row['selected_depth'][model]
                if type(depth) is not int or not 3 <= depth <= 20:
                    raise ValueError('Selected depth must be between three and twenty')
            if any((grade is not None if drop else type(grade) is not int or grade not in range(4))
                   for grade, drop in zip(ranked, excluded, strict=True)):
                raise ValueError('Only declared abstentions may lack grades')
    # The shared relevance reference defines answerability, never one arm's
    # retrieved prefix. Validate every source row before applying this filter.
    if answerable_only:
        rows = [row for row in rows if row['reference_grade_counts']['2'] + row['reference_grade_counts']['3'] > 0]
    output = {'queries': len(rows), 'models': {}}
    if not rows:
        return output  # No conditional score exists; do not substitute zeros.
    requested_cutoffs = tuple(k for k in CUTOFFS if k <= depth_limit)
    if include_selected:
        requested_cutoffs += ('selected',)
    for model in data['model_order']:
        cutoffs = {}
        for cutoff in requested_cutoffs:
            eligible = [row for row in rows if all(
                any(not drop for drop in row['excluded'][arm][:(
                    row['selected_depth'][arm] if cutoff == 'selected' else cutoff)])
                for arm in data['model_order']
            )]
            if not eligible:
                raise ValueError('No shared judged queries at this cutoff')
            per_query = []
            for row in eligible:
                counts = row['reference_grade_counts']
                ranked = row['ranked_grades'][model]
                excluded = row['excluded'][model]
                depth = row['selected_depth'][model] if cutoff == 'selected' else cutoff
                kept = [grade for grade, drop in zip(ranked[:depth], excluded[:depth], strict=True) if not drop]
                ideal_grades = [g for g in (3, 2, 1, 0) for _ in range(min(counts[str(g)], len(kept)))][:len(kept)]
                useful = sum(g >= 2 for g in kept)
                positives = counts['2'] + counts['3']
                ideal = dcg(ideal_grades, len(kept))
                per_query.append({
                    'hit': float(useful > 0), 'precision': useful / len(kept),
                    'useful': useful, 'retained': len(kept), 'excluded': depth - len(kept),
                    'ndcg': dcg(kept, len(kept)) / ideal if ideal else 0.0,
                    'known_recall': useful / positives if positives else None,
                })
            useful = sum(row['useful'] for row in per_query)
            retained = sum(row['retained'] for row in per_query)
            recalls = [row['known_recall'] for row in per_query if row['known_recall'] is not None]
            cutoffs[str(cutoff)] = {
                'scored_queries': len(eligible), 'excluded_queries': len(rows) - len(eligible),
                'hit': mean(row['hit'] for row in per_query), 'precision': useful / retained,
                'macro_precision': mean(row['precision'] for row in per_query),
                'ndcg': mean(row['ndcg'] for row in per_query),
                'known_positive_recall': mean(recalls) if recalls else None,
                'recall_queries': len(recalls), 'useful': useful, 'retained': retained,
                'excluded_positions': sum(row['excluded'] for row in per_query),
                'mean_useful': useful / len(eligible), 'mean_retained': retained / len(eligible),
            }
        output['models'][model] = cutoffs
    return output


def summarize_promotion(data: dict) -> dict:
    """Hybrid rankings and the separately measured public dense panels."""
    output = {**_summarize_ranked_grades(data, include_selected=True), 'public': {}}
    for dataset, panel in data['public'].items():
        if not panel['rows'] or len({row['id'] for row in panel['rows']}) != len(panel['rows']):
            raise ValueError('Public panel requires distinct query identities')
        measured = panel['model_order']
        for row in panel['rows']:
            if set(row['ndcg@10']) != set(measured) or any(
                isinstance(value, bool) or not isinstance(value, int | float)
                or not math.isfinite(value) or not 0 <= value <= 1
                for value in row['ndcg@10'].values()
            ):
                raise ValueError('Public nDCG values must be finite measured scores')
        output['public'][dataset] = {
            'queries': len(panel['rows']),
            'ndcg@10': {model: mean(row['ndcg@10'][model] for row in panel['rows']) for model in measured},
        }
    return output


def summarize_serving(data: dict) -> dict:
    """Current hybrid quality uses the same reference-defined eligible queries."""
    return {**summarize_promotion(data),
            'answerable': _summarize_ranked_grades(data, include_selected=True, answerable_only=True)}


def summarize_pipeline(data: dict) -> dict:
    """Matched component changes on a frozen, known-answerable query cohort."""
    if type(data['corpus_passages']) is not int or data['corpus_passages'] <= 0:
        raise ValueError('Pipeline comparison requires a positive corpus size')
    output = _summarize_ranked_grades(data, include_selected=False)
    if any(row['reference_grade_counts']['2'] + row['reference_grade_counts']['3'] == 0
           for row in data['rows']):
        raise ValueError('Pipeline cohort requires known useful evidence for every query')
    return {**output, 'corpus_passages': data['corpus_passages'],
            'reference_pairs': sum(sum(row['reference_grade_counts'].values()) for row in data['rows'])}


SUMMARIZERS = {
    'classifier': summarize_classifier, 'reranker': summarize_reranker,
    'retriever': summarize_retriever, 'dimensions': summarize_dimensions,
    'promotion': summarize_promotion, 'serving': summarize_serving, 'pipeline': summarize_pipeline,
}


def summarize_record(name: str, data: dict) -> dict:
    """Check anonymous publication records before recomputing their summary."""
    fields = {
        'classifier': {'id', 'group', 'labels', 'scores'},
        'reranker': {'id', 'group', 'grades', 'scores'},
        'retriever': {'id', 'group', 'reference_grade_counts', 'ranked_grades', 'excluded'},
        'dimensions': {'id', 'dataset', 'ndcg@10'},
        'promotion': {'id', 'group', 'reference_grade_counts', 'ranked_grades', 'excluded', 'selected_depth'},
        'serving': {'id', 'group', 'reference_grade_counts', 'ranked_grades', 'excluded', 'selected_depth'},
        'pipeline': {'id', 'group', 'reference_grade_counts', 'ranked_grades', 'excluded'},
    }[name]
    rows = data['rows']
    if not rows or any(set(row) != fields for row in rows):
        raise ValueError(f'Unexpected evaluation row fields in {name}')
    ids = [row['id'] for row in rows]
    if len(set(ids)) != len(ids):
        raise ValueError(f'Repeated query/passage identity in {name}')
    for row in rows:
        identity = row['id']
        if not isinstance(identity, str) or identity[:1] not in ('p', 'q') or not identity[1:].isdigit():
            raise ValueError(f'Non-anonymous row identity in {name}')
        if 'group' in row and (not isinstance(row['group'], str) or not row['group'].startswith('g') or not row['group'][1:].isdigit()):
            raise ValueError(f'Non-anonymous group identity in {name}')
    if name != 'dimensions' and set(data['model_order']) != set(data['models']):
        raise ValueError(f'Model identities do not match in {name}')
    models = set(data.get('model_order', ()))
    for row in rows:
        for field in ('scores', 'ranked_grades', 'excluded', 'selected_depth'):
            if field in row and set(row[field]) != models:
                raise ValueError(f'Unexpected model fields in {name}.{field}')
        if name == 'classifier':
            required = {facet for facet, label in row['labels'].items() if label is not None}
            if set(row['labels']) != set(FACETS) or any(not required <= set(scores) <= set(FACETS) for scores in row['scores'].values()):
                raise ValueError('Unexpected classifier facet fields')
            if any(type(value) not in (int, float) or not math.isfinite(value)
                   for scores in row['scores'].values() for value in scores.values()):
                raise ValueError('Classifier scores must be finite numbers, including unresolved facets')
    for panel in data.get('public', {}).values():
        if set(panel) != {'model_order', 'rows'} or not panel['rows']:
            raise ValueError('Unexpected public-panel fields')
        panel_models = set(panel['model_order'])
        panel_ids = []
        for row in panel['rows']:
            if set(row) != {'id', 'ndcg@10'} or set(row['ndcg@10']) != panel_models:
                raise ValueError('Unexpected public-panel row fields')
            identity = row['id']
            if not isinstance(identity, str) or identity[:1] != 'q' or not identity[1:].isdigit():
                raise ValueError('Non-anonymous public-panel row identity')
            if any(type(value) not in (int, float) or not math.isfinite(value) or not 0 <= value <= 1 for value in row['ndcg@10'].values()):
                raise ValueError('Invalid public-panel score')
            panel_ids.append(identity)
        if len(panel_ids) != len(set(panel_ids)):
            raise ValueError('Repeated public-panel query identity')
    text = json.dumps(data).lower()
    if any(value in text for value in ('c:\\', 'eval/private', 'qrel_target_id', 'query_text', 'passage_text')):
        raise ValueError(f'Private payload in {name}')
    return SUMMARIZERS[name](data)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--directory", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--model", choices=tuple(SUMMARIZERS), help="Recompute one comparison; default: every record in this package")
    args = parser.parse_args()
    names = [args.model] if args.model else [
        name for name in SUMMARIZERS if (args.directory / f"{name}.json").is_file()
    ]
    if not names:
        parser.error("No evaluation records found in the selected directory")
    result = {}
    for name in names:
        data = json.loads((args.directory / f"{name}.json").read_text(encoding="utf-8"))
        result[name] = summarize_record(name, data)
    text = json.dumps(result, indent=2, allow_nan=False) + "\n"
    if args.output:
        args.output.write_text(text, encoding="utf-8")
    else:
        print(text, end="")


if __name__ == "__main__":
    main()