#!/usr/bin/env python3
"""Dependency-free protocol smoke test for reranking and hard negatives.

This is deliberately a transparent linear pair scorer, not a Transformer and
not a reproduction of MS MARCO numbers. It verifies the data contract, binary
cross-entropy, candidate-only reranking, MRR/NDCG, and the recall ceiling.
"""

from __future__ import annotations

import argparse
import math
import random
import re
from dataclasses import dataclass


SEED = 58


@dataclass(frozen=True)
class Case:
    query: str
    positive: str
    hard: str
    easy: str


CASES = [
    Case(
        "where was marie curie born",
        "marie curie was born in warsaw",
        "marie curie was born in paris and studied physics",
        "plants absorb carbon dioxide during photosynthesis",
    ),
    Case(
        "who wrote pride and prejudice",
        "jane austen wrote pride and prejudice",
        "charles dickens wrote pride and prejudice adaptations",
        "tokyo is the capital of japan",
    ),
    Case(
        "what is the capital of japan",
        "tokyo is the capital of japan",
        "kyoto was the capital of japan for centuries",
        "leonardo da vinci painted the mona lisa",
    ),
    Case(
        "which gas do plants absorb",
        "plants absorb carbon dioxide during photosynthesis",
        "plants absorb oxygen during respiration",
        "jane austen wrote pride and prejudice",
    ),
]


def tokens(text: str) -> list[str]:
    return re.findall(r"[a-z]+", text.lower())


def make_vocab(cases: list[Case]) -> list[str]:
    texts = [value for case in cases for value in (case.positive, case.hard, case.easy)]
    return sorted({token for text in texts for token in tokens(text)})


def pair_features(query: str, passage: str, vocab: list[str]) -> list[float]:
    q = set(tokens(query))
    p = set(tokens(passage))
    overlap = len(q & p) / max(len(q), 1)
    # Passage unigrams plus a query-passage interaction feature make both
    # inputs observable to one scorer. A real Cross-Encoder learns far richer
    # token interactions through self-attention.
    return [1.0, overlap] + [float(token in p) for token in vocab]


def sigmoid(value: float) -> float:
    if value >= 0:
        return 1.0 / (1.0 + math.exp(-value))
    exp_value = math.exp(value)
    return exp_value / (1.0 + exp_value)


def score(weights: list[float], features: list[float]) -> float:
    return sum(weight * feature for weight, feature in zip(weights, features))


def binary_cross_entropy(logit: float, label: int) -> float:
    # Stable BCE-with-logits expression.
    return max(logit, 0.0) - logit * label + math.log1p(math.exp(-abs(logit)))


def examples(cases: list[Case], include_hard: bool) -> list[tuple[str, str, int, str]]:
    rows = []
    for case in cases:
        rows.append((case.query, case.positive, 1, "positive"))
        rows.append((case.query, case.easy, 0, "easy"))
        if include_hard:
            rows.append((case.query, case.hard, 0, "hard"))
    return rows


def train(cases: list[Case], vocab: list[str], include_hard: bool, steps: int) -> tuple[list[float], float, float]:
    rng = random.Random(SEED)
    weights = [rng.uniform(-0.01, 0.01) for _ in range(2 + len(vocab))]
    rows = examples(cases, include_hard)

    def mean_loss() -> float:
        return sum(
            binary_cross_entropy(score(weights, pair_features(query, passage, vocab)), label)
            for query, passage, label, _ in rows
        ) / len(rows)

    initial_loss = mean_loss()
    learning_rate = 0.08
    for _ in range(steps):
        rng.shuffle(rows)
        for query, passage, label, _ in rows:
            features = pair_features(query, passage, vocab)
            prediction = sigmoid(score(weights, features))
            for i, feature in enumerate(features):
                weights[i] -= learning_rate * (prediction - label) * feature
    return weights, initial_loss, mean_loss()


def reciprocal_rank(rank: int | None) -> float:
    return 0.0 if rank is None else 1.0 / rank


def ndcg_at_3(rank: int | None) -> float:
    return 0.0 if rank is None or rank > 3 else 1.0 / math.log2(rank + 1)


def evaluate(weights: list[float], vocab: list[str], drop_last_positive: bool) -> dict[str, float]:
    ranks: list[int | None] = []
    for index, case in enumerate(CASES):
        candidates = [("hard", case.hard), ("positive", case.positive), ("easy", case.easy)]
        if drop_last_positive and index == len(CASES) - 1:
            candidates = [(name, text) for name, text in candidates if name != "positive"]
        ranked = sorted(
            candidates,
            key=lambda item: score(weights, pair_features(case.query, item[1], vocab)),
            reverse=True,
        )
        labels = [name for name, _ in ranked]
        rank = labels.index("positive") + 1 if "positive" in labels else None
        ranks.append(rank)
        printable_rank = "missing" if rank is None else str(rank)
        print(f"q{index}: rank={printable_rank} order={','.join(labels)}")

    n = len(ranks)
    return {
        "candidate_recall@3": sum(rank is not None for rank in ranks) / n,
        "mrr": sum(reciprocal_rank(rank) for rank in ranks) / n,
        "ndcg@3": sum(ndcg_at_3(rank) for rank in ranks) / n,
    }


def run(steps: int, check_only: bool) -> None:
    vocab = make_vocab(CASES)
    easy_weights, easy_start, easy_end = train(CASES, vocab, include_hard=False, steps=steps)
    hard_weights, hard_start, hard_end = train(CASES, vocab, include_hard=True, steps=steps)

    print(f"seed={SEED} vocab={len(vocab)} pair_features=(batch, {len(vocab) + 2})")
    print(f"easy-only rows={len(examples(CASES, False))} loss={easy_start:.4f}->{easy_end:.4f}")
    print(f"mixed rows={len(examples(CASES, True))} loss={hard_start:.4f}->{hard_end:.4f}")
    print("easy-only reranking:")
    easy_metrics = evaluate(easy_weights, vocab, drop_last_positive=False)
    print("hard-negative reranking:")
    hard_metrics = evaluate(hard_weights, vocab, drop_last_positive=False)
    print("hard-negative reranking with one missing positive:")
    ceiling_metrics = evaluate(hard_weights, vocab, drop_last_positive=True)
    for name, values in (("easy", easy_metrics), ("hard", hard_metrics), ("ceiling", ceiling_metrics)):
        print(name + ": " + " ".join(f"{key}={value:.3f}" for key, value in values.items()))

    assert len(pair_features(CASES[0].query, CASES[0].positive, vocab)) == len(vocab) + 2
    assert easy_end < easy_start and hard_end < hard_start
    assert hard_metrics["mrr"] > easy_metrics["mrr"]
    assert hard_metrics["candidate_recall@3"] == 1.0
    assert ceiling_metrics["candidate_recall@3"] == 0.75
    assert ceiling_metrics["mrr"] <= ceiling_metrics["candidate_recall@3"]
    if check_only:
        print("check-only: PASS")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=180)
    parser.add_argument("--check-only", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    run(steps=args.steps, check_only=args.check_only)