from __future__ import annotations

import argparse
import csv
import json
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Callable, Iterable

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]
    rows_words: np.ndarray
    queries_words: np.ndarray
    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 missing_dataset_message(dataset_path: Path) -> str:
    message = f"retrieval dataset not found: {dataset_path}"
    match = re.fullmatch(
        r"(?P<dataset>cifar10|cifar100)_hash(?P<hash_bits>\d+)_rows(?P<num_rows>\d+)_queries(?P<max_queries>\d+)\.npz",
        dataset_path.name,
    )
    if match is None:
        return message

    groups = match.groupdict()
    command = (
        "python scripts/prepare_cam_retrieval_dataset.py "
        f"--dataset {groups['dataset']} "
        f"--num-rows {groups['num_rows']} "
        f"--max-queries {groups['max_queries']} "
        f"--hash-bits {groups['hash_bits']}"
    )
    return (
        f"{message}\n"
        "The requested benchmark artifact has not been generated yet. "
        "Create it first, then rerun this benchmark:\n"
        f"  {command}"
    )


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(missing_dataset_message(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 = np.asarray(rows_words, dtype=np.uint64)
    queries_words = np.asarray(queries_words, dtype=np.uint64)
    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,
        rows_words=rows_words,
        queries_words=queries_words,
        hash_bits=int(rows_words.shape[1] * 64),
        num_classes=len(set(row_labels)),
    )


def match_topk_numpy(query_words: np.ndarray, rows_words: np.ndarray, *, width: int, k: int) -> list[int]:
    if k <= 0:
        raise ValueError("k must be greater than 0")
    if width <= 0:
        raise ValueError("width must be greater than 0")
    query_words = np.asarray(query_words, dtype=np.uint64)
    rows_words = np.asarray(rows_words, dtype=np.uint64)
    if rows_words.ndim != 2:
        raise ValueError("rows_words must have shape [N, words]")
    if query_words.shape != (rows_words.shape[1],):
        raise ValueError("query_words must have shape [words]")

    return _match_topk_numpy_batch(query_words[np.newaxis, :], rows_words, width=width, k=k)[0]


def _match_topk_numpy_batch(
    queries_words: np.ndarray,
    rows_words: np.ndarray,
    *,
    width: int,
    k: int,
) -> list[list[int]]:
    if k <= 0:
        raise ValueError("k must be greater than 0")
    if width <= 0:
        raise ValueError("width must be greater than 0")
    queries_words = np.asarray(queries_words, dtype=np.uint64)
    rows_words = np.asarray(rows_words, dtype=np.uint64)
    if queries_words.ndim != 2:
        raise ValueError("queries_words must have shape [Q, words]")
    if rows_words.ndim != 2:
        raise ValueError("rows_words must have shape [N, words]")
    if queries_words.shape[1] != rows_words.shape[1]:
        raise ValueError("queries_words and rows_words must use the same word width")

    xor = np.bitwise_xor(queries_words[:, np.newaxis, :], rows_words[np.newaxis, :, :])
    distances = np.bitwise_count(xor).sum(axis=2, dtype=np.int64)
    scores = int(width) - distances
    row_indices = np.arange(rows_words.shape[0], dtype=np.int64)
    topk_count = min(k, rows_words.shape[0])
    return [
        [int(idx) for idx in np.lexsort((row_indices, -score_row))[:topk_count]]
        for score_row in scores
    ]


def match_all_topk_numpy(
    queries_words: np.ndarray,
    rows_words: np.ndarray,
    *,
    width: int,
    k: int,
    workers: int = 1,
) -> list[list[int]]:
    if workers <= 0:
        raise ValueError("workers must be greater than 0")
    queries_words = np.asarray(queries_words, dtype=np.uint64)
    if workers == 1 or len(queries_words) <= 1:
        return _match_topk_numpy_batch(queries_words, rows_words, width=width, k=k)

    chunks = [
        chunk
        for chunk in np.array_split(queries_words, min(workers, len(queries_words)))
        if len(chunk)
    ]
    with ThreadPoolExecutor(max_workers=workers) as executor:
        futures = [
            executor.submit(_match_topk_numpy_batch, chunk, rows_words, width=width, k=k)
            for chunk in chunks
        ]
        parts = [future.result() for future in futures]
    return [topk for part in parts for topk in part]


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,
    workers: int = 1,
    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")
    if workers <= 0:
        raise ValueError("workers must be greater than 0")

    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_all_topk_numpy(
        dataset.queries_words,
        dataset.rows_words,
        width=hash_bits,
        k=max_k,
        workers=workers,
    )
    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-numpy"

    return {
        "run_id": resolved_run_id,
        "mode": "software-numpy",
        "status": "pass",
        "params": {
            "num_rows": len(dataset.rows),
            "hash_bits": int(hash_bits),
            "topk_k": max_k,
            "workers": int(workers),
            "engine": "numpy",
        },
        "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", "workers", "engine", "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"],
                "workers": result["params"].get("workers", 1),
                "engine": result["params"].get("engine", "numpy"),
                "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"- workers: `{result['params'].get('workers', 1)}`",
        f"- engine: `{result['params'].get('engine', 'numpy')}`",
        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` 中的 little-endian uint64 words 使用 NumPy bitwise_count 执行汉明距离 / 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 NumPy software 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("--workers", type=int, default=1, help="Number of software query worker threads")
    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)
    try:
        result = run_benchmark(
            args.dataset,
            hash_bits=args.hash_bits,
            topk_values=topk_values,
            run_id=args.run_id,
            workers=args.workers,
        )
    except FileNotFoundError as exc:
        print(str(exc), file=sys.stderr)
        raise SystemExit(2) from None
    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"workers={result['params']['workers']} "
        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()