"""
CDNA Code Statistics Utilities

Histogram analysis and deterministic regeneration for CDNA indices.
Used to test hypothesis: Can codes be predicted from (seed, tensor_name, histogram)?
"""

import hashlib
import struct
from pathlib import Path
from typing import Dict, Tuple, Optional

import numpy as np


def build_histogram(codes: np.ndarray) -> np.ndarray:
    """
    Build code frequency histogram [256] from CDNA indices.

    Args:
        codes: uint8 array of CDNA indices (any shape)

    Returns:
        np.ndarray: Histogram of shape [256], dtype int64
    """
    return np.bincount(codes.ravel(), minlength=256).astype(np.int64)


def histogram_entropy(histogram: np.ndarray) -> float:
    """
    Compute Shannon entropy of code distribution.

    Args:
        histogram: Code frequency histogram [256]

    Returns:
        Entropy in bits (max 8.0 for uniform distribution)
    """
    total = histogram.sum()
    if total == 0:
        return 0.0

    probs = histogram / total
    # Avoid log(0)
    nonzero = probs > 0
    entropy = -np.sum(probs[nonzero] * np.log2(probs[nonzero]))
    return float(entropy)


def shake256_rng(seed: str, total_bytes: int) -> np.ndarray:
    """
    Generate deterministic random bytes using SHAKE256.

    Args:
        seed: String seed for SHAKE256
        total_bytes: Number of bytes to generate

    Returns:
        np.ndarray: uint8 array of random bytes
    """
    h = hashlib.shake_256(seed.encode("utf-8"))
    random_bytes = h.digest(total_bytes)
    return np.frombuffer(random_bytes, dtype=np.uint8)


def regen_from_histogram(
    histogram: np.ndarray,
    shape: Tuple[int, int],
    seed: str,
) -> np.ndarray:
    """
    Regenerate CDNA codes by sampling from histogram distribution using SHAKE256.

    This tests the hypothesis: If codes are predictable from histogram + seed,
    we should get high match rates against real codes.

    Algorithm:
    1. Build CDF from histogram
    2. Generate deterministic random bytes with SHAKE256
    3. Map each random byte to a code value via inverse CDF

    Args:
        histogram: Code frequency histogram [256]
        shape: Target shape (rows, cols)
        seed: Seed string for SHAKE256

    Returns:
        np.ndarray: Predicted codes of shape [rows, cols], dtype uint8
    """
    rows, cols = shape
    total_elements = rows * cols

    # Build cumulative distribution
    total_count = histogram.sum()
    if total_count == 0:
        raise ValueError("Empty histogram")

    cdf = np.cumsum(histogram).astype(np.float64) / total_count

    # Generate random values in [0, 1) via SHAKE256
    # Use 2 bytes per sample for better resolution
    random_bytes = shake256_rng(seed, total_elements * 2)
    random_u16 = np.frombuffer(random_bytes.tobytes(), dtype=np.uint16)
    random_uniform = random_u16.astype(np.float64) / 65536.0

    # Inverse CDF sampling
    predicted = np.searchsorted(cdf, random_uniform).astype(np.uint8)
    # Clip to [0, 255] (searchsorted can return 256 if uniform == 1.0)
    predicted = np.clip(predicted, 0, 255)

    return predicted.reshape(shape)


