"""Tests for full GPU pipeline — zero Python in the hot path.

Every step is a Sound Block:
  1. build_matrix_gpu (sb_world_model.wgsl)
  2. GPU homeostasis (sb_homeostasis.wgsl)
  3. Colony modulation (sb_colony_block.wgsl)
  4. GPU clustering (sb_clustering.wgsl)
  5. extract_strategies_gpu (sb_strategy_extract.wgsl)
"""

import sys
import time
import unittest
import numpy as np
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from afs.gpu.sb_colony_block_kernel import SBColonyBlockKernel, gpu_available


@unittest.skipUnless(gpu_available(), "GPU not available")
class TestFullGPUPipeline(unittest.TestCase):
    """Full pipeline: prompt → strategies, zero Python in hot path."""

    def _market_decision(self):
        """The market segments decision from the session."""
        return {
            "question": "Which market segment for Sound Blocks?",
            "options": [
                {"name": "Hardware Emulation", "market_size": 85, "price_sensitivity": 15,
                 "technical_readiness": 70, "competition": 20, "revenue_per_customer": 95},
                {"name": "FPGA Prototyping SaaS", "market_size": 70, "price_sensitivity": 45,
                 "technical_readiness": 80, "competition": 35, "revenue_per_customer": 60},
                {"name": "Decision API", "market_size": 60, "price_sensitivity": 50,
                 "technical_readiness": 95, "competition": 55, "revenue_per_customer": 35},
            ],
            "factors": [
                {"name": "market_size", "range": [0, 100]},
                {"name": "price_sensitivity", "range": [0, 100], "goal": "minimize"},
                {"name": "technical_readiness", "range": [0, 100]},
                {"name": "competition", "range": [0, 100], "goal": "minimize"},
                {"name": "revenue_per_customer", "range": [0, 100]},
            ],
            "relationships": [
                {"from": "technical_readiness", "to": "competition", "weight": -0.4},
                {"from": "competition", "to": "price_sensitivity", "weight": 0.7},
                {"from": "market_size", "to": "competition", "weight": 0.5},
                {"from": "price_sensitivity", "to": "revenue_per_customer", "weight": -0.7},
            ],
        }

    def _simple_decision(self):
        return {
            "question": "A vs B vs C",
            "options": [
                {"name": "A", "speed": 80, "cost": 20, "quality": 90},
                {"name": "B", "speed": 40, "cost": 80, "quality": 60},
                {"name": "C", "speed": 60, "cost": 50, "quality": 70},
            ],
            "factors": [
                {"name": "speed", "range": [0, 100]},
                {"name": "cost", "range": [0, 100], "goal": "minimize"},
                {"name": "quality", "range": [0, 100]},
            ],
            "relationships": [
                {"from": "speed", "to": "cost", "weight": 0.6},
                {"from": "quality", "to": "cost", "weight": 0.5},
                {"from": "speed", "to": "quality", "weight": -0.4},
            ],
        }

    # ─── Phase 1+2: GPU Matrix Build + Preflight ────────────────────────────

    def test_01_build_matrix_gpu(self):
        """Influence matrix built entirely on GPU."""
        inp = self._simple_decision()
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)
        matrix, attractors, preflight = k.build_matrix_gpu(
            inp["factors"], inp["relationships"], inp["options"])

        self.assertEqual(matrix.shape, (3, 3))
        # speed→cost = 0.6 (primary)
        self.assertAlmostEqual(matrix[0, 1], 0.6, places=2)
        # quality→cost = 0.5 (primary)
        self.assertAlmostEqual(matrix[2, 1], 0.5, places=2)
        # speed→quality = -0.4 (primary, no feedback — quality is outgoing)
        self.assertAlmostEqual(matrix[0, 2], -0.4, places=2)
        # cost→speed = 0.24 (leaf feedback: 0.6 * 0.4)
        self.assertAlmostEqual(matrix[1, 0], 0.24, places=2)
        # cost→quality = 0.2 (leaf feedback: 0.5 * 0.4)
        self.assertAlmostEqual(matrix[1, 2], 0.2, places=2)
        # 3 attractors
        self.assertEqual(attractors.shape[0], 3)
        print(f"\n  Matrix:\n{matrix}")
        print(f"  Preflight: {preflight}")

    def test_02_preflight_diagnostics_gpu(self):
        """Preflight diagnostics computed on GPU."""
        inp = self._market_decision()
        k = SBColonyBlockKernel(num_dims=5, max_probes=64, max_attractors=8)
        _, _, preflight = k.build_matrix_gpu(
            inp["factors"], inp["relationships"], inp["options"])

        self.assertIn("min_attractor_distance", preflight)
        self.assertIn("matrix_density", preflight)
        self.assertIn("polarity_ratio", preflight)
        self.assertGreater(preflight["min_attractor_distance"], 0)
        self.assertGreater(preflight["positive_count"], 0)
        print(f"\n  Preflight 5D: {preflight}")

    def test_03_leaf_feedback_on_gpu(self):
        """Leaf factor auto-feedback generated on GPU."""
        inp = {
            "factors": [
                {"name": "a", "range": [0, 100]},
                {"name": "b", "range": [0, 100]},
                {"name": "c", "range": [0, 100]},  # leaf — no outgoing
            ],
            "relationships": [
                {"from": "a", "to": "b", "weight": 0.6},
                {"from": "a", "to": "c", "weight": 0.5},
            ],
            "options": [{"name": "X", "a": 80, "b": 50, "c": 30}],
        }
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)
        matrix, _, _ = k.build_matrix_gpu(
            inp["factors"], inp["relationships"], inp["options"])

        # Primary: a→b = 0.6, a→c = 0.5
        self.assertAlmostEqual(matrix[0, 1], 0.6, places=2)
        self.assertAlmostEqual(matrix[0, 2], 0.5, places=2)
        # Leaf feedback: b→a = 0.6*0.4=0.24, c→a = 0.5*0.4=0.2
        self.assertAlmostEqual(matrix[1, 0], 0.24, places=2)
        self.assertAlmostEqual(matrix[2, 0], 0.2, places=2)
        print(f"\n  Leaf feedback matrix:\n{matrix}")

    # ─── Phase 5+6: GPU Strategy Extraction + Output WM ─────────────────────

    def test_04_extract_strategies_gpu(self):
        """Strategy extraction (denorm + match + Pearson) on GPU."""
        inp = self._simple_decision()
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)

        # Mock clusters
        clusters = [
            {"centroid": np.array([0.8, 0.2, 0.9], dtype=np.float32), "cell_count": 10},
            {"centroid": np.array([0.4, 0.8, 0.6], dtype=np.float32), "cell_count": 8},
            {"centroid": np.array([0.6, 0.5, 0.7], dtype=np.float32), "cell_count": 5},
        ]
        att = np.array([[0.8, 0.2, 0.9], [0.4, 0.8, 0.6], [0.6, 0.5, 0.7]], dtype=np.float32)

        strategies, owm = k.extract_strategies_gpu(
            clusters, inp["factors"], ["speed", "cost", "quality"],
            inp["options"], att, 3)

        self.assertEqual(len(strategies), 3)
        # First cluster should match option A (speed=80, cost=20, quality=90)
        self.assertEqual(strategies[0]["label"], "A")
        # Real values should be denormalized
        self.assertAlmostEqual(strategies[0]["values"]["speed"], 80.0, places=0)
        self.assertAlmostEqual(strategies[0]["values"]["cost"], 20.0, places=0)

        # Output WM should have factors and relationships
        self.assertIsNotNone(owm)
        self.assertEqual(len(owm["factors"]), 3)
        self.assertGreater(len(owm["options"]), 0)
        print(f"\n  Strategies: {[s['label'] for s in strategies]}")
        print(f"  Output WM relationships: {owm['relationships']}")

    def test_05_pearson_gpu(self):
        """Pearson correlation computed on GPU matches Python."""
        from afs.core.output_wm import output_to_world_model

        inp = self._simple_decision()
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)

        clusters = [
            {"centroid": np.array([0.9, 0.1, 0.8], dtype=np.float32), "cell_count": 15},
            {"centroid": np.array([0.2, 0.9, 0.3], dtype=np.float32), "cell_count": 10},
            {"centroid": np.array([0.5, 0.5, 0.5], dtype=np.float32), "cell_count": 5},
        ]
        att = np.array([[0.9, 0.1, 0.8]], dtype=np.float32)

        strategies, owm = k.extract_strategies_gpu(
            clusters, inp["factors"], ["speed", "cost", "quality"],
            inp["options"], att, 3)

        # Compare with Python output_to_world_model
        py_owm = output_to_world_model(strategies)

        if owm and py_owm:
            gpu_rels = {(r["from"], r["to"]): r["strength"]
                        for r in owm["relationships"]}
            py_rels = {(r["from"], r["to"]): r["strength"]
                       for r in py_owm["relationships"]}

            # Same relationship pairs should appear
            for pair in set(gpu_rels.keys()) | set(py_rels.keys()):
                g = gpu_rels.get(pair, 0)
                p = py_rels.get(pair, 0)
                print(f"  {pair}: GPU={g:.4f}, Python={p:.4f}")
                # Allow small numerical difference from GPU float32
                self.assertAlmostEqual(g, p, places=1,
                    msg=f"Pearson mismatch for {pair}: GPU={g}, Python={p}")

    # ─── Full Pipeline: run_full_gpu ─────────────────────────────────────────

    def test_06_full_pipeline_3d(self):
        """Full pipeline: 3D, zero Python in hot path."""
        inp = self._simple_decision()
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)

        t0 = time.perf_counter()
        result = k.run_full_gpu(inp, config={"probe_count": 32, "max_passes": 40})
        total_ms = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        self.assertIn("preflight", result)
        self.assertIn("output_world_model", result)
        self.assertIn("timings", result)

        t = result["timings"]
        print(f"\n  3D full GPU pipeline: {total_ms:.1f}ms")
        print(f"    build_matrix_gpu: {t.get('build_matrix_gpu', 0)*1000:.1f}ms")
        print(f"    calibrate_gpu:    {t.get('calibrate_gpu', 0)*1000:.1f}ms")
        print(f"    initialize:       {t.get('initialize', 0)*1000:.1f}ms")
        print(f"    modulation:       {t.get('modulation', 0)*1000:.1f}ms")
        print(f"    clustering:       {t.get('clustering', 0)*1000:.1f}ms")
        print(f"    extraction_gpu:   {t.get('extraction_gpu', 0)*1000:.1f}ms")
        print(f"  Strategies: {len(result['strategies'])}")
        print(f"  Converged: {result['converged']} at pass {result['pass_count']}")

    def test_07_full_pipeline_5d(self):
        """Full pipeline: 5D market decision."""
        inp = self._market_decision()
        k = SBColonyBlockKernel(num_dims=5, max_probes=64, max_attractors=8)

        t0 = time.perf_counter()
        result = k.run_full_gpu(inp, config={"probe_count": 64, "max_passes": 60})
        total_ms = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        self.assertIsNotNone(result["output_world_model"])

        t = result["timings"]
        print(f"\n  5D full GPU pipeline: {total_ms:.1f}ms")
        for phase, dur in t.items():
            print(f"    {phase}: {dur*1000:.1f}ms")
        print(f"  Strategies: {[s['label'] for s in result['strategies']]}")
        print(f"  Preflight warnings: {result['preflight']['warnings']}")

    def test_08_full_pipeline_10d(self):
        """Full pipeline: 10D stress test."""
        factors = [{"name": f"f{i}", "range": [0, 100]} for i in range(10)]
        relationships = [
            {"from": f"f{i}", "to": f"f{(i+1)%10}", "weight": 0.5}
            for i in range(8)
        ] + [
            {"from": "f0", "to": "f5", "weight": -0.6},
            {"from": "f3", "to": "f7", "weight": -0.4},
        ]
        options = [
            {"name": f"Opt{j}", **{f"f{i}": 20 + j * 25 + i * 5 for i in range(10)}}
            for j in range(3)
        ]

        k = SBColonyBlockKernel(num_dims=10, max_probes=128, max_attractors=8)
        t0 = time.perf_counter()
        result = k.run_full_gpu(
            {"factors": factors, "relationships": relationships, "options": options},
            config={"probe_count": 64, "max_passes": 60})
        total_ms = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  10D full GPU pipeline: {total_ms:.1f}ms")
        print(f"  Strategies: {len(result['strategies'])}")

    def test_09_output_wm_feeds_back(self):
        """Output WM can be fed back into run_full_gpu (self-referential)."""
        inp = self._simple_decision()
        k = SBColonyBlockKernel(num_dims=3, max_probes=64, max_attractors=8)

        # First run
        r1 = k.run_full_gpu(inp, config={"probe_count": 32, "max_passes": 40})
        owm = r1["output_world_model"]
        self.assertIsNotNone(owm)

        # Feed output WM back as input
        k2 = SBColonyBlockKernel(num_dims=len(owm["factors"]),
                                  max_probes=64, max_attractors=8)
        r2 = k2.run_full_gpu(owm, config={"probe_count": 32, "max_passes": 40})

        self.assertGreater(len(r2["strategies"]), 0)
        print(f"\n  Self-referential: {len(r1['strategies'])} → {len(r2['strategies'])} strategies")

    def test_10_benchmark_full_vs_legacy(self):
        """Benchmark: full GPU vs legacy run()."""
        inp = self._market_decision()

        # Full GPU pipeline
        k = SBColonyBlockKernel(num_dims=5, max_probes=128, max_attractors=8)
        cfg = {"probe_count": 64, "max_passes": 60}

        gpu_times = []
        for _ in range(3):
            t0 = time.perf_counter()
            k.run_full_gpu(inp, config=cfg)
            gpu_times.append((time.perf_counter() - t0) * 1000)

        # Legacy pipeline (Python matrix build + extract)
        legacy_times = []
        for _ in range(3):
            t0 = time.perf_counter()
            k.run(inp, config=cfg)
            legacy_times.append((time.perf_counter() - t0) * 1000)

        gpu_med = sorted(gpu_times)[1]
        legacy_med = sorted(legacy_times)[1]
        speedup = legacy_med / gpu_med if gpu_med > 0 else 0

        print(f"\n  Full GPU:  median={gpu_med:.1f}ms ({gpu_times})")
        print(f"  Legacy:    median={legacy_med:.1f}ms ({legacy_times})")
        print(f"  Speedup:   {speedup:.1f}x")


if __name__ == "__main__":
    unittest.main(verbosity=2)
