"""
FP8-style block affine quantization and compute-on-compressed kernels.

This simulates hardware-style float8 with per-block scale/bias:
  real_w = (int8_val / 127) * scale + bias

The reference path stays in the compressed domain:
  - Never materializes the full float32 weight matrix
  - Supports streaming block-by-block
  - Ready for future C/CUDA/Triton backends
"""

from __future__ import annotations

import json
import struct
import zlib
from pathlib import Path
from typing import Dict, Optional, Tuple

import numpy as np

from helix_cdc.regrow.cdna_stream import CDNA_MAGIC, CDNA_VERSION


def quantize_fp8_block_affine(
    tensor: np.ndarray,
    block_size: int = 256,
    scale_dtype: str = "f16",
    bias_dtype: str = "f16",
    clip_percentile: Optional[float] = None,
    align_to_rows: bool = True,
    eps: float = 1e-6,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, float]]:
    """
    Quantize a tensor into int8 with per-block scale/bias (FP8-style).

    Args:
        tensor: float array to quantize
        block_size: number of weights per block
        scale_dtype: 'f16' or 'f32' for stored scales
        bias_dtype: 'f16' or 'f32' for stored biases
        clip_percentile: optional percentile (e.g., 99.995) to set scale from |residual|;
            if None, uses stddev-based scaling
        align_to_rows: if True, blocks are per-row (blocks_per_row = ceil(cols/block_size));
            avoids cross-row scale sharing when in_dim not divisible by block_size
        eps: minimum stddev to avoid div-by-zero

    Returns:
        indices: uint8, same shape as tensor (stored offset by +128)
        scales: per-block scales (float16/float32)
        biases: per-block biases (float16/float32)
        stats: dict of error metrics
    """
    flat = np.asarray(tensor, dtype=np.float32)
    if align_to_rows and flat.ndim == 2 and clip_percentile is not None and flat.shape[1] % block_size == 0:
        # Fast vectorized path when rows align perfectly with block_size
        rows, cols = flat.shape
        blocks_per_row = cols // block_size
        n_blocks = rows * blocks_per_row

        blocks = flat.reshape(rows, blocks_per_row, block_size)
        bias = blocks.mean(axis=2)
        residual = blocks - bias[..., None]

        residual_abs = np.abs(residual)
        if clip_percentile >= 99.9:
            spread = residual_abs.max(axis=2)
        else:
            sorted_abs = np.sort(residual_abs, axis=2)
            k = int(np.ceil((clip_percentile / 100.0) * block_size)) - 1
            k = max(0, min(block_size - 1, k))
            spread = sorted_abs[:, :, k]
        scale = np.maximum(spread / 127.0, eps)

        q = np.clip(np.round(residual / scale[..., None] * 127.0), -128, 127).astype(np.int8)
        indices = (q.astype(np.int16) + 128).astype(np.uint8).reshape(rows, cols)

        recon = ((q.astype(np.float32) / 127.0) * scale[..., None] + bias[..., None]).reshape(rows, cols)
        err = np.abs(recon - flat)
        max_abs_err = float(err.max())
        mean_abs_err = float(err.mean())

        scales = scale.astype(np.float16 if scale_dtype == "f16" else np.float32).reshape(-1)
        biases = bias.astype(np.float16 if bias_dtype == "f16" else np.float32).reshape(-1)
    elif align_to_rows and flat.ndim == 2:
        rows, cols = flat.shape
        blocks_per_row = (cols + block_size - 1) // block_size
        n_blocks = rows * blocks_per_row

        indices = np.zeros((rows, cols), dtype=np.uint8)
        scales = np.zeros(n_blocks, dtype=np.float32)
        biases = np.zeros(n_blocks, dtype=np.float32)

        max_abs_err = 0.0
        mean_abs_err = 0.0

        block_idx = 0
        for r in range(rows):
            row = flat[r]
            for b in range(blocks_per_row):
                start = b * block_size
                end = min((b + 1) * block_size, cols)
                block = row[start:end]

                bias = float(block.mean())
                if clip_percentile is None:
                    spread = float(block.std())
                else:
                    spread = float(np.percentile(np.abs(block - bias), clip_percentile))
                scale = max(spread / 127.0 if clip_percentile else spread, eps)

                q = np.clip(np.round((block - bias) / scale * 127.0), -128, 127).astype(np.int8)
                indices[r, start:end] = (q.astype(np.int16) + 128).astype(np.uint8)

                recon = (q.astype(np.float32) / 127.0) * scale + bias
                block_err = np.abs(recon - block)
                max_abs_err = max(max_abs_err, float(block_err.max()))
                mean_abs_err += float(block_err.mean())

                scales[block_idx] = scale
                biases[block_idx] = bias
                block_idx += 1

        mean_abs_err /= n_blocks

        scales = scales.astype(np.float16 if scale_dtype == "f16" else np.float32)
        biases = biases.astype(np.float16 if bias_dtype == "f16" else np.float32)
    else:
        flat_vec = flat.ravel()
        rows, cols = flat.shape if flat.ndim == 2 else (1, flat_vec.size)
        blocks_per_row = None
        n_blocks = (flat_vec.size + block_size - 1) // block_size

        indices = np.zeros_like(flat_vec, dtype=np.uint8)
        scales = np.zeros(n_blocks, dtype=np.float32)
        biases = np.zeros(n_blocks, dtype=np.float32)

        max_abs_err = 0.0
        mean_abs_err = 0.0

        for i in range(n_blocks):
            start = i * block_size
            end = min((i + 1) * block_size, flat_vec.size)
            block = flat_vec[start:end]

            bias = float(block.mean())
            if clip_percentile is None:
                spread = float(block.std())
            else:
                spread = float(np.percentile(np.abs(block - bias), clip_percentile))
            scale = max(spread / 127.0 if clip_percentile else spread, eps)

            q = np.clip(np.round((block - bias) / scale * 127.0), -128, 127).astype(
                np.int8
            )
            indices[start:end] = (q.astype(np.int16) + 128).astype(np.uint8)

            recon = (q.astype(np.float32) / 127.0) * scale + bias
            block_err = np.abs(recon - block)
            max_abs_err = max(max_abs_err, float(block_err.max()))
            mean_abs_err += float(block_err.mean())

            scales[i] = scale
            biases[i] = bias

        if n_blocks:
            mean_abs_err /= n_blocks

        scale_arr = scales.astype(np.float16 if scale_dtype == "f16" else np.float32)
        bias_arr = biases.astype(np.float16 if bias_dtype == "f16" else np.float32)

        stats = {
            "n_blocks": n_blocks,
            "block_size": block_size,
            "scale_dtype": str(scale_arr.dtype),
            "bias_dtype": str(bias_arr.dtype),
            "clip_percentile": clip_percentile,
            "align_to_rows": align_to_rows,
            "max_abs_error": max_abs_err,
            "mean_abs_error": mean_abs_err,
        }

        return indices.reshape(flat.shape), scale_arr, bias_arr, stats

    stats = {
        "n_blocks": n_blocks,
        "block_size": block_size,
        "scale_dtype": str(scales.dtype),
        "bias_dtype": str(biases.dtype),
        "clip_percentile": clip_percentile,
        "align_to_rows": align_to_rows,
        "max_abs_error": max_abs_err,
        "mean_abs_error": mean_abs_err,
    }

    return indices, scales, biases, stats


