from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable, Sequence
import numpy as np


@dataclass(frozen=True)
class MatchResult:
    top1_index: int
    top1_score: int
    scores: np.ndarray


def popcount_int(x: int) -> int:
    return int(x.bit_count())


def mask_width(width: int) -> int:
    return (1 << width) - 1


def xnor_popcount_score(query: int, stored: int, width: int = 512) -> int:
    same_bits = ~(query ^ stored) & mask_width(width)
    return popcount_int(same_bits)


def apply_noise(stored: int, noise_mask: int) -> int:
    return stored ^ noise_mask


def match_top1(
    query: int,
    rows: Sequence[int],
    *,
    width: int = 512,
    noise_masks: Sequence[int] | None = None,
) -> MatchResult:
    scores = np.zeros(len(rows), dtype=np.int32)

    best_index = 0
    best_score = -1

    for idx, row in enumerate(rows):
        effective = row if noise_masks is None else apply_noise(row, int(noise_masks[idx]))
        score = xnor_popcount_score(int(query), int(effective), width)
        scores[idx] = score

        # Tie-break: choose the smallest row index.
        if score > best_score:
            best_score = score
            best_index = idx

    return MatchResult(
        top1_index=int(best_index),
        top1_score=int(best_score),
        scores=scores,
    )


def random_hashes(
    rng: np.random.Generator,
    n: int,
    *,
    width: int = 512,
) -> list[int]:
    words = (width + 63) // 64
    out: list[int] = []

    for _ in range(n):
        value = 0
        for w in range(words):
            value |= int(rng.integers(0, 1 << 64, dtype=np.uint64)) << (64 * w)
        out.append(value & mask_width(width))

    return out


def random_noise_masks(
    rng: np.random.Generator,
    n: int,
    *,
    width: int = 512,
    bit_flip_rate: float = 0.0,
) -> list[int]:
    if not (0.0 <= bit_flip_rate <= 1.0):
        raise ValueError("bit_flip_rate must be in [0, 1]")

    masks: list[int] = []
    for _ in range(n):
        bits = rng.random(width) < bit_flip_rate
        value = 0
        for i, bit in enumerate(bits):
            if bool(bit):
                value |= 1 << i
        masks.append(value)
    return masks


def pack_lanes_flat(masks: Sequence[int], *, width: int = 512) -> int:
    flat = 0
    lane_mask = mask_width(width)
    for lane, mask in enumerate(masks):
        flat |= (int(mask) & lane_mask) << (lane * width)
    return flat


def unpack_score_debug_flat(flat: int, num_rows: int, score_bits: int) -> np.ndarray:
    mask = (1 << score_bits) - 1
    return np.array(
        [(int(flat) >> (row * score_bits)) & mask for row in range(num_rows)],
        dtype=np.int32,
    )


def split_hash_to_words_le(value: int, *, width: int = 512, word_bits: int = 32) -> list[int]:
    n_words = width // word_bits
    word_mask = (1 << word_bits) - 1
    return [(int(value) >> (word_bits * i)) & word_mask for i in range(n_words)]


def join_hash_words_le(words: Sequence[int], *, word_bits: int = 32) -> int:
    value = 0
    word_mask = (1 << word_bits) - 1
    for i, word in enumerate(words):
        value |= (int(word) & word_mask) << (word_bits * i)
    return value