from __future__ import annotations

import argparse
import csv
import json
import sys
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Callable, Iterable, Sequence

import numpy as np


BENCHMARK_KS = (1, 5)
PROJECT_ROOT = Path(__file__).resolve().parents[1]
HW_SIM_DIR = PROJECT_ROOT / "hw" / "sim"
if str(HW_SIM_DIR) not in sys.path:
    sys.path.insert(0, str(HW_SIM_DIR))

from model.ref_model import match_topk as ref_match_topk  # noqa: E402


@dataclass(frozen=True)
class RetrievalDataset:
    rows: list[int]
    row_labels: list[int]
    queries: list[int]
    query_labels: list[int]
    hash_bits: int
    num_classes: int
    positives_per_class: int = 0
    queries_per_class: int = 0
    seed: int = 0


@dataclass(frozen=True)
class MetricAccumulator:
    precision_sum: float = 0.0
    recall_sum: float = 0.0
    f1_sum: float = 0.0
    exact_matches: int = 0
    label_hits: int = 0
    count: int = 0

    def add(self, precision: float, recall: float, f1: float, label_hit: bool, exact: bool) -> "MetricAccumulator":
        return MetricAccumulator(
            precision_sum=self.precision_sum + precision,
            recall_sum=self.recall_sum + recall,
            f1_sum=self.f1_sum + f1,
            exact_matches=self.exact_matches + int(exact),
            label_hits=self.label_hits + int(label_hit),
            count=self.count + 1,
        )

    def as_dict(self) -> dict[str, float]:
        if self.count == 0:
            return {
                "macro_precision": 0.0,
                "retrieval_recall": 0.0,
                "macro_f1": 0.0,
                "exact_match_rate": 0.0,
                "recall@k": 0.0,
            }
        return {
            "macro_precision": self.precision_sum / self.count,
            "retrieval_recall": self.recall_sum / self.count,
            "macro_f1": self.f1_sum / self.count,
            "exact_match_rate": self.exact_matches / self.count,
            "recall@k": self.label_hits / self.count,
        }


def project_root() -> Path:
    return PROJECT_ROOT


def words_le_to_int(words: np.ndarray) -> int:
    value = 0
    for idx, word in enumerate(words.tolist()):
        value |= int(word) << (64 * idx)
    return value


def load_retrieval_dataset_npz(path: str | Path) -> RetrievalDataset:
    dataset_path = Path(path)
    if not dataset_path.is_absolute():
        dataset_path = project_root() / dataset_path
    if not dataset_path.exists():
        raise FileNotFoundError(f"retrieval dataset not found: {dataset_path}")

    loaded = np.load(dataset_path)
    rows_words = loaded["rows_words"]
    queries_words = loaded["queries_words"]
    if rows_words.ndim != 2 or queries_words.ndim != 2:
        raise ValueError("rows_words and queries_words must have shape [N, words]")
    if rows_words.shape[1] != queries_words.shape[1]:
        raise ValueError("rows_words and queries_words must use the same word width")

    rows = [words_le_to_int(words) for words in rows_words]
    queries = [words_le_to_int(words) for words in queries_words]
    row_labels = [int(x) for x in loaded["row_labels"].tolist()]
    query_labels = [int(x) for x in loaded["query_labels"].tolist()]
    return RetrievalDataset(
        rows=rows,
        row_labels=row_labels,
        queries=queries,
        query_labels=query_labels,
        hash_bits=int(rows_words.shape[1] * 64),
        num_classes=len(set(row_labels)),
    )


def hamming_similarity_score(query_row: int, stored_row: int, *, width: int) -> int:
    if width <= 0:
        raise ValueError("width must be greater than 0")
    mask = (1 << width) - 1
    distance = ((int(query_row) ^ int(stored_row)) & mask).bit_count()
    return int(width - distance)


