#!/usr/bin/env python3
"""
CDNA Code Regeneration Experiment

Tests whether CDNA indices can be regenerated from (seed, tensor_name, histogram).

Four experiments:
A. Histogram-Only Regeneration - Match rate when regenerating from histogram + seed
B. Raw Code Compressibility - How small can indices get with zstd/lzma/brotli?
C. Residual/Correction Size - How large is sidecar needed for exact reconstruction?
D. Behavioral Equivalence - Do wrong codes + value sidecars produce equivalent weights?

Usage:
    python -m helix_cdc.echo_core.tools.cdna_regen_experiment \
        --shard seeds/cdna_shards_f16/blk.0.attn_k.weight.cdna.hxz \
        --seed "cdna_regen_v1"

    # Batch mode: all shards
    python -m helix_cdc.echo_core.tools.cdna_regen_experiment \
        --shard-dir seeds/cdna_shards_f16/ \
        --seed "cdna_regen_v1" \
        --output-dir receipts/cdna_regen_experiment/
"""

import argparse
import json
import lzma
import subprocess
import tempfile
import zlib
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import numpy as np

# Import HXZO sidecar reader for Experiment D
try:
    from helix_cdc.regrow.hxzo_sidecar import read_outlier_sidecar, inspect_hxzo_header
    HXZO_AVAILABLE = True
except ImportError:
    HXZO_AVAILABLE = False

# Import from sister module
from .cdna_code_stats import (
    build_histogram,
    histogram_entropy,
    regen_from_histogram,
    compute_match_metrics,
    compute_residual,
    summarize_codebook,
)


# Constants
CDNA_MAGIC = b"HXZC"
CODEBOOK_SIZE = 256


def load_cdna_shard(cdna_path: Path) -> Tuple[np.ndarray, np.ndarray, Dict[str, Any]]:
    """
    Load CDNA shard: indices [M, K] and codebook [256].

    Returns:
        Tuple of:
            - indices: np.ndarray [rows, cols], uint8
            - codebook: np.ndarray [256], float32
            - header: Dict with metadata
    """
    import struct

    with open(cdna_path, "rb") as f:
        # Magic
        magic = f.read(4)
        if magic != CDNA_MAGIC:
            raise ValueError(f"Not CDNA file (magic: {magic!r})")

        # Version
        version = struct.unpack(">H", f.read(2))[0]

        # Header length
        header_len = struct.unpack(">H", f.read(2))[0]

        # Header JSON
        header_json = f.read(header_len)
        header = json.loads(header_json.decode("utf-8"))

        # Codebook: 256 * float32 = 1024 bytes
        codebook_bytes = f.read(CODEBOOK_SIZE * 4)
        codebook = np.frombuffer(codebook_bytes, dtype=np.float32).copy()

        # Indices blob (rest of file)
        indices_blob = f.read()

    # Decompress indices
    encoding = header.get("encoding", "zlib")
    if encoding == "zlib":
        indices_raw = zlib.decompress(indices_blob)
    elif encoding == "raw":
        indices_raw = indices_blob
    else:
        raise ValueError(f"Unknown encoding: {encoding}")

    shape = tuple(header["shape"])
    rows, cols = shape
    indices = np.frombuffer(indices_raw, dtype=np.uint8).copy().reshape(shape)

    return indices, codebook, header


def try_compress_zstd(data: bytes, level: int = 22) -> Optional[int]:
    """
    Compress with zstd --ultra and return compressed size.
    Returns None if zstd not available.
    """
    try:
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp_in:
            tmp_in.write(data)
            tmp_in_path = tmp_in.name

        with tempfile.NamedTemporaryFile(suffix=".zst", delete=False) as tmp_out:
            tmp_out_path = tmp_out.name

        result = subprocess.run(
            ["zstd", f"-{level}", "--ultra", "-f", tmp_in_path, "-o", tmp_out_path],
            capture_output=True,
            timeout=300,
        )

        if result.returncode == 0:
            size = Path(tmp_out_path).stat().st_size
            Path(tmp_in_path).unlink(missing_ok=True)
            Path(tmp_out_path).unlink(missing_ok=True)
            return size
        else:
            return None
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return None
    finally:
        Path(tmp_in_path).unlink(missing_ok=True) if 'tmp_in_path' in dir() else None
        Path(tmp_out_path).unlink(missing_ok=True) if 'tmp_out_path' in dir() else None