def fp8_block_affine_linear_forward(
    indices: np.ndarray,
    scales: np.ndarray,
    biases: np.ndarray,
    x: np.ndarray,
    block_size: int = 256,
    align_to_rows: bool = True,
) -> np.ndarray:
    """
    Compute y = (int8 * scale + bias) @ x without expanding W.

    Args:
        indices: [out_dim, in_dim] uint8 (stored int8 offset by +128)
        scales: [n_blocks] scales for each block
        biases: [n_blocks] biases for each block
        x: [in_dim] or [in_dim, batch] input matrix (also accepts [batch, in_dim])
        block_size: number of weights per block
        align_to_rows: if True, block id = row * blocks_per_row + j // block_size; keeps blocks per-row

    Returns:
        y with shape [out_dim] (if x 1D) or [out_dim, batch]
    """
    idx = np.asarray(indices, dtype=np.uint8)
    scales = np.asarray(scales)
    biases = np.asarray(biases)
    x_arr = np.asarray(x, dtype=np.float32)

    out_dim, in_dim = idx.shape
    blocks_per_row = (in_dim + block_size - 1) // block_size if align_to_rows else None
    expected_blocks = out_dim * blocks_per_row if align_to_rows else (idx.size + block_size - 1) // block_size
    if len(scales) != expected_blocks or len(biases) != expected_blocks:
        raise ValueError(
            f"Scale/bias length mismatch: expected {expected_blocks} blocks, "
            f"got scales={len(scales)}, biases={len(biases)}"
        )

    if x_arr.ndim == 1:
        x_arr = x_arr.reshape(in_dim, 1)
    elif x_arr.shape[0] != in_dim and x_arr.shape[1] == in_dim:
        x_arr = x_arr.T

    if x_arr.shape[0] != in_dim:
        raise ValueError(
            f"Input dimension mismatch: indices expect {in_dim} features, got {x_arr.shape[0]}"
        )

    batch_size = x_arr.shape[1]
    y = np.zeros((out_dim, batch_size), dtype=np.float32)

    for b in range(batch_size):
        x_vec = x_arr[:, b]
        for i in range(out_dim):
            acc = 0.0
            row_offset = i * in_dim
            base_block = i * blocks_per_row if align_to_rows else 0
            for j in range(in_dim):
                block = base_block + (j // block_size) if align_to_rows else (row_offset + j) // block_size
                v = (int(idx[i, j]) - 128) / 127.0
                w = v * float(scales[block]) + float(biases[block])
                acc += w * x_vec[j]
            y[i, b] = acc

    return y[:, 0] if batch_size == 1 else y


def write_fp8_block_affine_hxz(
    indices: np.ndarray,
    scales: np.ndarray,
    biases: np.ndarray,
    tensor_name: str,
    output_path: str,
    block_size: int = 256,
    encoding: str = "zlib",
    compression_level: int = 3,
    align_to_rows: bool = True,
    extra_meta: Optional[Dict[str, object]] = None,
) -> Dict[str, object]:
    """
    Serialize FP8 block-affine weights into a CDNA-style .hxz container.

    Layout:
      [magic HXZC][version][header_len][header_json][scales][biases][compressed indices]
    """
    idx = np.asarray(indices, dtype=np.uint8)
    scales_arr = np.asarray(scales)
    biases_arr = np.asarray(biases)

    if scales_arr.shape != biases_arr.shape:
        raise ValueError("scales and biases must have identical shape")

    rows, cols = idx.shape
    blocks_per_row = (cols + block_size - 1) // block_size if align_to_rows else None
    n_blocks = rows * blocks_per_row if align_to_rows else (idx.size + block_size - 1) // block_size
    if len(scales_arr) != n_blocks:
        raise ValueError(
            f"Expected {n_blocks} blocks, got {len(scales_arr)} scale entries"
        )

    scale_dtype = "f16" if scales_arr.dtype == np.float16 else "f32"
    bias_dtype = "f16" if biases_arr.dtype == np.float16 else "f32"

    header = {
        "schema": "cdna_fp8_block_affine_v1",
        "tensor_name": tensor_name,
        "shape": [rows, cols],
        "block_size": block_size,
        "quantization": "int8",
        "scale_dtype": scale_dtype,
        "bias_dtype": bias_dtype,
        "encoding": encoding,
        "align_to_rows": align_to_rows,
    }
    if extra_meta:
        header.update(extra_meta)

    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")

    indices_bytes = idx.tobytes()
    if encoding == "zlib":
        compressed = zlib.compress(indices_bytes, level=compression_level)
    elif encoding == "raw":
        compressed = indices_bytes
    else:
        raise ValueError(f"Unsupported encoding: {encoding}")

    payload = bytearray()
    payload.extend(CDNA_MAGIC)
    payload.extend(struct.pack(">H", CDNA_VERSION))
    payload.extend(struct.pack(">H", len(header_bytes)))
    payload.extend(header_bytes)
    payload.extend(scales_arr.tobytes())
    payload.extend(biases_arr.tobytes())
    payload.extend(compressed)

    Path(output_path).write_bytes(bytes(payload))

    return {
        "schema": "write_fp8_block_affine_hxz_v1",
        "tensor_name": tensor_name,
        "shape": [rows, cols],
        "block_size": block_size,
        "scale_dtype": scale_dtype,
        "bias_dtype": bias_dtype,
        "encoding": encoding,
        "compression_level": compression_level,
        "indices_bytes": len(indices_bytes),
        "compressed_indices_bytes": len(compressed),
        "scales_bytes": scales_arr.nbytes,
        "biases_bytes": biases_arr.nbytes,
        "total_bytes": len(payload),
        "path": str(output_path),
    }


def read_fp8_block_affine_hxz(
    hxz_path: str,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, object]]:
    """Load a block-affine FP8 shard written by write_fp8_block_affine_hxz."""
    path = Path(hxz_path)
    with path.open("rb") as f:
        magic = f.read(4)
        if magic != CDNA_MAGIC:
            raise ValueError(f"Not an HXZ file: magic {magic!r}")

        version = struct.unpack(">H", f.read(2))[0]
        if version != CDNA_VERSION:
            raise ValueError(f"Unsupported version: {version}")

        header_len = struct.unpack(">H", f.read(2))[0]
        header = json.loads(f.read(header_len).decode("utf-8"))

        rows, cols = header["shape"]
        block_size = header.get("block_size", 256)
        align_to_rows = header.get("align_to_rows", True)
        blocks_per_row = (cols + block_size - 1) // block_size if align_to_rows else None
        n_blocks = rows * blocks_per_row if align_to_rows else (rows * cols + block_size - 1) // block_size

        scale_dtype = np.float16 if header.get("scale_dtype", "f16") == "f16" else np.float32
        bias_dtype = np.float16 if header.get("bias_dtype", "f16") == "f16" else np.float32

        scale_bytes = n_blocks * np.dtype(scale_dtype).itemsize
        bias_bytes = n_blocks * np.dtype(bias_dtype).itemsize

        scales = np.frombuffer(f.read(scale_bytes), dtype=scale_dtype)
        biases = np.frombuffer(f.read(bias_bytes), dtype=bias_dtype)

        compressed = f.read()
        if header.get("encoding", "zlib") == "zlib":
            indices_bytes = zlib.decompress(compressed)
        else:
            indices_bytes = compressed

    indices = np.frombuffer(indices_bytes, dtype=np.uint8).reshape(rows, cols)

    meta = {
        "schema": header.get("schema"),
        "tensor_name": header.get("tensor_name"),
        "shape": header.get("shape"),
        "block_size": block_size,
        "quantization": header.get("quantization"),
        "scale_dtype": np.dtype(scale_dtype).name,
        "bias_dtype": np.dtype(bias_dtype).name,
        "encoding": header.get("encoding", "zlib"),
        "n_blocks": n_blocks,
        "align_to_rows": align_to_rows,
    }

    return indices, scales, biases, meta