def match_topk_hamming(query: int, rows: Sequence[int], *, width: int, k: int) -> list[int]:
    if k <= 0:
        raise ValueError("k must be greater than 0")
    scored = [
        (-hamming_similarity_score(query, row, width=width), row_index)
        for row_index, row in enumerate(rows)
    ]
    scored.sort()
    return [row_index for _, row_index in scored[: min(k, len(scored))]]


def compute_metrics(topk_indices: list[int], row_labels: list[int], query_label: int, k: int) -> tuple[float, float, float]:
    retrieved = topk_indices[:k]
    relevant = {idx for idx, label in enumerate(row_labels) if label == query_label}
    tp = len(set(retrieved) & relevant)
    precision = tp / float(k)
    recall = tp / float(len(relevant)) if relevant else 0.0
    f1 = 0.0 if precision + recall == 0 else (2.0 * precision * recall) / (precision + recall)
    return precision, recall, f1


def _normalized_topk_values(topk_values: Iterable[int]) -> tuple[int, ...]:
    values = tuple(sorted({int(k) for k in topk_values}))
    if not values or values[0] <= 0:
        raise ValueError("topk_values must contain positive integers")
    return values


def run_benchmark(
    dataset_path: str | Path,
    *,
    hash_bits: int = 512,
    topk_values: Iterable[int] = BENCHMARK_KS,
    run_id: str | None = None,
    timer_ns: Callable[[], int] = time.perf_counter_ns,
) -> dict:
    dataset = load_retrieval_dataset_npz(dataset_path)
    if not dataset.rows:
        raise ValueError("cannot benchmark an empty row set")
    if hash_bits != dataset.hash_bits:
        raise ValueError(
            f"hash_bits={hash_bits} does not match dataset width {dataset.hash_bits}"
        )
    if len(dataset.queries) != len(dataset.query_labels):
        raise ValueError("queries and query_labels must have the same length")

    ks = _normalized_topk_values(topk_values)
    max_k = max(ks)
    if max_k > len(dataset.rows):
        raise ValueError("topk_values cannot exceed the number of dataset rows")

    start_ns = timer_ns()
    all_topk = [
        match_topk_hamming(query, dataset.rows, width=hash_bits, k=max_k)
        for query in dataset.queries
    ]
    end_ns = timer_ns()

    golden_topk = [
        ref_match_topk(query, dataset.rows, width=hash_bits, k=max_k)[0]
        for query in dataset.queries
    ]

    accumulators = {k: MetricAccumulator() for k in ks}
    for topk_indices, golden_indices, query_label in zip(all_topk, golden_topk, dataset.query_labels):
        for k in ks:
            precision, recall, f1 = compute_metrics(topk_indices, dataset.row_labels, query_label, k)
            retrieved_labels = [dataset.row_labels[idx] for idx in topk_indices[:k]]
            label_hit = query_label in retrieved_labels
            exact = topk_indices[:k] == golden_indices[:k]
            accumulators[k] = accumulators[k].add(precision, recall, f1, label_hit, exact)

    metrics = {str(k): accumulators[k].as_dict() for k in ks}

    elapsed_ns = max(0, int(end_ns - start_ns))
    num_queries = len(dataset.queries)
    ns_per_query = (elapsed_ns / float(num_queries)) if num_queries else 0.0
    qps = (1_000_000_000.0 / ns_per_query) if ns_per_query > 0 else 0.0
    resolved_run_id = run_id or f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-software-hamming"

    return {
        "run_id": resolved_run_id,
        "mode": "software-hamming",
        "status": "pass",
        "params": {
            "num_rows": len(dataset.rows),
            "hash_bits": int(hash_bits),
            "topk_k": max_k,
        },
        "dataset": {
            "num_classes": dataset.num_classes,
            "positives_per_class": dataset.positives_per_class,
            "queries_per_class": dataset.queries_per_class,
            "num_queries": num_queries,
            "seed": dataset.seed,
        },
        "metrics": metrics,
        "performance": {
            "total_elapsed_ns": elapsed_ns,
            "total_elapsed_sec": elapsed_ns / 1_000_000_000.0,
            "ns_per_query": ns_per_query,
            "queries_per_second": qps,
        },
    }


