#!/usr/bin/env python3
"""
substrate_runner_v01.py - Substrate-aware matmul runner

Minimal runner that selects between:
  - 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

Usage:
    PYTHONPATH=. python3 helix_cdc/drivers/substrate_runner_v01.py \
        --backend cdna \
        --hxz seeds/shards/blk.0.attn_k.weight.hxz \
        --cdna seeds/cdna_shards/blk.0.attn_k.weight.cdna.hxz

Work Order: WO-SUBSTRATE-INTEGRATION-01
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import numpy as np

# Lazy imports
_backends_module = None


def _lazy_backends():
    global _backends_module
    if _backends_module is None:
        from helix_cdc.backends import substrate_backends
        _backends_module = substrate_backends
    return _backends_module


@dataclass
class SubstrateConfig:
    """Configuration for substrate matmul."""
    backend: str  # "dense", "hxz_stream", "cdna_stream"
    hxz_path: Optional[str] = None
    cdna_path: Optional[str] = None
    block_rows: int = 32
    seed: int = 42


@dataclass
class MatmulResult:
    """Result from a substrate matmul."""
    backend: str
    tensor_name: str
    shape: Tuple[int, int]
    elapsed_ms: float
    peak_memory_bytes: int
    memory_reduction: float
    mse_vs_dense: float
    rel_error_vs_dense: float
    w_decoded_to_float32: bool
    output_mean: float
    output_std: float


class SubstrateRunner:
    """
    Substrate-aware matmul runner.

    Wraps a single HXZ/CDNA shard and provides matmul via:
      - DENSE: Full tensor in RAM
      - HXZ_STREAM: Row-block streaming from HXZ1
      - CDNA_STREAM: TRUE CDC via uint8 indices
    """

    def __init__(self, config: SubstrateConfig):
        """
        Initialize substrate runner.

        Args:
            config: SubstrateConfig with paths and backend choice
        """
        self.config = config
        self._validate_config()

        # Get shape from HXZ metadata
        self._load_metadata()

    def _validate_config(self) -> None:
        """Validate configuration."""
        if self.config.backend == "cdna_stream":
            if not self.config.cdna_path:
                raise ValueError("CDNA backend requires --cdna path")
            if not Path(self.config.cdna_path).is_file():
                raise FileNotFoundError(f"CDNA not found: {self.config.cdna_path}")

        if self.config.backend in ["dense", "hxz_stream"]:
            if not self.config.hxz_path:
                raise ValueError(f"{self.config.backend} backend requires --hxz path")
            if not Path(self.config.hxz_path).is_file():
                raise FileNotFoundError(f"HXZ not found: {self.config.hxz_path}")

    def _load_metadata(self) -> None:
        """Load tensor metadata from HXZ/CDNA."""
        backends = _lazy_backends()

        if self.config.cdna_path and Path(self.config.cdna_path).is_file():
            from helix_cdc.regrow.cdna_stream import inspect_cdna_hxz_header
            info = inspect_cdna_hxz_header(self.config.cdna_path)
            self.tensor_name = info.get("tensor_name", "unknown")
            self.shape = tuple(info.get("shape", [0, 0]))
        elif self.config.hxz_path and Path(self.config.hxz_path).is_file():
            from helix_cdc.regrow.hxz_model_loader import load_hxz_to_state_dict
            state_dict, meta = load_hxz_to_state_dict(self.config.hxz_path)
            tensor_name, tensor = next(iter(state_dict.items()))
            self.tensor_name = tensor_name
            self.shape = tuple(tensor.shape)
        else:
            self.tensor_name = "unknown"
            self.shape = (0, 0)

    @property
    def input_dim(self) -> int:
        """Input dimension (columns)."""
        return self.shape[1] if len(self.shape) >= 2 else 0

    @property
    def output_dim(self) -> int:
        """Output dimension (rows)."""
        return self.shape[0] if len(self.shape) >= 1 else 0

    def matmul(
        self,
        x: np.ndarray,
        y_dense_ref: Optional[np.ndarray] = None,
    ) -> Tuple[np.ndarray, MatmulResult]:
        """
        Run matmul y = W @ x using configured backend.

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

        Returns:
            (y [M], MatmulResult)
        """
        backends = _lazy_backends()

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

        if self.config.backend == "dense":
            y, result = backends.run_dense_matmul(
                Path(self.config.hxz_path), x, y_dense_ref
            )
        elif self.config.backend == "hxz_stream":
            y, result = backends.run_hxz_stream_matmul(
                Path(self.config.hxz_path), x, y_dense_ref,
                block_rows=self.config.block_rows
            )
        elif self.config.backend == "cdna_stream":
            y, result = backends.run_cdna_stream_matmul(
                Path(self.config.cdna_path), x, y_dense_ref,
                block_rows=self.config.block_rows
            )
        else:
            raise ValueError(f"Unknown backend: {self.config.backend}")

        # Determine if W was decoded
        w_decoded = self.config.backend == "dense"

        return y, MatmulResult(
            backend=self.config.backend,
            tensor_name=self.tensor_name,
            shape=self.shape,
            elapsed_ms=result.elapsed_ms,
            peak_memory_bytes=result.peak_memory_bytes,
            memory_reduction=result.memory_reduction,
            mse_vs_dense=result.mse_vs_dense,
            rel_error_vs_dense=result.rel_error_vs_dense,
            w_decoded_to_float32=w_decoded,
            output_mean=float(np.mean(y)),
            output_std=float(np.std(y)),
        )

    def run_all_backends(
        self,
        x: np.ndarray,
    ) -> Dict[str, MatmulResult]:
        """
        Run matmul on all available backends for comparison.

        Args:
            x: Input vector [K]

        Returns:
            Dict of backend_name -> MatmulResult
        """
        backends = _lazy_backends()

        results = {}

        # Dense first (reference)
        if self.config.hxz_path and Path(self.config.hxz_path).is_file():
            y_dense, result = backends.run_dense_matmul(
                Path(self.config.hxz_path), x
            )
            results["dense"] = MatmulResult(
                backend="dense",
                tensor_name=self.tensor_name,
                shape=self.shape,
                elapsed_ms=result.elapsed_ms,
                peak_memory_bytes=result.peak_memory_bytes,
                memory_reduction=1.0,
                mse_vs_dense=0.0,
                rel_error_vs_dense=0.0,
                w_decoded_to_float32=True,
                output_mean=float(np.mean(y_dense)),
                output_std=float(np.std(y_dense)),
            )

            # HXZ stream
            y_hxz, hxz_result = backends.run_hxz_stream_matmul(
                Path(self.config.hxz_path), x, y_dense,
                block_rows=self.config.block_rows
            )
            results["hxz_stream"] = MatmulResult(
                backend="hxz_stream",
                tensor_name=self.tensor_name,
                shape=self.shape,
                elapsed_ms=hxz_result.elapsed_ms,
                peak_memory_bytes=hxz_result.peak_memory_bytes,
                memory_reduction=hxz_result.memory_reduction,
                mse_vs_dense=hxz_result.mse_vs_dense,
                rel_error_vs_dense=hxz_result.rel_error_vs_dense,
                w_decoded_to_float32=False,
                output_mean=float(np.mean(y_hxz)),
                output_std=float(np.std(y_hxz)),
            )

        # CDNA stream
        if self.config.cdna_path and Path(self.config.cdna_path).is_file():
            y_dense_ref = y_dense if "dense" in results else None
            if y_dense_ref is None:
                # Need to compute dense for reference
                y_dense_ref, _ = backends.run_dense_matmul(
                    Path(self.config.hxz_path), x
                )

            y_cdna, cdna_result = backends.run_cdna_stream_matmul(
                Path(self.config.cdna_path), x, y_dense_ref,
                block_rows=self.config.block_rows
            )
            results["cdna_stream"] = MatmulResult(
                backend="cdna_stream",
                tensor_name=self.tensor_name,
                shape=self.shape,
                elapsed_ms=cdna_result.elapsed_ms,
                peak_memory_bytes=cdna_result.peak_memory_bytes,
                memory_reduction=cdna_result.memory_reduction,
                mse_vs_dense=cdna_result.mse_vs_dense,
                rel_error_vs_dense=cdna_result.rel_error_vs_dense,
                w_decoded_to_float32=False,
                output_mean=float(np.mean(y_cdna)),
                output_std=float(np.std(y_cdna)),
            )

        return results


def main() -> int:
    """CLI entry point."""
    parser = argparse.ArgumentParser(
        description="Substrate-aware matmul runner"
    )
    parser.add_argument(
        "--backend",
        choices=["dense", "hxz_stream", "cdna_stream", "all"],
        default="all",
        help="Which backend to use (default: all)"
    )
    parser.add_argument(
        "--hxz",
        default=None,
        help="Path to HXZ1 shard (required for dense/hxz_stream)"
    )
    parser.add_argument(
        "--cdna",
        default=None,
        help="Path to CDNA shard (required for cdna_stream)"
    )
    parser.add_argument(
        "--block-rows",
        type=int,
        default=32,
        help="Block rows for streaming (default: 32)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for input vector (default: 42)"
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output as JSON"
    )

    args = parser.parse_args()

    # Resolve CDNA path from HXZ if not provided
    if args.cdna is None and args.hxz is not None:
        hxz_path = Path(args.hxz)
        cdna_path = Path("seeds/cdna_shards") / (hxz_path.stem + ".cdna.hxz")
        if cdna_path.is_file():
            args.cdna = str(cdna_path)

    config = SubstrateConfig(
        backend=args.backend if args.backend != "all" else "dense",
        hxz_path=args.hxz,
        cdna_path=args.cdna,
        block_rows=args.block_rows,
        seed=args.seed,
    )

    try:
        runner = SubstrateRunner(config)
    except (FileNotFoundError, ValueError) as e:
        print(f"[ERROR] {e}", file=sys.stderr)
        return 1

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

    if not args.json:
        print("=" * 60)
        print("SUBSTRATE RUNNER")
        print("=" * 60)
        print()
        print(f"Tensor: {runner.tensor_name}")
        print(f"Shape: {runner.shape}")
        print(f"Input dim: {runner.input_dim}")
        print(f"Output dim: {runner.output_dim}")
        print()

    # Run backends
    if args.backend == "all":
        results = runner.run_all_backends(x)

        if args.json:
            print(json.dumps({k: asdict(v) for k, v in results.items()}, indent=2))
        else:
            print(f"{'Backend':<15} {'MSE':<12} {'Rel Err %':<12} {'Memory KB':<12} {'Time ms':<10}")
            print("-" * 60)
            for name, r in results.items():
                print(f"{name:<15} {r.mse_vs_dense:<12.2e} {r.rel_error_vs_dense*100:<12.4f} {r.peak_memory_bytes/1024:<12.1f} {r.elapsed_ms:<10.1f}")
            print()
    else:
        y, result = runner.matmul(x)

        if args.json:
            print(json.dumps(asdict(result), indent=2))
        else:
            print(f"Backend: {result.backend}")
            print(f"Elapsed: {result.elapsed_ms:.1f}ms")
            print(f"Memory: {result.peak_memory_bytes / 1024:.1f} KB")
            print(f"Output mean: {result.output_mean:.4f}")
            print(f"Output std: {result.output_std:.4f}")
            print()

    return 0


if __name__ == "__main__":
    sys.exit(main())
