#!/usr/bin/env python3
"""Dependency-free toy reproduction of dual-tower dense retrieval.

This is a protocol check, not a reproduction of DPR/Contriever benchmark
numbers. It makes the score matrix, in-batch contrastive loss, offline passage
index, online top-k search, and retrieval metrics directly inspectable.
"""

from __future__ import annotations

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


SEED = 57
DIM = 24


@dataclass(frozen=True)
class Example:
    query: str
    passage: str


TRAIN = [
    Example("where was marie curie born", "marie curie was born in warsaw poland"),
    Example("who wrote pride and prejudice", "jane austen wrote pride and prejudice"),
    Example("what planet is called the red planet", "mars is known as the red planet"),
    Example("which gas do plants absorb", "plants absorb carbon dioxide during photosynthesis"),
    Example("what is the capital of japan", "tokyo is the capital city of japan"),
    Example("who painted the mona lisa", "leonardo da vinci painted the mona lisa"),
]

CORPUS = [ex.passage for ex in TRAIN] + [
    "pierre curie was born in paris france",
    "venus is sometimes called earths sister planet",
    "plants release oxygen as a product of photosynthesis",
    "kyoto was an imperial capital of japan",
]
QRELS = {i: i for i in range(len(TRAIN))}


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


def build_vocab(texts: list[str]) -> list[str]:
    return sorted({token for text in texts for token in tokenize(text)})


def init_table(vocab: list[str], rng: random.Random) -> dict[str, list[float]]:
    return {token: [rng.uniform(-0.08, 0.08) for _ in range(DIM)] for token in vocab}


def mean_vector(text: str, table: dict[str, list[float]]) -> list[float]:
    tokens = tokenize(text)
    return [sum(table[token][j] for token in tokens) / len(tokens) for j in range(DIM)]


def dot(left: list[float], right: list[float]) -> float:
    return sum(a * b for a, b in zip(left, right))


def softmax(row: list[float]) -> list[float]:
    offset = max(row)
    exps = [math.exp(value - offset) for value in row]
    total = sum(exps)
    return [value / total for value in exps]


def score_matrix(queries: list[list[float]], passages: list[list[float]], temperature: float) -> list[list[float]]:
    return [[dot(query, passage) / temperature for passage in passages] for query in queries]


def contrastive_loss(scores: list[list[float]]) -> float:
    return -sum(math.log(max(softmax(row)[i], 1e-12)) for i, row in enumerate(scores)) / len(scores)


def train_step(
    query_table: dict[str, list[float]],
    passage_table: dict[str, list[float]],
    learning_rate: float,
    temperature: float,
) -> float:
    queries = [mean_vector(ex.query, query_table) for ex in TRAIN]
    passages = [mean_vector(ex.passage, passage_table) for ex in TRAIN]
    scores = score_matrix(queries, passages, temperature)
    batch = len(TRAIN)
    d_scores = []
    for i, row in enumerate(scores):
        probs = softmax(row)
        probs[i] -= 1.0
        d_scores.append([value / batch for value in probs])

    d_queries = [
        [sum(d_scores[i][k] * passages[k][j] for k in range(batch)) / temperature for j in range(DIM)]
        for i in range(batch)
    ]
    d_passages = [
        [sum(d_scores[i][k] * queries[i][j] for i in range(batch)) / temperature for j in range(DIM)]
        for k in range(batch)
    ]

    for i, ex in enumerate(TRAIN):
        query_tokens = tokenize(ex.query)
        passage_tokens = tokenize(ex.passage)
        for token in query_tokens:
            for j in range(DIM):
                query_table[token][j] -= learning_rate * d_queries[i][j] / len(query_tokens)
        for token in passage_tokens:
            for j in range(DIM):
                passage_table[token][j] -= learning_rate * d_passages[i][j] / len(passage_tokens)
    return contrastive_loss(scores)


def retrieval_metrics(scores: list[list[float]], qrels: dict[int, int]) -> dict[str, float]:
    reciprocal_ranks = []
    hit_at_1 = 0
    hit_at_3 = 0
    for query_id, positive_id in qrels.items():
        ranking = sorted(range(len(scores[query_id])), key=lambda i: scores[query_id][i], reverse=True)
        rank = ranking.index(positive_id) + 1
        reciprocal_ranks.append(1.0 / rank)
        hit_at_1 += int(rank <= 1)
        hit_at_3 += int(rank <= 3)
    n = len(qrels)
    return {"recall@1": hit_at_1 / n, "recall@3": hit_at_3 / n, "mrr": sum(reciprocal_ranks) / n}


def run(steps: int, check_only: bool) -> None:
    rng = random.Random(SEED)
    vocab = build_vocab([ex.query for ex in TRAIN] + CORPUS)
    query_table = init_table(vocab, rng)
    passage_table = init_table(vocab, rng)  # independent DPR-style tower
    temperature = 0.10

    initial_queries = [mean_vector(ex.query, query_table) for ex in TRAIN]
    initial_passages = [mean_vector(ex.passage, passage_table) for ex in TRAIN]
    initial_loss = contrastive_loss(score_matrix(initial_queries, initial_passages, temperature))
    for _ in range(steps):
        train_step(query_table, passage_table, learning_rate=0.35, temperature=temperature)

    query_vectors = [mean_vector(ex.query, query_table) for ex in TRAIN]
    positive_vectors = [mean_vector(ex.passage, passage_table) for ex in TRAIN]
    final_matrix = score_matrix(query_vectors, positive_vectors, temperature)
    final_loss = contrastive_loss(final_matrix)
    passage_index = [mean_vector(passage, passage_table) for passage in CORPUS]  # offline
    retrieval_scores = score_matrix(query_vectors, passage_index, temperature=1.0)  # online exact search
    metrics = retrieval_metrics(retrieval_scores, QRELS)

    print(f"seed={SEED} vocab={len(vocab)}")
    print(f"query_vectors=({len(query_vectors)}, {DIM}) passage_index=({len(passage_index)}, {DIM})")
    print(f"train_score_matrix=({len(final_matrix)}, {len(final_matrix[0])})")
    print(f"loss: {initial_loss:.4f} -> {final_loss:.4f}")
    print("metrics: " + " ".join(f"{key}={value:.3f}" for key, value in metrics.items()))
    for i, ex in enumerate(TRAIN):
        top_id = max(range(len(CORPUS)), key=lambda j: retrieval_scores[i][j])
        print(f"q{i}: top1={top_id} positive={QRELS[i]} text={CORPUS[top_id]}")

    assert len(query_vectors) == 6 and len(query_vectors[0]) == DIM
    assert len(passage_index) == 10 and len(passage_index[0]) == DIM
    assert len(final_matrix) == 6 and len(final_matrix[0]) == 6
    assert final_loss < initial_loss
    assert all(0.0 <= value <= 1.0 for value in metrics.values())
    if check_only:
        assert metrics["recall@1"] == 1.0
        print("check-only: PASS")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--steps", type=int, default=240)
    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)