def compute_match_metrics(
    real_codes: np.ndarray,
    predicted_codes: np.ndarray,
    codebook: np.ndarray,
) -> Dict[str, float]:
    """
    Compute metrics comparing real vs predicted codes.

    Args:
        real_codes: Ground truth CDNA indices [M, K], uint8
        predicted_codes: Predicted CDNA indices [M, K], uint8
        codebook: CDNA codebook [256], float32

    Returns:
        Dict with:
            - exact_match_rate: % of positions where predicted == real
            - mse_dequantized: MSE after applying codebook
            - max_abs_error: Worst-case dequantization error
            - histogram_match: Whether histograms are identical
    """
    if real_codes.shape != predicted_codes.shape:
        raise ValueError(f"Shape mismatch: {real_codes.shape} vs {predicted_codes.shape}")

    total = real_codes.size
    matches = np.sum(real_codes == predicted_codes)
    exact_match_rate = matches / total

    # Dequantize both
    real_weights = codebook[real_codes.ravel()]
    predicted_weights = codebook[predicted_codes.ravel()]

    diff = real_weights - predicted_weights
    mse = float(np.mean(diff ** 2))
    max_abs_error = float(np.max(np.abs(diff)))

    # Check histogram match
    real_hist = build_histogram(real_codes)
    pred_hist = build_histogram(predicted_codes)
    histogram_match = np.array_equal(real_hist, pred_hist)

    return {
        "exact_match_rate": float(exact_match_rate),
        "exact_match_count": int(matches),
        "total_positions": int(total),
        "mse_dequantized": mse,
        "rmse_dequantized": float(np.sqrt(mse)),
        "max_abs_error": max_abs_error,
        "histogram_match": histogram_match,
    }


def compute_residual(
    real_codes: np.ndarray,
    predicted_codes: np.ndarray,
) -> Tuple[bytes, Dict[str, int]]:
    """
    Encode residual as sparse corrections.

    For each position where predicted != real, store (position, correct_code).
    Uses variable-length encoding for efficiency.

    Args:
        real_codes: Ground truth CDNA indices (flattened or 2D)
        predicted_codes: Predicted CDNA indices (same shape)

    Returns:
        Tuple of:
            - bytes: Encoded residual
            - Dict: Statistics about the residual
    """
    real_flat = real_codes.ravel()
    pred_flat = predicted_codes.ravel()

    # Find mismatches
    mismatch_mask = real_flat != pred_flat
    mismatch_indices = np.where(mismatch_mask)[0]
    correct_values = real_flat[mismatch_mask]

    num_corrections = len(mismatch_indices)
    total_positions = len(real_flat)

    # Encode as: [num_corrections (4B)] + [(position (4B), value (1B))...]
    # For positions, use 4 bytes (supports up to 4B elements)
    parts = [struct.pack("<I", num_corrections)]

    for idx, val in zip(mismatch_indices, correct_values):
        parts.append(struct.pack("<IB", int(idx), int(val)))

    residual_bytes = b"".join(parts)

    stats = {
        "num_corrections": num_corrections,
        "total_positions": total_positions,
        "correction_rate": num_corrections / total_positions if total_positions > 0 else 0.0,
        "residual_bytes_uncompressed": len(residual_bytes),
        "bytes_per_correction": 5,  # 4B position + 1B value
    }

    return residual_bytes, stats


def decode_residual(
    predicted_codes: np.ndarray,
    residual_bytes: bytes,
) -> np.ndarray:
    """
    Apply residual corrections to predicted codes.

    Args:
        predicted_codes: Predicted CDNA indices (any shape)
        residual_bytes: Encoded residual from compute_residual()

    Returns:
        np.ndarray: Corrected codes (same shape as input)
    """
    corrected = predicted_codes.ravel().copy()

    if len(residual_bytes) < 4:
        return corrected.reshape(predicted_codes.shape)

    num_corrections = struct.unpack("<I", residual_bytes[:4])[0]
    offset = 4

    for _ in range(num_corrections):
        idx, val = struct.unpack("<IB", residual_bytes[offset:offset+5])
        corrected[idx] = val
        offset += 5

    return corrected.reshape(predicted_codes.shape)


def summarize_codebook(codebook: np.ndarray) -> Dict[str, float]:
    """
    Summary statistics for a CDNA codebook.

    Args:
        codebook: CDNA codebook [256], float32

    Returns:
        Dict with min, max, mean, std, range stats
    """
    return {
        "min": float(codebook.min()),
        "max": float(codebook.max()),
        "mean": float(codebook.mean()),
        "std": float(codebook.std()),
        "range": float(codebook.max() - codebook.min()),
        "size": len(codebook),
    }


__all__ = [
    "build_histogram",
    "histogram_entropy",
    "shake256_rng",
    "regen_from_histogram",
    "compute_match_metrics",
    "compute_residual",
    "decode_residual",
    "summarize_codebook",
]
