"""
cdna_layer_stack.py - Multi-layer CDC Forward Pass

Chains multiple CDNA shards into a sequential forward pass:
  x → W0 → W1 → ... → WN → y

All layers use TRUE CDC: weights are NEVER decoded to float32.

Work Order: WO-CDNA-LAYER-STACK-01
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from helix_cdc.regrow.cdna_stream import (
    CDNAStreamReader,
    cdna_matmul_block,
    stream_cdna_matmul,
    inspect_cdna_hxz_header,
)


@dataclass
class LayerResult:
    """Result from a single layer matmul."""
    layer_idx: int
    tensor_name: str
    shape: Tuple[int, int]
    elapsed_ms: float
    peak_memory_bytes: int
    mse_vs_dense: float
    rel_error_vs_dense: float


@dataclass
class StackResult:
    """Result from full stack forward pass."""
    num_layers: int
    total_elapsed_ms: float
    total_peak_memory_bytes: int
    per_layer_results: List[LayerResult]
    final_mse_vs_dense: float
    final_rel_error_vs_dense: float
    w_decoded_to_float32: bool  # Always False for CDC


class CDNALayerStack:
    """
    Multi-layer CDC forward pass.

    Chains multiple CDNA shards sequentially:
      x → W0 @ x → W1 @ (W0 @ x) → ... → output

    All matmuls use TRUE CDC (bincount + dot, W never decoded).
    """

    def __init__(self, cdna_paths: List[str], block_rows: int = 32):
        """
        Initialize layer stack.

        Args:
            cdna_paths: List of CDNA .hxz paths in order [layer0, layer1, ...]
            block_rows: Block size for streaming

        Raises:
            ValueError: If layers have incompatible dimensions
        """
        self.cdna_paths = [Path(p) for p in cdna_paths]
        self.block_rows = block_rows

        # Validate existence and load headers
        self.layer_info: List[Dict] = []
        for path in self.cdna_paths:
            if not path.is_file():
                raise FileNotFoundError(f"CDNA shard not found: {path}")
            info = inspect_cdna_hxz_header(str(path))
            if not info.get("valid"):
                raise ValueError(f"Invalid CDNA file: {path} - {info.get('error')}")
            self.layer_info.append(info)

        # Validate dimension chain
        self._validate_dimensions()

    def _validate_dimensions(self) -> None:
        """Validate that layer dimensions chain correctly."""
        for i in range(len(self.layer_info) - 1):
            curr_rows = self.layer_info[i]["rows"]  # output dim of current
            next_cols = self.layer_info[i + 1]["cols"]  # input dim of next
            if curr_rows != next_cols:
                raise ValueError(
                    f"Dimension mismatch between layer {i} and {i+1}: "
                    f"layer {i} outputs {curr_rows}, layer {i+1} expects {next_cols}"
                )

    @property
    def input_dim(self) -> int:
        """Input dimension of first layer."""
        return self.layer_info[0]["cols"]

    @property
    def output_dim(self) -> int:
        """Output dimension of last layer."""
        return self.layer_info[-1]["rows"]

    @property
    def num_layers(self) -> int:
        """Number of layers in stack."""
        return len(self.cdna_paths)

    def forward_cdc(
        self,
        x: np.ndarray,
        y_dense_ref: Optional[np.ndarray] = None,
    ) -> Tuple[np.ndarray, StackResult]:
        """
        Forward pass through all layers using CDC.

        Args:
            x: Input vector [input_dim]
            y_dense_ref: Reference dense output for error computation

        Returns:
            (y [output_dim], StackResult)
        """
        import time

        x = np.asarray(x, dtype=np.float32).ravel()

        if len(x) != self.input_dim:
            raise ValueError(
                f"Input dimension mismatch: expected {self.input_dim}, got {len(x)}"
            )

        # Forward through each layer
        h = x
        layer_results: List[LayerResult] = []
        total_peak_memory = 0
        total_elapsed = 0.0

        for i, path in enumerate(self.cdna_paths):
            t0 = time.time()
            y, stats = stream_cdna_matmul(str(path), h, block_rows=self.block_rows)
            elapsed = (time.time() - t0) * 1000

            # Track memory
            peak_mem = stats["max_block_bytes_uint8"] + 256 * 4 + len(h) * 4 + len(y) * 4
            total_peak_memory = max(total_peak_memory, peak_mem)
            total_elapsed += elapsed

            # Placeholder for per-layer error (computed later if reference available)
            layer_results.append(LayerResult(
                layer_idx=i,
                tensor_name=stats["tensor_name"],
                shape=tuple(stats["tensor_shape"]),
                elapsed_ms=elapsed,
                peak_memory_bytes=peak_mem,
                mse_vs_dense=0.0,  # Filled later
                rel_error_vs_dense=0.0,  # Filled later
            ))

            h = y

        # Final output
        y_final = h

        # Compute final error vs dense reference
        if y_dense_ref is not None:
            mse = float(np.mean((y_final - y_dense_ref) ** 2))
            denom = max(float(np.abs(y_dense_ref).max()), 1e-10)
            rel_err = float(np.max(np.abs(y_final - y_dense_ref)) / denom)
        else:
            mse = 0.0
            rel_err = 0.0

        result = StackResult(
            num_layers=self.num_layers,
            total_elapsed_ms=total_elapsed,
            total_peak_memory_bytes=total_peak_memory,
            per_layer_results=layer_results,
            final_mse_vs_dense=mse,
            final_rel_error_vs_dense=rel_err,
            w_decoded_to_float32=False,  # THE KEY CLAIM
        )

        return y_final, result

    def forward_dense(
        self,
        x: np.ndarray,
    ) -> Tuple[np.ndarray, float]:
        """
        Forward pass through all layers using dense matmul (for comparison).

        Args:
            x: Input vector [input_dim]

        Returns:
            (y [output_dim], elapsed_ms)
        """
        import time
        from helix_cdc.regrow.hxz_model_loader import load_hxz_to_state_dict

        x = np.asarray(x, dtype=np.float32).ravel()

        t0 = time.time()
        h = x

        for path in self.cdna_paths:
            # For dense comparison, we need to reconstruct W from CDNA
            # This is the ANTI-PATTERN (decoding W) that CDC avoids
            reader = CDNAStreamReader(str(path))
            blob = reader._ensure_indices_blob()
            indices = np.frombuffer(blob, dtype=np.uint8).reshape(reader.rows, reader.cols)
            W = reader.codebook[indices]  # THIS IS THE DECODE STEP CDC AVOIDS
            h = W @ h

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

        return h, elapsed


def run_layer_stack_comparison(
    cdna_paths: List[str],
    seed: int = 42,
    block_rows: int = 32,
) -> Dict[str, Any]:
    """
    Run comparison between CDC and dense stack forward pass.

    Args:
        cdna_paths: List of CDNA shard paths
        seed: Random seed for input vector
        block_rows: Block size for CDC streaming

    Returns:
        Comparison results dict
    """
    stack = CDNALayerStack(cdna_paths, block_rows=block_rows)

    # Generate random input
    rng = np.random.default_rng(seed)
    x = rng.standard_normal(stack.input_dim).astype(np.float32)

    # Run dense first (reference)
    y_dense, dense_elapsed = stack.forward_dense(x)

    # Run CDC
    y_cdc, cdc_result = stack.forward_cdc(x, y_dense_ref=y_dense)

    # Compute per-layer error accumulation (would need to run layer-by-layer)
    # For now, just report final error

    return {
        "num_layers": stack.num_layers,
        "input_dim": stack.input_dim,
        "output_dim": stack.output_dim,
        "dense_elapsed_ms": dense_elapsed,
        "cdc_elapsed_ms": cdc_result.total_elapsed_ms,
        "cdc_peak_memory_bytes": cdc_result.total_peak_memory_bytes,
        "final_mse": cdc_result.final_mse_vs_dense,
        "final_rel_error": cdc_result.final_rel_error_vs_dense,
        "w_decoded_to_float32": cdc_result.w_decoded_to_float32,
        "per_layer": [
            {
                "layer_idx": lr.layer_idx,
                "tensor_name": lr.tensor_name,
                "shape": list(lr.shape),
                "elapsed_ms": lr.elapsed_ms,
            }
            for lr in cdc_result.per_layer_results
        ],
    }
