#!/usr/bin/env python3
"""
CC4: Centroid Regrow Tests + Receipts

Tests for determinism and correctness of Q4_K_M centroid regeneration
via the FlowTorch bridge.

Single source of truth: helix_cdc/regrow_q4km.py
Bridge: helix_cdc/regrow/flowtorch_bridge.py
"""
import sys
import hashlib
from pathlib import Path
from io import BytesIO

import pytest
import numpy as np

# Add helix-cdc to path
HELIX_CDC = Path(__file__).parent.parent
sys.path.insert(0, str(HELIX_CDC))

# Test constants (from ~/.echo/models.d/mistral.yaml)
TEST_DNA_HEX = "1f647d91f6"
TEST_ENTROPY_HEX = "bf690b54"
TEST_SEED_HEX = hashlib.sha256(f"{TEST_ENTROPY_HEX}:{TEST_DNA_HEX}".encode()).hexdigest()
TEST_ATLAS_PATH = HELIX_CDC / "artifacts" / "centroids" / "mistral_v1" / "crystal_centroids_v1.npy"

# Fallback to v0 atlas if v1 doesn't exist
if not TEST_ATLAS_PATH.exists():
    TEST_ATLAS_PATH = HELIX_CDC / "helix_cdc" / "centroid_atlas_v0.npy"


class TestCentroidRegrowDeterminism:
    """Test that centroid regrow is deterministic."""

    def test_centroid_regrow_deterministic_bridge(self):
        """
        Same seed/config → same tensor bytes via FlowTorch bridge.

        CC4 Test #1: Proves bit-exact determinism through the bridge.
        """
        from helix_cdc.regrow.flowtorch_bridge import (
            regen_tensor_q4km_from_seed,
            TensorSpec,
        )

        # Small tensor for quick test (256x256 = 65536 values)
        spec = TensorSpec(
            name="test.attn.q",
            shape=(256, 256),
            dtype="float16",
        )

        # First regen
        vals1, receipt1 = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
            use_fibpi=True,
        )

        # Second regen with identical args
        vals2, receipt2 = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
            use_fibpi=True,
        )

        # Assert shape/dtype match
        assert vals1.shape == vals2.shape == (256, 256), "Shape mismatch"
        assert vals1.dtype == vals2.dtype == np.float16, "Dtype mismatch"

        # Assert byte-for-byte equality
        bytes1 = vals1.tobytes()
        bytes2 = vals2.tobytes()
        assert bytes1 == bytes2, "Regenerated tensors differ!"

        # Assert receipts have same output hash
        hash1 = receipt1.get("bridge", {}).get("output_hash") or hashlib.sha256(bytes1).hexdigest()[:16]
        hash2 = receipt2.get("bridge", {}).get("output_hash") or hashlib.sha256(bytes2).hexdigest()[:16]
        assert hash1 == hash2, "Receipt hashes differ!"

        print(f"✅ Determinism verified: {hash1}")

    def test_centroid_regrow_matches_proven_path(self):
        """
        FlowTorch bridge output matches the original regrow path.

        CC4 Test #2: Proves the bridge is literally just a wrapper.
        """
        from helix_cdc.regrow.flowtorch_bridge import (
            regen_tensor_q4km_from_seed,
            TensorSpec,
            _decode_q4km_blocks_to_floats,
        )
        from helix_cdc.regrow_q4km import (
            load_centroid_atlas,
            grow_q4km_weights,
        )

        # Load atlas
        atlas = load_centroid_atlas(TEST_ATLAS_PATH)

        # Test spec
        spec = TensorSpec(
            name="test.layer.weight",
            shape=(128, 128),
            dtype="float32",  # Use float32 for precision comparison
        )

        # Compute expected bytes via direct regrow path
        numel = 128 * 128
        num_blocks = (numel + 31) // 32
        size_bytes = num_blocks * 24

        reference_buffer = BytesIO()
        _ = grow_q4km_weights(
            dna_hex=TEST_DNA_HEX,
            entropy_seed=TEST_ENTROPY_HEX,
            label=spec.name,
            offset=0,
            size=size_bytes,
            atlas=atlas,
            output_file=reference_buffer,
        )

        # Decode reference to floats
        reference_buffer.seek(0)
        raw_bytes = reference_buffer.read()
        reference_vals = _decode_q4km_blocks_to_floats(raw_bytes, atlas, numel)
        reference_vals = reference_vals.reshape(spec.shape)

        # Now get bridge output
        bridge_vals, receipt = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
            use_fibpi=False,  # Use non-parallel for exact match
        )

        # Compare bytes
        assert reference_vals.tobytes() == bridge_vals.tobytes(), \
            "Bridge output differs from proven path!"

        print(f"✅ Bridge matches proven path: sha256={hashlib.sha256(bridge_vals.tobytes()).hexdigest()[:16]}")


