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 match_topk_from_scores(scores: Sequence[int], k: int) -> list[int]:
    """Return row indices sorted by score desc, row index asc (HW tie-break)."""
    if k <= 0:
        raise ValueError("k must be greater than 0")
    return sorted(range(len(scores)), key=lambda idx: (-int(scores[idx]), idx))[: min(k, len(scores))]


def match_topk(
    query: int,
    rows: Sequence[int],
    *,
    width: int = 512,
    k: int = 5,
) -> tuple[list[int], np.ndarray]:
    """Pure Top-K matching — noise is already baked into rows if needed.

    Returns (list of row indices in rank order, NumPy score array).
    """
    scores = np.zeros(len(rows), dtype=np.int32)
    for idx, row in enumerate(rows):
        scores[idx] = xnor_popcount_score(int(query), int(row), width)
    return match_topk_from_scores(scores, k), 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 generate_grouped_flip_mask(
    *,
    random_value: int,
    hash_bits: int,
    noise_bits: int,
    rate_num: int,
    rate_den: int,
) -> int:
    """Generate a grouped flip mask from one 128-bit value.

    This is the shared write/read noise model: 8 default 64-bit groups, one
    candidate flip per group, 6-bit bit index and 8-bit threshold sample.
    It is not independent Bernoulli sampling over all 512 bits.
    """
    assert noise_bits > 0
    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
    assert rate_den > 0
    assert 0 <= rate_num <= rate_den

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

    for group_idx in range(noise_bits):
        group_rand = (random_value >> (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


def lane_seed_128(seed: int, lane: int) -> int:
    """Derive a nonzero 128-bit lane seed matching the RTL salt convention."""
    mask128 = (1 << 128) - 1
    salt = ((lane + 1) * 0x9E37_79B9_7F4A_7C15) & ((1 << 64) - 1)
    mixed64 = (int(seed) ^ salt) & ((1 << 64) - 1)
    state = ((mixed64 << 64) | mixed64) & mask128
    assert state != 0
    return state


def generate_read_lane_masks(
    lane_states: list[int],
    *,
    hash_bits: int,
    noise_bits: int,
    rate_num: int,
    rate_den: int,
    lane_valid: list[bool],
) -> tuple[list[int], list[int]]:
    """Advance valid lane PRNG states once and return one mask per lane."""
    next_states: list[int] = []
    masks: list[int] = []

    for lane, state in enumerate(lane_states):
        if lane_valid[lane]:
            next_state = xorshift128(state)
            mask = generate_grouped_flip_mask(
                random_value=next_state,
                hash_bits=hash_bits,
                noise_bits=noise_bits,
                rate_num=rate_num,
                rate_den=rate_den,
            )
        else:
            next_state = state
            mask = 0
        next_states.append(next_state)
        masks.append(mask)

    return masks, next_states


def score_rows_with_read_noise(
    query: int,
    rows: Sequence[int],
    *,
    lane_states: Sequence[int],
    width: int = 512,
    lanes: int = 8,
    noise_bits: int = 8,
    rate_num: int = 1,
    rate_den: int = 100,
) -> tuple[np.ndarray, list[int]]:
    """Score one query with read noise and return updated lane PRNG states.

    Unlike match_top1_with_read_noise(), this helper is stateful across calls:
    callers pass current lane states in and receive the next states back.
    This matches a DUT that is reset once, then serves multiple queries.
    """
    assert lanes > 0
    assert len(rows) % lanes == 0
    assert len(lane_states) == lanes

    scores = np.zeros(len(rows), dtype=np.int32)
    next_lane_states = [int(state) for state in lane_states]

    for base in range(0, len(rows), lanes):
        lane_valid = [True] * lanes
        masks, next_lane_states = generate_read_lane_masks(
            next_lane_states,
            hash_bits=width,
            noise_bits=noise_bits,
            rate_num=rate_num,
            rate_den=rate_den,
            lane_valid=lane_valid,
        )

        for lane in range(lanes):
            row_idx = base + lane
            noisy_row = int(rows[row_idx]) ^ int(masks[lane])
            scores[row_idx] = xnor_popcount_score(int(query), noisy_row, width)

    return scores, next_lane_states


def match_top1_with_read_noise(
    query: int,
    rows: Sequence[int],
    *,
    width: int = 512,
    lanes: int = 8,
    noise_bits: int = 8,
    rate_num: int = 1,
    rate_den: int = 100,
    seed: int = 0x6A09_E667_F3BC_C909,
) -> MatchResult:
    """Top-1 matching with dynamic read noise, one query in flight."""
    assert lanes > 0
    assert len(rows) % lanes == 0

    scores = np.zeros(len(rows), dtype=np.int32)
    best_index = 0
    best_score = -1
    lane_states = [lane_seed_128(seed, lane) for lane in range(lanes)]

    for base in range(0, len(rows), lanes):
        lane_valid = [True] * lanes
        masks, lane_states = generate_read_lane_masks(
            lane_states,
            hash_bits=width,
            noise_bits=noise_bits,
            rate_num=rate_num,
            rate_den=rate_den,
            lane_valid=lane_valid,
        )

        for lane in range(lanes):
            row_idx = base + lane
            noisy_row = int(rows[row_idx]) ^ masks[lane]
            score = xnor_popcount_score(int(query), noisy_row, width)
            scores[row_idx] = score
            if score > best_score:
                best_score = score
                best_index = row_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 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