def try_compress_brotli(data: bytes, level: int = 11) -> Optional[int]:
    """
    Compress with brotli and return compressed size.
    Returns None if brotli not available.
    """
    try:
        import brotli
        compressed = brotli.compress(data, quality=level)
        return len(compressed)
    except ImportError:
        return None


def compress_lzma(data: bytes, preset: int = 9) -> int:
    """Compress with LZMA and return compressed size."""
    compressed = lzma.compress(data, preset=preset)
    return len(compressed)


def compress_zlib(data: bytes, level: int = 9) -> int:
    """Compress with zlib and return compressed size."""
    compressed = zlib.compress(data, level=level)
    return len(compressed)


class CDNARegenExperiment:
    """
    Test whether CDNA codes can be regenerated from seeds.

    This experiment produces receipts showing:
    1. Match rate when predicting codes from histogram
    2. Compressibility of raw indices
    3. Size of correction sidecar for exact reconstruction
    """

    def __init__(self, seed: str = "cdna_regen_v1"):
        self.seed = seed

    def run_experiment_a(
        self,
        indices: np.ndarray,
        codebook: np.ndarray,
        tensor_name: str,
    ) -> Dict[str, Any]:
        """
        Experiment A: Histogram-Only Regeneration Sanity Test

        Method:
        1. Build histogram from real codes
        2. Regenerate codes using histogram + seed via SHAKE256
        3. Measure match rate and dequantized MSE

        Args:
            indices: Real CDNA indices [M, K]
            codebook: CDNA codebook [256]
            tensor_name: Tensor name for seed derivation

        Returns:
            Dict with experiment results
        """
        shape = indices.shape

        # Build histogram from real codes
        histogram = build_histogram(indices)
        entropy = histogram_entropy(histogram)

        # Regenerate from histogram + seed
        derived_seed = f"{self.seed}:{tensor_name}"
        predicted = regen_from_histogram(histogram, shape, derived_seed)

        # Compute metrics
        metrics = compute_match_metrics(indices, predicted, codebook)

        # Expected baseline: if codes were random, match ~1/256 = 0.39%
        expected_random_match = 1.0 / 256

        return {
            "experiment": "A_histogram_only_regen",
            "tensor_name": tensor_name,
            "shape": list(shape),
            "total_elements": int(indices.size),
            "seed_used": derived_seed,
            "histogram_entropy_bits": entropy,
            "expected_random_match_rate": expected_random_match,
            **metrics,
            "interpretation": (
                "high" if metrics["exact_match_rate"] > 0.1
                else "low" if metrics["exact_match_rate"] < 0.01
                else "marginal"
            ),
        }

    def run_experiment_b(
        self,
        indices: np.ndarray,
        tensor_name: str,
    ) -> Dict[str, Any]:
        """
        Experiment B: Raw Code Compressibility Test

        How small can indices get with standard compression?
        Tests: zstd --ultra, lzma, brotli, zlib

        Args:
            indices: Raw CDNA indices [M, K]
            tensor_name: Tensor name for reporting

        Returns:
            Dict with compression results
        """
        raw_bytes = indices.tobytes()
        raw_size = len(raw_bytes)

        results = {
            "experiment": "B_raw_compressibility",
            "tensor_name": tensor_name,
            "raw_size_bytes": raw_size,
        }

        # zlib (baseline - already used in CDNA format)
        zlib_size = compress_zlib(raw_bytes)
        results["zlib_size"] = zlib_size
        results["zlib_ratio"] = raw_size / zlib_size

        # LZMA (usually better)
        lzma_size = compress_lzma(raw_bytes)
        results["lzma_size"] = lzma_size
        results["lzma_ratio"] = raw_size / lzma_size

        # zstd --ultra (if available)
        zstd_size = try_compress_zstd(raw_bytes)
        if zstd_size is not None:
            results["zstd_size"] = zstd_size
            results["zstd_ratio"] = raw_size / zstd_size
        else:
            results["zstd_size"] = None
            results["zstd_note"] = "zstd not available"

        # brotli (if available)
        brotli_size = try_compress_brotli(raw_bytes)
        if brotli_size is not None:
            results["brotli_size"] = brotli_size
            results["brotli_ratio"] = raw_size / brotli_size
        else:
            results["brotli_size"] = None
            results["brotli_note"] = "brotli not available"

        # Best compression
        valid_sizes = [
            ("zlib", zlib_size),
            ("lzma", lzma_size),
        ]
        if zstd_size:
            valid_sizes.append(("zstd", zstd_size))
        if brotli_size:
            valid_sizes.append(("brotli", brotli_size))

        best_name, best_size = min(valid_sizes, key=lambda x: x[1])
        results["best_compressor"] = best_name
        results["best_size"] = best_size
        results["best_ratio"] = raw_size / best_size

        return results

    def run_experiment_c(
        self,
        indices: np.ndarray,
        codebook: np.ndarray,
        tensor_name: str,
    ) -> Dict[str, Any]:
        """
        Experiment C: Residual/Correction Size Test

        How large is the sidecar needed for exact reconstruction?

        Args:
            indices: Real CDNA indices [M, K]
            codebook: CDNA codebook [256]
            tensor_name: Tensor name for seed derivation

        Returns:
            Dict with residual size analysis
        """
        shape = indices.shape

        # Regenerate predicted codes (same as Experiment A)
        histogram = build_histogram(indices)
        derived_seed = f"{self.seed}:{tensor_name}"
        predicted = regen_from_histogram(histogram, shape, derived_seed)

        # Compute residual
        residual_bytes, residual_stats = compute_residual(indices, predicted)

        # Compress residual
        residual_zlib = compress_zlib(residual_bytes)
        residual_lzma = compress_lzma(residual_bytes)

        # Total "seed + sidecar" size
        # Seed components: histogram (256 * 8B = 2KB) + shape (16B) + seed string (~50B)
        histogram_size = 256 * 8  # int64 histogram
        shape_size = 16
        seed_size = len(derived_seed.encode())
        codebook_size = 256 * 4  # float32 codebook

        seed_total = histogram_size + shape_size + seed_size + codebook_size

        results = {
            "experiment": "C_residual_size",
            "tensor_name": tensor_name,
            "shape": list(shape),
            **residual_stats,
            "residual_zlib_bytes": residual_zlib,
            "residual_lzma_bytes": residual_lzma,
            "seed_overhead_bytes": seed_total,
            "total_seed_plus_residual_zlib": seed_total + residual_zlib,
            "total_seed_plus_residual_lzma": seed_total + residual_lzma,
            "raw_indices_size": indices.nbytes,
        }

        # Compute effective compression vs raw indices
        results["effective_ratio_zlib"] = indices.nbytes / results["total_seed_plus_residual_zlib"]
        results["effective_ratio_lzma"] = indices.nbytes / results["total_seed_plus_residual_lzma"]

        # Verdict
        if residual_stats["correction_rate"] > 0.9:
            results["verdict"] = "histogram_insufficient"
        elif residual_stats["correction_rate"] < 0.1:
            results["verdict"] = "histogram_predictive"
        else:
            results["verdict"] = "partial_structure"

        return results

    def run_experiment_d(
        self,
        indices: np.ndarray,
        codebook: np.ndarray,
        tensor_name: str,
        sidecar_path: Optional[Path] = None,
    ) -> Dict[str, Any]:
        """
        Experiment D: Behavioral Equivalence with Wrong Codes + Value Sidecars

        Tests the hypothesis: Is code arrangement NOISE if outlier VALUES are preserved?

        Method:
        1. Generate WRONG codes from histogram + seed
        2. Dequantize with codebook: wrong_weights = codebook[wrong_codes]
        3. Apply EXISTING value sidecars at outlier positions
        4. Compare to original weights (with same outlier corrections)

        If behavioral equivalence passes (low MSE), arrangement is noise.
        If it fails, arrangement carries information.

        Args:
            indices: Real CDNA indices [M, K]
            codebook: CDNA codebook [256]
            tensor_name: Tensor name for seed derivation
            sidecar_path: Optional path to .hxzo sidecar file

        Returns:
            Dict with behavioral equivalence results
        """
        if not HXZO_AVAILABLE:
            return {
                "experiment": "D_behavioral_equivalence",
                "tensor_name": tensor_name,
                "skipped": True,
                "reason": "HXZO module not available",
            }

        shape = indices.shape
        flat_size = indices.size

        # Generate wrong codes from histogram + seed
        histogram = build_histogram(indices)
        derived_seed = f"{self.seed}:{tensor_name}"
        wrong_codes = regen_from_histogram(histogram, shape, derived_seed)

        # Dequantize both sets of codes (WITHOUT sidecar first)
        real_weights_raw = codebook[indices.ravel()].astype(np.float32)
        wrong_weights_raw = codebook[wrong_codes.ravel()].astype(np.float32)

        # === CONTROL 2: Baseline - how bad are wrong codes WITHOUT sidecar? ===
        real_norm_raw = np.linalg.norm(real_weights_raw)
        wrong_norm_raw = np.linalg.norm(wrong_weights_raw)
        if real_norm_raw > 0 and wrong_norm_raw > 0:
            control2_cosine = float(np.dot(real_weights_raw, wrong_weights_raw) / (real_norm_raw * wrong_norm_raw))
        else:
            control2_cosine = 0.0

        # Check for sidecar
        sidecar_found = False
        outlier_count = 0
        sidecar_size_bytes = 0
        control1_cosine = None  # Will be set if sidecar exists

        # Make copies for sidecar application
        real_weights = real_weights_raw.copy()
        wrong_weights = wrong_weights_raw.copy()

        if sidecar_path and sidecar_path.exists():
            try:
                positions, values, meta = read_outlier_sidecar(str(sidecar_path))
                sidecar_found = True
                outlier_count = len(positions)
                sidecar_size_bytes = sidecar_path.stat().st_size

                # === CONTROL 1: Verify sidecar path - real vs real+sidecar ===
                # Should be ~1.0 since sidecar only patches 0.2%
                real_with_sidecar = real_weights_raw.copy()
                real_with_sidecar[positions] = values.astype(np.float32)
                real_sidecar_norm = np.linalg.norm(real_with_sidecar)
                if real_norm_raw > 0 and real_sidecar_norm > 0:
                    control1_cosine = float(np.dot(real_weights_raw, real_with_sidecar) / (real_norm_raw * real_sidecar_norm))
                else:
                    control1_cosine = 0.0

                # Apply outlier corrections to both sets of weights
                # These are the exact fp16 values that should override codebook lookups
                real_weights[positions] = values.astype(np.float32)
                wrong_weights[positions] = values.astype(np.float32)

            except Exception as e:
                return {
                    "experiment": "D_behavioral_equivalence",
                    "tensor_name": tensor_name,
                    "error": f"Failed to load sidecar: {e}",
                    "sidecar_path": str(sidecar_path),
                }

        # === MAIN TEST: Compute error metrics after sidecar corrections ===
        diff = real_weights - wrong_weights
        mse = float(np.mean(diff ** 2))
        rmse = float(np.sqrt(mse))
        max_abs_error = float(np.max(np.abs(diff)))
        mean_abs_error = float(np.mean(np.abs(diff)))

        # Cosine similarity (for behavioral equivalence)
        real_norm = np.linalg.norm(real_weights)
        wrong_norm = np.linalg.norm(wrong_weights)
        if real_norm > 0 and wrong_norm > 0:
            cosine_sim = float(np.dot(real_weights, wrong_weights) / (real_norm * wrong_norm))
        else:
            cosine_sim = 0.0

        # Relative error
        rel_mse = mse / (np.mean(real_weights ** 2) + 1e-12)

        # Calculate sidecar recovery: how much did sidecar help?
        # If control2 ≈ main_cosine, sidecar didn't help
        sidecar_delta = cosine_sim - control2_cosine if control2_cosine else 0.0

        # Verdict based on cosine similarity threshold
        # > 0.999 = essentially equivalent (arrangement is noise)
        # > 0.99  = highly similar
        # > 0.9   = similar but potentially degraded
        # < 0.9   = significant difference (arrangement matters)
        if cosine_sim > 0.999:
            verdict = "arrangement_is_noise"
        elif cosine_sim > 0.99:
            verdict = "mostly_equivalent"
        elif cosine_sim > 0.9:
            verdict = "degraded_but_functional"
        else:
            verdict = "arrangement_matters"

        results = {
            "experiment": "D_behavioral_equivalence",
            "tensor_name": tensor_name,
            "shape": list(shape),
            "sidecar_found": sidecar_found,
            "sidecar_path": str(sidecar_path) if sidecar_path else None,
            "sidecar_size_bytes": sidecar_size_bytes,
            "outlier_count": outlier_count,
            "outlier_rate": outlier_count / flat_size if flat_size > 0 else 0,
            # Control measurements
            "control1_real_vs_real_sidecar": control1_cosine,  # Should be ~1.0 (sidecar path sanity)
            "control2_real_vs_wrong_no_sidecar": control2_cosine,  # Baseline: how bad are wrong codes?
            # Main measurement
            "cosine_similarity": cosine_sim,  # Main test: real+sidecar vs wrong+sidecar
            "sidecar_recovery_delta": sidecar_delta,  # main - control2: did sidecar help?
            # Error metrics
            "mse_after_sidecar": mse,
            "rmse_after_sidecar": rmse,
            "max_abs_error": max_abs_error,
            "mean_abs_error": mean_abs_error,
            "relative_mse": rel_mse,
            "verdict": verdict,
            "interpretation": {
                "arrangement_is_noise": "Outlier VALUES are the knowledge. Code arrangement doesn't matter.",
                "mostly_equivalent": "Sidecar captures most signal. Slight degradation acceptable.",
                "degraded_but_functional": "Arrangement carries some information. May need code corrections.",
                "arrangement_matters": "Code arrangement IS the information. Sidecars insufficient.",
            }.get(verdict, "unknown"),
        }

        return results

    def run_all_experiments(
        self,
        cdna_path: Path,
        sidecar_dir: Optional[Path] = None,
    ) -> Dict[str, Any]:
        """
        Run all four experiments on a single CDNA shard.

        Args:
            cdna_path: Path to .cdna.hxz file

        Returns:
            Combined results dict
        """
        # Load shard
        indices, codebook, header = load_cdna_shard(cdna_path)
        tensor_name = header.get("tensor_name", cdna_path.stem)

        # Run experiments A-C
        exp_a = self.run_experiment_a(indices, codebook, tensor_name)
        exp_b = self.run_experiment_b(indices, tensor_name)
        exp_c = self.run_experiment_c(indices, codebook, tensor_name)

        # Find matching sidecar for Experiment D
        # Convert tensor name: blk.0.attn_k.weight -> blk_0_attn_k_weight
        sidecar_name = tensor_name.replace(".", "_") + ".hxzo"
        sidecar_path = None

        # Search in known sidecar directories
        if sidecar_dir:
            candidate = sidecar_dir / sidecar_name
            if candidate.exists():
                sidecar_path = candidate
        else:
            # Try default sidecar directories
            for sidecar_base in [
                cdna_path.parent.parent / "sidecars_hybrid_v2",
                cdna_path.parent.parent / "sidecars_hybrid_v2_pilot",
                cdna_path.parent / "sidecars",
            ]:
                candidate = sidecar_base / sidecar_name
                if candidate.exists():
                    sidecar_path = candidate
                    break

        # Run experiment D
        exp_d = self.run_experiment_d(indices, codebook, tensor_name, sidecar_path)

        return {
            "cdna_path": str(cdna_path),
            "tensor_name": tensor_name,
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "seed": self.seed,
            "codebook_stats": summarize_codebook(codebook),
            "experiment_a": exp_a,
            "experiment_b": exp_b,
            "experiment_c": exp_c,
            "experiment_d": exp_d,
        }