class TestFlowTorchIntegration:
    """Test FlowTorch engine integration."""

    def test_centroid_regrow_tensor_from_seed_roundtrip(self):
        """
        tensor_from_seed() with CRYSTAL_Q4KM produces same bytes as bridge.

        CC4 Test #3: End-to-end FlowTorch → bridge → regrow roundtrip.
        """
        # Add echo-box runtimes to path
        ECHO_BOX = Path.home() / "repos" / "echo-box"
        sys.path.insert(0, str(ECHO_BOX))

        from runtimes.flowtorch.engine import (
            tensor_from_seed,
            RegrowConfig,
            RegrowStrategy,
        )
        from helix_cdc.regrow.flowtorch_bridge import (
            regen_tensor_q4km_from_seed,
            TensorSpec,
        )

        # Build config
        config = RegrowConfig(
            strategy=RegrowStrategy.CRYSTAL_Q4KM,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
            use_fibpi=True,
        )

        shape = (256, 256)
        tensor_name = "test.roundtrip.weight"

        # FlowTorch path
        flowtorch_arr = tensor_from_seed(
            seed_hex=TEST_SEED_HEX,
            shape=shape,
            dtype="float16",
            config=config,
            tensor_name=tensor_name,
        )

        # Convert to numpy if torch tensor
        if hasattr(flowtorch_arr, 'numpy'):
            flowtorch_arr = flowtorch_arr.numpy()

        # Direct bridge path
        spec = TensorSpec(
            name=tensor_name,
            shape=shape,
            dtype="float16",
        )
        bridge_vals, _ = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
            use_fibpi=True,
        )

        # Compare
        assert flowtorch_arr.shape == bridge_vals.shape, "Shape mismatch"
        assert flowtorch_arr.tobytes() == bridge_vals.tobytes(), \
            "FlowTorch output differs from bridge!"

        print(f"✅ FlowTorch roundtrip verified: {hashlib.sha256(flowtorch_arr.tobytes()).hexdigest()[:16]}")


class TestOffsetSemantics:
    """Test CC2-OFFSET: offset and size_bytes semantics."""

    @pytest.mark.xfail(reason="C backend (helix_regen_cli) may not use label/offset in RNG context - needs investigation")
    def test_different_names_different_rng(self):
        """
        Different tensor names → different RNG context, even with offset=0.

        This proves that unique tensor name provides RNG uniqueness.
        """
        from helix_cdc.regrow.flowtorch_bridge import (
            regen_tensor_q4km_from_seed,
            TensorSpec,
        )

        shape = (64, 64)

        # Tensor A
        spec_a = TensorSpec(name="blk.00.attn.q", shape=shape, dtype="float32")
        vals_a, _ = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec_a,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
        )

        # Tensor B (different name, same everything else)
        spec_b = TensorSpec(name="blk.00.attn.k", shape=shape, dtype="float32")
        vals_b, _ = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec_b,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
        )

        # They MUST differ
        assert vals_a.tobytes() != vals_b.tobytes(), \
            "Different tensor names produced identical output - RNG isolation broken!"

        print("✅ Different names → different RNG: isolation verified")

    @pytest.mark.xfail(reason="C backend (helix_regen_cli) may not use offset in RNG context - needs investigation")
    def test_explicit_offset_changes_output(self):
        """
        Same name but different offset → different output.

        This is for parallel slicing of a single tensor.
        """
        from helix_cdc.regrow.flowtorch_bridge import (
            regen_tensor_q4km_from_seed,
            TensorSpec,
        )

        shape = (64, 64)
        name = "test.parallel.chunk"

        # Offset 0
        spec_0 = TensorSpec(name=name, shape=shape, dtype="float32", offset=0)
        vals_0, _ = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec_0,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
        )

        # Offset 1024
        spec_1024 = TensorSpec(name=name, shape=shape, dtype="float32", offset=1024)
        vals_1024, _ = regen_tensor_q4km_from_seed(
            seed_hex=TEST_SEED_HEX,
            spec=spec_1024,
            atlas_path=str(TEST_ATLAS_PATH),
            dna_hex=TEST_DNA_HEX,
            entropy_hex=TEST_ENTROPY_HEX,
        )

        # They MUST differ
        assert vals_0.tobytes() != vals_1024.tobytes(), \
            "Different offsets produced identical output!"

        print("✅ Different offset → different output: FibPi3D slicing verified")


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
