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 xorshift64(state: int) -> int:
    """64-bit XOR-shift PRNG, single step. Matches RTL behavior."""
    mask64 = (1 << 64) - 1
    s = state & mask64
    s ^= (s << 13) & mask64
    s ^= (s >> 7) & mask64
    s ^= (s << 17) & mask64
    return s


def generate_write_flip_mask(
    prng_state: int,
    hash_bits: int,
    noise_gen_bits: int,
    noise_sample_bits: int,
    rate_num: int,
    rate_den: int,
) -> tuple[int, int]:
    """
    Generate write-noise flip mask.
    Returns (flip_mask, next_prng_state).
    Matches RTL multi-cycle GEN_MASK behavior.

    Each cycle processes noise_gen_bits bit decisions:
    - Advance xorshift64 → 64-bit output
    - Split into noise_gen_bits x noise_sample_bits-bit samples
    - Each sample < THRESHOLD → that bit flips
    """
    assert hash_bits % noise_gen_bits == 0
    assert noise_gen_bits * noise_sample_bits == 64
    mask = 0
    state = prng_state
    sample_range = 1 << noise_sample_bits
    threshold = (rate_num * sample_range) // rate_den

    for bit_offset in range(0, hash_bits, noise_gen_bits):
        # Advance PRNG
        state = xorshift64(state)

        # Split into noise_gen_bits independent samples
        for b in range(noise_gen_bits):
            sample_b = (state >> (b * noise_sample_bits)) & (sample_range - 1)
            if sample_b < threshold:
                mask |= (1 << (bit_offset + b))

    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