def run_batch_experiment(
    shard_dir: Path,
    seed: str,
    output_dir: Path,
    limit: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Run experiments on all shards in a directory.

    Args:
        shard_dir: Directory containing .cdna.hxz files
        seed: Seed string for regeneration
        output_dir: Directory to write receipts
        limit: Optional limit on number of shards to process

    Returns:
        Summary results dict
    """
    shards = sorted(shard_dir.glob("*.cdna.hxz"))
    if limit:
        shards = shards[:limit]

    print(f"Found {len(shards)} CDNA shards")

    experiment = CDNARegenExperiment(seed=seed)
    results = []

    # Aggregators
    total_raw_size = 0
    total_best_compressed = 0
    total_residual_lzma = 0
    match_rates = []
    correction_rates = []
    cosine_sims = []
    exp_d_verdicts = []

    for i, shard_path in enumerate(shards):
        print(f"[{i+1}/{len(shards)}] Processing {shard_path.name}...")
        try:
            result = experiment.run_all_experiments(shard_path)
            results.append(result)

            # Aggregate
            exp_b = result["experiment_b"]
            exp_c = result["experiment_c"]
            exp_d = result.get("experiment_d", {})

            total_raw_size += exp_b["raw_size_bytes"]
            total_best_compressed += exp_b["best_size"]
            total_residual_lzma += exp_c["residual_lzma_bytes"]
            match_rates.append(result["experiment_a"]["exact_match_rate"])
            correction_rates.append(exp_c["correction_rate"])

            # Experiment D aggregation
            if not exp_d.get("skipped") and "cosine_similarity" in exp_d:
                cosine_sims.append(exp_d["cosine_similarity"])
                exp_d_verdicts.append(exp_d.get("verdict", "unknown"))

        except Exception as e:
            print(f"  ERROR: {e}")
            results.append({
                "cdna_path": str(shard_path),
                "error": str(e),
            })

    # Experiment D verdict counts
    verdict_counts = {}
    for v in exp_d_verdicts:
        verdict_counts[v] = verdict_counts.get(v, 0) + 1

    # Summary
    summary = {
        "receipt_type": "cdna_regen_experiment_batch",
        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
        "seed": seed,
        "shards_processed": len(results),
        "shards_successful": len([r for r in results if "error" not in r]),
        "totals": {
            "raw_indices_bytes": total_raw_size,
            "best_compressed_bytes": total_best_compressed,
            "residual_lzma_bytes": total_residual_lzma,
            "compression_ratio_best": total_raw_size / total_best_compressed if total_best_compressed > 0 else 0,
        },
        "averages": {
            "match_rate": float(np.mean(match_rates)) if match_rates else 0,
            "match_rate_std": float(np.std(match_rates)) if match_rates else 0,
            "correction_rate": float(np.mean(correction_rates)) if correction_rates else 0,
            "correction_rate_std": float(np.std(correction_rates)) if correction_rates else 0,
        },
        "experiment_d_summary": {
            "shards_with_sidecars": len(cosine_sims),
            "avg_cosine_similarity": float(np.mean(cosine_sims)) if cosine_sims else None,
            "std_cosine_similarity": float(np.std(cosine_sims)) if cosine_sims else None,
            "min_cosine_similarity": float(np.min(cosine_sims)) if cosine_sims else None,
            "verdict_counts": verdict_counts,
        },
    }

    # Write receipts
    output_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    summary_path = output_dir / f"summary_{timestamp}.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\nSummary written to: {summary_path}")

    # Write per-shard results
    per_shard_dir = output_dir / f"per_shard_{timestamp}"
    per_shard_dir.mkdir(exist_ok=True)
    for result in results:
        if "error" not in result:
            tensor_name = result["tensor_name"].replace(".", "_")
            shard_path = per_shard_dir / f"{tensor_name}.json"
            with open(shard_path, "w") as f:
                json.dump(result, f, indent=2)

    print(f"Per-shard results written to: {per_shard_dir}")

    return summary


def main():
    parser = argparse.ArgumentParser(
        description="CDNA Code Regeneration Experiment",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # Single shard
    python -m helix_cdc.echo_core.tools.cdna_regen_experiment \\
        --shard seeds/cdna_shards_f16/blk.0.attn_k.weight.cdna.hxz

    # All shards (batch mode)
    python -m helix_cdc.echo_core.tools.cdna_regen_experiment \\
        --shard-dir seeds/cdna_shards_f16/ \\
        --output-dir receipts/cdna_regen_experiment/

    # Limited batch (first 10 shards)
    python -m helix_cdc.echo_core.tools.cdna_regen_experiment \\
        --shard-dir seeds/cdna_shards_f16/ \\
        --limit 10
        """,
    )

    parser.add_argument(
        "--shard",
        type=Path,
        help="Path to single .cdna.hxz file",
    )
    parser.add_argument(
        "--shard-dir",
        type=Path,
        help="Directory containing .cdna.hxz files (batch mode)",
    )
    parser.add_argument(
        "--seed",
        type=str,
        default="cdna_regen_v1",
        help="Seed string for regeneration (default: cdna_regen_v1)",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("receipts/cdna_regen_experiment"),
        help="Output directory for receipts",
    )
    parser.add_argument(
        "--limit",
        type=int,
        help="Limit number of shards to process in batch mode",
    )

    args = parser.parse_args()

    if args.shard:
        # Single shard mode
        experiment = CDNARegenExperiment(seed=args.seed)
        result = experiment.run_all_experiments(args.shard)

        # Print results
        print("\n" + "=" * 60)
        print(f"CDNA Regeneration Experiment: {args.shard.name}")
        print("=" * 60)

        exp_a = result["experiment_a"]
        print(f"\nExperiment A: Histogram-Only Regeneration")
        print(f"  Match rate:    {exp_a['exact_match_rate']*100:.4f}%")
        print(f"  Expected random: {exp_a['expected_random_match_rate']*100:.2f}%")
        print(f"  MSE (dequant): {exp_a['mse_dequantized']:.6e}")
        print(f"  Entropy:       {exp_a['histogram_entropy_bits']:.2f} bits")

        exp_b = result["experiment_b"]
        print(f"\nExperiment B: Raw Code Compressibility")
        print(f"  Raw size:      {exp_b['raw_size_bytes']:,} bytes")
        print(f"  zlib:          {exp_b['zlib_size']:,} ({exp_b['zlib_ratio']:.2f}x)")
        print(f"  lzma:          {exp_b['lzma_size']:,} ({exp_b['lzma_ratio']:.2f}x)")
        if exp_b.get("zstd_size"):
            print(f"  zstd:          {exp_b['zstd_size']:,} ({exp_b['zstd_ratio']:.2f}x)")
        if exp_b.get("brotli_size"):
            print(f"  brotli:        {exp_b['brotli_size']:,} ({exp_b['brotli_ratio']:.2f}x)")
        print(f"  Best:          {exp_b['best_compressor']} ({exp_b['best_ratio']:.2f}x)")

        exp_c = result["experiment_c"]
        print(f"\nExperiment C: Residual/Correction Size")
        print(f"  Corrections:   {exp_c['num_corrections']:,} / {exp_c['total_positions']:,}")
        print(f"  Correction %:  {exp_c['correction_rate']*100:.2f}%")
        print(f"  Residual raw:  {exp_c['residual_bytes_uncompressed']:,} bytes")
        print(f"  Residual lzma: {exp_c['residual_lzma_bytes']:,} bytes")
        print(f"  Verdict:       {exp_c['verdict']}")

        exp_d = result.get("experiment_d", {})
        print(f"\nExperiment D: Behavioral Equivalence (Wrong Codes + Value Sidecars)")
        if exp_d.get("skipped"):
            print(f"  SKIPPED: {exp_d.get('reason', 'unknown')}")
        elif exp_d.get("error"):
            print(f"  ERROR: {exp_d['error']}")
        elif not exp_d.get("sidecar_found"):
            print(f"  No sidecar found for this tensor")
            print(f"  Control 2 (baseline): cosine={exp_d.get('control2_real_vs_wrong_no_sidecar', 0):.6f}")
            print(f"  (Cannot test sidecar recovery without .hxzo file)")
        else:
            print(f"  Sidecar:       {exp_d['sidecar_size_bytes']:,} bytes")
            print(f"  Outliers:      {exp_d['outlier_count']:,} ({exp_d['outlier_rate']*100:.3f}%)")
            print(f"  --- Controls ---")
            print(f"  Control 1 (real vs real+sidecar): {exp_d['control1_real_vs_real_sidecar']:.6f}")
            print(f"    → Should be ~1.0 (verifies sidecar path)")
            print(f"  Control 2 (real vs wrong, NO sidecar): {exp_d['control2_real_vs_wrong_no_sidecar']:.6f}")
            print(f"    → Baseline: how bad are wrong codes?")
            print(f"  --- Main Test ---")
            print(f"  Main (real+sidecar vs wrong+sidecar): {exp_d['cosine_similarity']:.6f}")
            print(f"    → Did sidecar fix wrong codes?")
            print(f"  Sidecar delta: {exp_d['sidecar_recovery_delta']:+.6f}")
            print(f"    → main - control2: sidecar recovery")
            print(f"  MSE after:     {exp_d['mse_after_sidecar']:.6e}")
            print(f"  Verdict:       {exp_d['verdict']}")
            print(f"  Meaning:       {exp_d['interpretation']}")

        # Write receipt
        args.output_dir.mkdir(parents=True, exist_ok=True)
        receipt_path = args.output_dir / f"{args.shard.stem}_receipt.json"
        with open(receipt_path, "w") as f:
            json.dump(result, f, indent=2)
        print(f"\nReceipt written to: {receipt_path}")

    elif args.shard_dir:
        # Batch mode
        summary = run_batch_experiment(
            shard_dir=args.shard_dir,
            seed=args.seed,
            output_dir=args.output_dir,
            limit=args.limit,
        )

        print("\n" + "=" * 60)
        print("BATCH SUMMARY")
        print("=" * 60)
        print(f"Shards processed: {summary['shards_processed']}")
        print(f"Shards successful: {summary['shards_successful']}")
        print(f"\nTotals:")
        print(f"  Raw indices:     {summary['totals']['raw_indices_bytes']:,} bytes")
        print(f"  Best compressed: {summary['totals']['best_compressed_bytes']:,} bytes")
        print(f"  Compression:     {summary['totals']['compression_ratio_best']:.2f}x")
        print(f"  Residual (lzma): {summary['totals']['residual_lzma_bytes']:,} bytes")
        print(f"\nAverages:")
        print(f"  Match rate:      {summary['averages']['match_rate']*100:.4f}% (±{summary['averages']['match_rate_std']*100:.4f}%)")
        print(f"  Correction rate: {summary['averages']['correction_rate']*100:.2f}% (±{summary['averages']['correction_rate_std']*100:.2f}%)")

        exp_d_sum = summary.get("experiment_d_summary", {})
        if exp_d_sum.get("shards_with_sidecars", 0) > 0:
            print(f"\nExperiment D (Behavioral Equivalence):")
            print(f"  Shards w/ sidecar: {exp_d_sum['shards_with_sidecars']}")
            print(f"  Avg cosine sim:    {exp_d_sum['avg_cosine_similarity']:.6f} (±{exp_d_sum['std_cosine_similarity']:.6f})")
            print(f"  Min cosine sim:    {exp_d_sum['min_cosine_similarity']:.6f}")
            print(f"  Verdicts:")
            for verdict, count in exp_d_sum.get("verdict_counts", {}).items():
                print(f"    {verdict}: {count}")
        else:
            print(f"\nExperiment D: No sidecars found")

    else:
        parser.print_help()
        print("\nError: Must specify either --shard or --shard-dir")
        return 1

    return 0


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