from __future__ import annotations

from dataclasses import dataclass
from typing import 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 match_top1(
    query: int,
    rows: Sequence[int],
    *,
    width: int = 512,
) -> MatchResult:
    """Pure matching — noise is already baked into rows at write time."""
    scores = np.zeros(len(rows), dtype=np.int32)

    best_index = 0
    best_score = -1

    for idx, row in enumerate(rows):
        score = xnor_popcount_score(int(query), int(row), 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 xorshift128(state: int) -> int:
    """128-bit xorshift PRNG, single step. Matches random128.sv."""
    mask32 = (1 << 32) - 1
    mask128 = (1 << 128) - 1
    s = state & mask128
    x = (s >> 96) & mask32
    y = (s >> 64) & mask32
    z = (s >> 32) & mask32
    w = s & mask32
    t = (x ^ ((x << 11) & mask32)) & mask32
    next_x = y
    next_y = z
    next_z = w
    next_w = (w ^ (w >> 19) ^ t ^ (t >> 8)) & mask32
    return ((next_x << 96) | (next_y << 64) | (next_z << 32) | next_w) & mask128


def generate_write_flip_mask(
    prng_state: int,
    hash_bits: int,
    noise_bits: int,
    rate_num: int,
    rate_den: int,
) -> tuple[int, int]:
    """Generate one write-noise flip mask using one xorshift128 step."""
    assert hash_bits % noise_bits == 0
    group_bits = hash_bits // noise_bits
    bit_index_bits = 6
    sample_bits = 8
    group_random_bits = bit_index_bits + sample_bits
    assert group_bits == 64
    assert noise_bits * group_random_bits <= 128

    sample_range = 1 << sample_bits
    threshold = (rate_num * sample_range) // rate_den

    state = xorshift128(prng_state)
    mask = 0
    for group_idx in range(noise_bits):
        group_rand = (state >> (group_idx * group_random_bits)) & ((1 << group_random_bits) - 1)
        bit_idx = group_rand & ((1 << bit_index_bits) - 1)
        sample = (group_rand >> bit_index_bits) & (sample_range - 1)
        if sample < threshold:
            mask |= 1 << (group_idx * group_bits + bit_idx)

    return mask, state


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 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