#!/usr/bin/env python3
"""
Packed-Domain Attention Kernels

Compute attention WITHOUT full fp32 KV materialization.
Uses on-the-fly tile dequantization only where needed.
"""
import numpy as np
from typing import Tuple
from ..packing.layouts import Packed4


def attention_packed(
    Q: np.ndarray,
    K_packed: Packed4,
    V_packed: Packed4,
    K_shape_orig: Tuple[int, int, int],
    V_shape_orig: Tuple[int, int, int]
) -> Tuple[np.ndarray, dict]:
    """Compute attention output with packed KV cache

    Args:
        Q: Query tensor (H, D)
        K_packed: Packed key cache
        V_packed: Packed value cache
        K_shape_orig: Original K shape (H, seq_len, D)
        V_shape_orig: Original V shape (H, seq_len, D)

    Returns:
        (output, stats)
        output: Attention output (H, D)
        stats: Dictionary with I/O stats and intermediate values
    """
    H, D = Q.shape
    _, seq_len, _ = K_shape_orig

    # Stats tracking
    stats = {
        'tiles_touched': 0,
        'bytes_read': 0,
        'scores_raw': None,  # For signal-preserving metrics
        'attn_weights': None
    }

    # Step 1: Compute attention scores (Q @ K.T)
    # Shape: (H, seq_len)
    scores = np.zeros((H, seq_len), dtype=np.float32)

    # Dequantize K row-by-row (on-the-fly)
    for h in range(H):
        for t in range(seq_len):
            # Get K[h, t, :] via tile dequantization
            # K is stored as (H * seq_len, D) after reshape
            row_idx = h * seq_len + t

            # Dequantize this row from packed format
            K_row = _dequant_row_from_packed(K_packed, row_idx)
            stats['tiles_touched'] += 1
            stats['bytes_read'] += _bytes_per_row(K_packed)

            # Compute score: Q[h, :] @ K[h, t, :]
            scores[h, t] = Q[h, :] @ K_row

    stats['scores_raw'] = scores.copy()

    # Step 2: Softmax over seq_len dimension
    scores_max = scores.max(axis=1, keepdims=True)
    scores_exp = np.exp(scores - scores_max)
    scores_sum = scores_exp.sum(axis=1, keepdims=True)
    attn_weights = scores_exp / scores_sum  # (H, seq_len)

    stats['attn_weights'] = attn_weights.copy()

    # Step 3: Apply attention weights to V
    # output = attn_weights @ V  (H, seq_len) @ (H, seq_len, D) → (H, D)
    output = np.zeros((H, D), dtype=np.float32)

    for h in range(H):
        for t in range(seq_len):
            # Get V[h, t, :] via tile dequantization
            row_idx = h * seq_len + t

            V_row = _dequant_row_from_packed(V_packed, row_idx)
            stats['tiles_touched'] += 1
            stats['bytes_read'] += _bytes_per_row(V_packed)

            # Accumulate weighted value
            output[h, :] += attn_weights[h, t] * V_row

    return output, stats


def _dequant_row_from_packed(packed: Packed4, row_idx: int) -> np.ndarray:
    """Dequantize a single row from packed format

    Args:
        packed: Packed4 instance (M × N)
        row_idx: Row index to extract

    Returns:
        Dequantized row (N,) float32
    """
    N = packed.N
    tile_h = packed.tile_h
    tile_w = packed.tile_w

    result = np.zeros(N, dtype=np.float32)

    # Determine which tile row this belongs to
    tile_i = row_idx // tile_h
    local_i = row_idx % tile_h

    # Iterate over tile columns
    num_tiles_w = packed.num_tiles_w

    for tile_j in range(num_tiles_w):
        # Get tile parameters
        scale, zero = packed.get_tile_params(tile_i, tile_j)

        # Column range for this tile
        col_start = tile_j * tile_w
        col_end = min(col_start + tile_w, N)

        # Dequantize this row within the tile
        for j in range(col_start, col_end):
            # Get packed nybble
            idx = row_idx * N + j
            byte_idx = idx // 2
            nybble_shift = (idx % 2) * 4

            quant_val = (packed.data[byte_idx] >> nybble_shift) & 0x0F

            # Dequantize
            result[j] = (quant_val - zero) * scale

    return result


def _bytes_per_row(packed: Packed4) -> int:
    """Estimate bytes read for one row dequantization

    Includes data nybbles + amortized scale/zero overhead
    """
    # Data: N/2 bytes per row (2 nybbles per byte)
    # Plus amortized scale/zero from tiles this row touches
    N = packed.N
    num_tiles_w = packed.num_tiles_w

    data_bytes = (N + 1) // 2
    scale_bytes = num_tiles_w * 4  # float32 per tile column
    zero_bytes = num_tiles_w * 1   # int8 per tile column

    return data_bytes + scale_bytes + zero_bytes