def write_outputs(out_dir: Path, result: dict) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "metrics.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")

    fieldnames = [
        "run_id", "mode", "num_rows", "hash_bits", "topk_k", "num_queries", "k",
        "macro_precision", "retrieval_recall", "macro_f1", "recall@k", "exact_match_rate",
        "total_elapsed_ns", "ns_per_query", "queries_per_second", "status",
    ]
    with (out_dir / "metrics.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for k, metrics in result["metrics"].items():
            writer.writerow({
                "run_id": result["run_id"],
                "mode": result["mode"],
                "num_rows": result["params"]["num_rows"],
                "hash_bits": result["params"]["hash_bits"],
                "topk_k": result["params"]["topk_k"],
                "num_queries": result["dataset"]["num_queries"],
                "k": int(k),
                "macro_precision": metrics["macro_precision"],
                "retrieval_recall": metrics["retrieval_recall"],
                "macro_f1": metrics["macro_f1"],
                "recall@k": metrics["recall@k"],
                "exact_match_rate": metrics["exact_match_rate"],
                "total_elapsed_ns": result["performance"]["total_elapsed_ns"],
                "ns_per_query": result["performance"]["ns_per_query"],
                "queries_per_second": result["performance"]["queries_per_second"],
                "status": result["status"],
            })

    lines = [
        "# Software CAM Retrieval Benchmark Summary",
        "",
        f"- run_id: `{result['run_id']}`",
        f"- mode: `{result['mode']}`",
        f"- status: `{result['status']}`",
        f"- num_queries: `{result['dataset']['num_queries']}`",
        f"- ns_per_query: `{result['performance']['ns_per_query']}`",
        f"- queries_per_second: `{result['performance']['queries_per_second']}`",
        "",
        "| k | macro_precision | retrieval_recall | macro_f1 | recall@k | exact_match_rate |",
        "|---:|---:|---:|---:|---:|---:|",
    ]
    for k, metrics in result["metrics"].items():
        lines.append(
            f"| {k} | {metrics['macro_precision']:.6f} | {metrics['retrieval_recall']:.6f} | "
            f"{metrics['macro_f1']:.6f} | {metrics['recall@k']:.6f} | {metrics['exact_match_rate']:.6f} |"
        )
    lines.extend([
        "",
        "说明：软件路径直接对 `.npz` 中的 CAM 行整数执行汉明距离 / XNOR-popcount 检索，不使用软件 CAM 时序仿真。",
    ])
    (out_dir / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def output_dir_for(run_id: str, output_root: Path) -> Path:
    return output_root / run_id


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run software Hamming CAM retrieval benchmark.")
    parser.add_argument("--dataset", required=True, help="Prepared CAM retrieval .npz artifact")
    parser.add_argument("--hash-bits", type=int, default=512, help="Hash width in bits")
    parser.add_argument("--topk-k", type=int, default=5, help="Maximum Top-K to report; reports k=1 and this value")
    parser.add_argument("--run-id", default=None, help="Output run id")
    parser.add_argument(
        "--output-root",
        type=Path,
        default=Path("outputs/sw_retrieval_benchmark"),
        help="Directory under which the run directory is written",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    topk_values = (1,) if args.topk_k == 1 else (1, args.topk_k)
    result = run_benchmark(
        args.dataset,
        hash_bits=args.hash_bits,
        topk_values=topk_values,
        run_id=args.run_id,
    )
    out_dir = output_dir_for(result["run_id"], args.output_root)
    write_outputs(out_dir, result)
    print(
        "SW_RETRIEVAL_RESULT "
        f"run_id={result['run_id']} "
        f"num_rows={result['params']['num_rows']} "
        f"hash_bits={result['params']['hash_bits']} "
        f"num_queries={result['dataset']['num_queries']} "
        f"ns_per_query={result['performance']['ns_per_query']:.3f} "
        f"queries_per_second={result['performance']['queries_per_second']:.3f} "
        f"output_dir={out_dir}"
    )


if __name__ == "__main__":
    main()