"""
substrate_backends.py
Backend selector for HXZ1 substrate matmul operations.

Provides three compute backends:
  - DENSE: Full W in RAM, standard matmul
  - HXZ_STREAM: Stream float32 row blocks from HXZ1
  - CDNA_STREAM: TRUE CDC - stream uint8 indices, W never decoded

Work Order: WO-CDNA-BACKEND-EXPERIMENT-01
"""

import time
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Tuple

import numpy as np

from helix_cdc.regrow.hxz_model_loader import load_hxz_to_state_dict
from helix_cdc.regrow.hxz1_stream import stream_matmul
from helix_cdc.regrow.cdna_stream import stream_cdna_matmul


class SubstrateBackend(str, Enum):
    """Available matmul backends."""
    DENSE = "dense"
    HXZ_STREAM = "hxz_stream"
    CDNA_STREAM = "cdna_stream"


@dataclass
class BackendResult:
    """Result from a backend matmul operation."""
    backend: str
    elapsed_ms: float
    peak_memory_bytes: int
    memory_reduction: float
    mse_vs_dense: float
    rel_error_vs_dense: float
    max_abs_vs_dense: float
    output_mean: float


def run_dense_matmul(
    hxz_path: Path,
    x: np.ndarray,
    y_dense_ref: np.ndarray = None,
) -> Tuple[np.ndarray, BackendResult]:
    """
    Dense matmul: load full W into RAM, compute y = W @ x.

    Args:
        hxz_path: Path to HXZ1 shard
        x: Input vector [K]
        y_dense_ref: Reference dense output for error computation (None = self)

    Returns:
        (y, BackendResult)
    """
    t0 = time.time()

    # Load full tensor
    state_dict, meta = load_hxz_to_state_dict(str(hxz_path))
    tensor_name, tensor = next(iter(state_dict.items()))
    W = np.asarray(tensor, dtype=np.float32)

    # Compute
    y = W @ x

    elapsed_ms = (time.time() - t0) * 1000

    # Memory: W + x + y
    M, K = W.shape
    peak_memory = W.nbytes + x.nbytes + y.nbytes

    # Errors
    if y_dense_ref is None:
        y_dense_ref = y

    mse = float(np.mean((y - y_dense_ref) ** 2))
    max_abs = float(np.max(np.abs(y - y_dense_ref)))
    denom = max(float(np.abs(y_dense_ref).max()), 1e-10)
    rel_error = float(np.max(np.abs(y - y_dense_ref)) / denom)

    return y, BackendResult(
        backend="dense",
        elapsed_ms=elapsed_ms,
        peak_memory_bytes=peak_memory,
        memory_reduction=1.0,
        mse_vs_dense=mse,
        rel_error_vs_dense=rel_error,
        max_abs_vs_dense=max_abs,
        output_mean=float(np.mean(y)),
    )


def run_hxz_stream_matmul(
    hxz_path: Path,
    x: np.ndarray,
    y_dense_ref: np.ndarray,
    block_rows: int = 32,
) -> Tuple[np.ndarray, BackendResult]:
    """
    HXZ1 streaming matmul: stream float32 row blocks from HXZ1.

    Args:
        hxz_path: Path to HXZ1 shard
        x: Input vector [K]
        y_dense_ref: Reference dense output for error computation
        block_rows: Rows per block

    Returns:
        (y, BackendResult)
    """
    t0 = time.time()

    y, stats = stream_matmul(str(hxz_path), x, block_rows=block_rows)

    elapsed_ms = (time.time() - t0) * 1000

    # Memory: block + x + y (W never fully in RAM)
    peak_memory = stats["max_block_bytes"] + x.nbytes + y.nbytes

    # Dense memory for comparison
    M = len(y)
    K = len(x)
    dense_memory = M * K * 4 + x.nbytes + y.nbytes  # float32 W

    # Errors
    mse = float(np.mean((y - y_dense_ref) ** 2))
    max_abs = float(np.max(np.abs(y - y_dense_ref)))
    denom = max(float(np.abs(y_dense_ref).max()), 1e-10)
    rel_error = float(np.max(np.abs(y - y_dense_ref)) / denom)

    return y, BackendResult(
        backend="hxz_stream",
        elapsed_ms=elapsed_ms,
        peak_memory_bytes=peak_memory,
        memory_reduction=dense_memory / peak_memory,
        mse_vs_dense=mse,
        rel_error_vs_dense=rel_error,
        max_abs_vs_dense=max_abs,
        output_mean=float(np.mean(y)),
    )


def run_cdna_stream_matmul(
    cdna_path: Path,
    x: np.ndarray,
    y_dense_ref: np.ndarray,
    block_rows: int = 32,
) -> Tuple[np.ndarray, BackendResult]:
    """
    CDNA streaming matmul: TRUE CDC - stream uint8 indices, W never decoded.

    Args:
        cdna_path: Path to CDNA HXZ seed
        x: Input vector [K]
        y_dense_ref: Reference dense output for error computation
        block_rows: Rows per block

    Returns:
        (y, BackendResult)
    """
    t0 = time.time()

    y, stats = stream_cdna_matmul(str(cdna_path), x, block_rows=block_rows)

    elapsed_ms = (time.time() - t0) * 1000

    # Memory: indices block (uint8) + codebook (256 * 4) + x + y
    peak_memory = stats["max_block_bytes_uint8"] + 256 * 4 + x.nbytes + y.nbytes

    # Dense memory for comparison
    M = len(y)
    K = len(x)
    dense_memory = M * K * 4 + x.nbytes + y.nbytes  # float32 W

    # Errors
    mse = float(np.mean((y - y_dense_ref) ** 2))
    max_abs = float(np.max(np.abs(y - y_dense_ref)))
    denom = max(float(np.abs(y_dense_ref).max()), 1e-10)
    rel_error = float(np.max(np.abs(y - y_dense_ref)) / denom)

    return y, BackendResult(
        backend="cdna_stream",
        elapsed_ms=elapsed_ms,
        peak_memory_bytes=peak_memory,
        memory_reduction=dense_memory / peak_memory,
        mse_vs_dense=mse,
        rel_error_vs_dense=rel_error,
        max_abs_vs_dense=max_abs,
        output_mean=float(np.mean(y)),
    )


def run_all_backends(
    hxz_path: Path,
    cdna_path: Path,
    x: np.ndarray,
    block_rows: int = 32,
) -> dict:
    """
    Run all three backends and compare results.

    Args:
        hxz_path: Path to HXZ1 shard
        cdna_path: Path to CDNA HXZ seed
        x: Input vector [K]
        block_rows: Rows per block for streaming

    Returns:
        Dict with results for each backend
    """
    # Dense first (reference)
    y_dense, dense_result = run_dense_matmul(hxz_path, x)

    # HXZ stream
    y_hxz, hxz_result = run_hxz_stream_matmul(
        hxz_path, x, y_dense, block_rows=block_rows
    )

    # CDNA stream
    y_cdna, cdna_result = run_cdna_stream_matmul(
        cdna_path, x, y_dense, block_rows=block_rows
    )

    return {
        "dense": dense_result,
        "hxz_stream": hxz_result,
        "cdna_stream": cdna_result,
        "claims": {
            "hxz_stream_lossless": hxz_result.mse_vs_dense < 1e-10,
            "cdna_true_cdc": cdna_result.memory_reduction > 100,
        },
    }
