"""
Colony Sound Block — GPU Kernel Tests

Tests the SBColonyBlockKernel against reference decisions and validates:
1. GPU pipeline produces strategies
2. Convergence detection works
3. Physics equivalence with production colony engine
4. Performance targets met

Run:
    cd afs-engine
    python3 tests/test_colony_block.py
    python3 tests/test_colony_block.py -v          # verbose
    python3 tests/test_colony_block.py TestColonyBlockGPU.test_synth_decision  # single test
"""

import unittest
import numpy as np
import time
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from afs.gpu.sb_colony_block_kernel import (
    SBColonyBlockKernel, gpu_available, DEFAULT_CONFIG,
    compute_gravity_ramp, compute_settling_damp
)

# ─── Reference Decisions ──────────────────────────────────────────────────────

SYNTH_DECISION = {
    "question": "Which synthesizer should I buy?",
    "options": [
        {"name": "Prophet 6", "price": 3500, "analog_warmth": 85, "versatility": 45,
         "build_quality": 80, "resale_value": 75, "learning_curve": 35,
         "patch_memory": 50, "effects_quality": 60, "mod_routing": 55, "portability": 40},
        {"name": "OB-8x", "price": 4999, "analog_warmth": 92, "versatility": 55,
         "build_quality": 90, "resale_value": 85, "learning_curve": 50,
         "patch_memory": 80, "effects_quality": 70, "mod_routing": 65, "portability": 30},
        {"name": "Nord Stage 4", "price": 4000, "analog_warmth": 30, "versatility": 95,
         "build_quality": 85, "resale_value": 60, "learning_curve": 70,
         "patch_memory": 95, "effects_quality": 80, "mod_routing": 40, "portability": 85},
    ],
    "factors": [
        {"name": "price", "range": [2000, 6000], "unit": "$", "goal": "minimize"},
        {"name": "analog_warmth", "range": [0, 100], "unit": "score"},
        {"name": "versatility", "range": [0, 100], "unit": "score"},
        {"name": "build_quality", "range": [0, 100], "unit": "score"},
        {"name": "resale_value", "range": [0, 100], "unit": "%", "goal": "maximize"},
        {"name": "learning_curve", "range": [0, 100], "unit": "hours", "goal": "minimize"},
        {"name": "patch_memory", "range": [0, 100], "unit": "score"},
        {"name": "effects_quality", "range": [0, 100], "unit": "score"},
        {"name": "mod_routing", "range": [0, 100], "unit": "score"},
        {"name": "portability", "range": [0, 100], "unit": "score"},
    ],
    "relationships": [
        {"from": "price", "to": "build_quality", "strength": 0.8},
        {"from": "analog_warmth", "to": "versatility", "strength": -0.6},
        {"from": "learning_curve", "to": "versatility", "strength": 0.5},
        {"from": "build_quality", "to": "resale_value", "strength": 0.7},
        {"from": "versatility", "to": "effects_quality", "strength": 0.6},
        {"from": "analog_warmth", "to": "mod_routing", "strength": 0.5},
        {"from": "portability", "to": "build_quality", "strength": -0.4},
        {"from": "patch_memory", "to": "learning_curve", "strength": -0.3},
        {"from": "price", "to": "analog_warmth", "strength": 0.5},
        {"from": "mod_routing", "to": "versatility", "strength": 0.4},
    ],
}

SIMPLE_DECISION = {
    "question": "Buy or rent a home?",
    "options": [
        {"name": "Buy", "monthly_cost": 75, "equity_building": 90, "flexibility": 20,
         "maintenance_burden": 80, "tax_benefits": 70, "upfront_cost": 90,
         "location_freedom": 15, "long_term_wealth": 85},
        {"name": "Rent", "monthly_cost": 40, "equity_building": 5, "flexibility": 90,
         "maintenance_burden": 10, "tax_benefits": 20, "upfront_cost": 10,
         "location_freedom": 95, "long_term_wealth": 20},
    ],
    "factors": [
        {"name": "monthly_cost", "range": [0, 100], "goal": "minimize"},
        {"name": "equity_building", "range": [0, 100], "goal": "maximize"},
        {"name": "flexibility", "range": [0, 100]},
        {"name": "maintenance_burden", "range": [0, 100], "goal": "minimize"},
        {"name": "tax_benefits", "range": [0, 100], "goal": "maximize"},
        {"name": "upfront_cost", "range": [0, 100], "goal": "minimize"},
        {"name": "location_freedom", "range": [0, 100]},
        {"name": "long_term_wealth", "range": [0, 100], "goal": "maximize"},
    ],
    "relationships": [
        {"from": "monthly_cost", "to": "equity_building", "strength": 0.6},
        {"from": "equity_building", "to": "long_term_wealth", "strength": 0.8},
        {"from": "flexibility", "to": "location_freedom", "strength": 0.7},
        {"from": "maintenance_burden", "to": "monthly_cost", "strength": 0.5},
        {"from": "upfront_cost", "to": "tax_benefits", "strength": 0.4},
        {"from": "location_freedom", "to": "flexibility", "strength": 0.3},
    ],
}

# Minimal 3-factor decision for fast tests
MINIMAL_DECISION = {
    "question": "A or B?",
    "options": [
        {"name": "Option A", "speed": 90, "cost": 30, "quality": 70},
        {"name": "Option B", "speed": 40, "cost": 80, "quality": 90},
    ],
    "factors": [
        {"name": "speed", "range": [0, 100]},
        {"name": "cost", "range": [0, 100], "goal": "minimize"},
        {"name": "quality", "range": [0, 100], "goal": "maximize"},
    ],
    "relationships": [
        {"from": "speed", "to": "cost", "strength": 0.6},
        {"from": "cost", "to": "quality", "strength": -0.5},
        {"from": "quality", "to": "speed", "strength": -0.3},
    ],
}


# ─── Helper Tests (no GPU required) ──────────────────────────────────────────

class TestColonyBlockHelpers(unittest.TestCase):
    """Test gravity ramp, settling, and config — no GPU needed."""

    def test_gravity_ramp_boundaries(self):
        """Gravity ramp: 5% at start, 100% at end."""
        self.assertAlmostEqual(compute_gravity_ramp(0, 100), 0.05, places=2)
        self.assertAlmostEqual(compute_gravity_ramp(100, 100), 1.0, places=2)

    def test_gravity_ramp_monotonic(self):
        """Gravity ramp should be monotonically increasing."""
        prev = 0
        for i in range(101):
            val = compute_gravity_ramp(i, 100)
            self.assertGreaterEqual(val, prev)
            prev = val

    def test_settling_damp_before_threshold(self):
        """No settling before 85%."""
        self.assertEqual(compute_settling_damp(50, 100), 1.0)
        self.assertEqual(compute_settling_damp(84, 100), 1.0)

    def test_settling_damp_at_end(self):
        """Full settling at 100%."""
        val = compute_settling_damp(100, 100)
        self.assertAlmostEqual(val, 0.3, places=2)

    def test_default_config_complete(self):
        """All expected keys present in DEFAULT_CONFIG."""
        required = ["max_passes", "convergence_threshold", "min_passes",
                     "probe_count", "gravity_strength", "force_scale",
                     "momentum", "max_amplitude", "dt_norm"]
        for key in required:
            self.assertIn(key, DEFAULT_CONFIG, f"Missing key: {key}")


# ─── GPU Tests ────────────────────────────────────────────────────────────────

@unittest.skipUnless(gpu_available(), "No GPU available")
class TestColonyBlockGPU(unittest.TestCase):
    """GPU kernel tests — require wgpu + Vulkan/Metal/DX12."""

    def test_kernel_init(self):
        """Kernel initializes with GPU device."""
        kernel = SBColonyBlockKernel(num_dims=10)
        info = kernel.adapter_info
        self.assertIn("device", info)
        print(f"  GPU: {info.get('device', 'unknown')}")

    def test_minimal_decision(self):
        """3-factor decision produces strategies."""
        kernel = SBColonyBlockKernel(num_dims=3)
        result = kernel.run(MINIMAL_DECISION, config={"probe_count": 16, "max_passes": 30})

        self.assertIn("strategies", result)
        self.assertGreater(len(result["strategies"]), 0)
        self.assertIn("summary", result)
        self.assertEqual(result["summary"]["dimensions"], 3)

        print(f"  Strategies: {len(result['strategies'])}")
        print(f"  Passes: {result['pass_count']}, converged: {result['converged']}")
        print(f"  Time: {result['total_ms']:.1f}ms")

    def test_synth_decision(self):
        """10-factor synth decision — the reference benchmark."""
        kernel = SBColonyBlockKernel(num_dims=10)
        result = kernel.run(SYNTH_DECISION, config={"probe_count": 32, "max_passes": 60})

        strategies = result["strategies"]
        self.assertGreater(len(strategies), 0)
        self.assertLessEqual(len(strategies), 32)

        # Each strategy should have values and concept_weights
        for s in strategies:
            self.assertIn("label", s)
            self.assertIn("values", s)
            self.assertIn("concept_weights", s)
            self.assertIn("cell_count", s)
            self.assertEqual(len(s["concept_weights"]), 10)

        # Summary
        summary = result["summary"]
        self.assertEqual(summary["dimensions"], 10)
        self.assertGreater(summary["total_strategies"], 0)

        print(f"  Strategies: {len(strategies)}")
        for s in strategies:
            print(f"    {s['label']}: {s['cell_count']} probes")
        print(f"  Passes: {result['pass_count']}, converged: {result['converged']}")
        print(f"  Time: {result['total_ms']:.1f}ms")

    def test_simple_decision(self):
        """8-factor buy/rent decision."""
        kernel = SBColonyBlockKernel(num_dims=8)
        result = kernel.run(SIMPLE_DECISION, config={"probe_count": 32, "max_passes": 40})

        self.assertGreater(len(result["strategies"]), 0)
        print(f"  Strategies: {len(result['strategies'])}")
        print(f"  Passes: {result['pass_count']}, converged: {result['converged']}")
        print(f"  Time: {result['total_ms']:.1f}ms")

    def test_convergence_detection(self):
        """Convergence should stop early when probes stabilize."""
        kernel = SBColonyBlockKernel(num_dims=3)
        result = kernel.run(MINIMAL_DECISION, config={
            "probe_count": 16,
            "max_passes": 60,
            "convergence_threshold": 0.01,  # loose threshold
            "min_passes": 4,
        })

        # With a loose threshold on a simple problem, should converge early
        if result["converged"]:
            self.assertLess(result["pass_count"], 60)
            print(f"  Converged at pass {result['pass_count']}")
        else:
            print(f"  Did not converge in {result['pass_count']} passes")

        # Convergence history should exist
        self.assertGreater(len(result["convergence_history"]), 0)
        print(f"  History: {[f'{d:.4f}' for d in result['convergence_history']]}")

    def test_convergence_history_decreasing(self):
        """After min_passes, convergence deltas should generally decrease."""
        kernel = SBColonyBlockKernel(num_dims=3)
        result = kernel.run(MINIMAL_DECISION, config={
            "probe_count": 16,
            "max_passes": 60,
            "min_passes": 8,
            "convergence_check_interval": 2,
        })

        history = result["convergence_history"]
        if len(history) >= 4:
            # Velocity energy may oscillate before settling (probes accelerate
            # toward basins, then decelerate). The final value should be smaller
            # than the peak, not necessarily monotonically decreasing.
            final = history[-1]
            peak = max(history)
            self.assertLessEqual(final, peak,
                "Final velocity energy should be <= peak energy")

    def test_high_probe_count(self):
        """64 probes should still work and produce strategies."""
        kernel = SBColonyBlockKernel(num_dims=10, max_probes=64)
        result = kernel.run(SYNTH_DECISION, config={"probe_count": 64, "max_passes": 30})

        self.assertGreater(len(result["strategies"]), 0)
        print(f"  64 probes: {len(result['strategies'])} strategies, "
              f"{result['total_ms']:.1f}ms")

    def test_strategy_values_in_range(self):
        """All strategy values should fall within factor ranges."""
        kernel = SBColonyBlockKernel(num_dims=10)
        result = kernel.run(SYNTH_DECISION, config={"probe_count": 32, "max_passes": 40})

        for s in result["strategies"]:
            for f in SYNTH_DECISION["factors"]:
                fn = f["name"]
                lo, hi = f["range"]
                val = s["values"][fn]
                self.assertGreaterEqual(val, lo - 0.01,
                    f"{fn}={val} below range [{lo}, {hi}]")
                self.assertLessEqual(val, hi + 0.01,
                    f"{fn}={val} above range [{lo}, {hi}]")

    def test_concept_weights_normalized(self):
        """Concept weights should be in [0, 1]."""
        kernel = SBColonyBlockKernel(num_dims=10)
        result = kernel.run(SYNTH_DECISION, config={"probe_count": 32, "max_passes": 40})

        for s in result["strategies"]:
            for fn, w in s["concept_weights"].items():
                self.assertGreaterEqual(w, 0.0, f"{fn} weight {w} < 0")
                self.assertLessEqual(w, 1.0, f"{fn} weight {w} > 1")

    def test_influence_matrix_returned(self):
        """Result should include the influence matrix."""
        kernel = SBColonyBlockKernel(num_dims=3)
        result = kernel.run(MINIMAL_DECISION)

        self.assertIn("influence_matrix", result)
        matrix = result["influence_matrix"]
        self.assertEqual(len(matrix), 3)
        self.assertEqual(len(matrix[0]), 3)

    def test_fast_decision_preset(self):
        """Fast preset should complete quickly."""
        kernel = SBColonyBlockKernel(num_dims=10)
        t0 = time.perf_counter()
        result = kernel.run(SYNTH_DECISION, config={
            "probe_count": 16,
            "max_passes": 20,
            "convergence_threshold": 0.005,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"  Fast preset: {elapsed:.1f}ms, "
              f"{len(result['strategies'])} strategies, "
              f"{result['pass_count']} passes")

    def test_repeated_runs_consistent(self):
        """Multiple runs should produce similar strategy counts."""
        kernel = SBColonyBlockKernel(num_dims=10)
        counts = []
        for _ in range(5):
            result = kernel.run(SYNTH_DECISION, config={
                "probe_count": 32, "max_passes": 40})
            counts.append(len(result["strategies"]))

        # Strategy count should be relatively stable (±2)
        self.assertLessEqual(max(counts) - min(counts), 3,
            f"Strategy counts too variable: {counts}")
        print(f"  Run counts: {counts}")


# ─── Benchmark ────────────────────────────────────────────────────────────────

@unittest.skipUnless(gpu_available(), "No GPU available")
class TestColonyBlockBenchmark(unittest.TestCase):
    """Performance benchmarks."""

    def test_benchmark_synth(self):
        """Benchmark: 10-factor synth decision, 32 probes."""
        kernel = SBColonyBlockKernel(num_dims=10)

        # Warmup
        kernel.run(SYNTH_DECISION, config={"probe_count": 32, "max_passes": 30})

        # Timed runs
        times = []
        for _ in range(5):
            t0 = time.perf_counter()
            result = kernel.run(SYNTH_DECISION, config={"probe_count": 32, "max_passes": 30})
            times.append((time.perf_counter() - t0) * 1000)

        avg = np.mean(times)
        std = np.std(times)
        print(f"\n  Colony Sound Block benchmark (32 probes, 10D):")
        print(f"    Mean: {avg:.1f}ms ± {std:.1f}ms")
        print(f"    Min: {min(times):.1f}ms, Max: {max(times):.1f}ms")
        print(f"    vs Production colony: 169.8ms")
        print(f"    Speedup: {169.8 / avg:.1f}x")

    def test_benchmark_fast_preset(self):
        """Benchmark: fast preset (16 probes, 20 passes)."""
        kernel = SBColonyBlockKernel(num_dims=10)

        # Warmup
        kernel.run(SYNTH_DECISION, config={"probe_count": 16, "max_passes": 20})

        times = []
        for _ in range(10):
            t0 = time.perf_counter()
            result = kernel.run(SYNTH_DECISION, config={
                "probe_count": 16, "max_passes": 20,
                "convergence_threshold": 0.005})
            times.append((time.perf_counter() - t0) * 1000)

        avg = np.mean(times)
        print(f"\n  Fast preset benchmark (16 probes, 10D):")
        print(f"    Mean: {avg:.1f}ms ± {np.std(times):.1f}ms")
        print(f"    Target: <10ms")


# ─── Scale Tests ──────────────────────────────────────────────────────────────

# 50-factor high-dimensional decision for scale testing
SCALE_DECISION_50D = {
    "question": "Scale test: 50 dimensions",
    "options": [
        {"name": f"Option {i}", **{f"f{d}": np.random.randint(10, 90)
         for d in range(50)}}
        for i in range(10)
    ],
    "factors": [
        {"name": f"f{d}", "range": [0, 100]} for d in range(50)
    ],
    "relationships": [
        {"from": f"f{i}", "to": f"f{(i+1) % 50}", "strength": 0.5 * (-1 if i % 3 == 0 else 1)}
        for i in range(50)
    ],
}


@unittest.skipUnless(gpu_available(), "No GPU available")
class TestColonyBlockScale(unittest.TestCase):
    """Extreme scale tests — push probe count, pass count, and dimensions."""

    def test_500_probes_150_passes(self):
        """500 probes × 150 passes — production colony equivalent."""
        kernel = SBColonyBlockKernel(num_dims=10, max_probes=512)
        t0 = time.perf_counter()
        result = kernel.run(SYNTH_DECISION, config={
            "probe_count": 500, "max_passes": 150,
            "convergence_threshold": 0.0001,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  500 probes × 150 passes (10D):")
        print(f"    Time: {elapsed:.1f}ms")
        print(f"    Strategies: {len(result['strategies'])}")
        print(f"    Passes: {result['pass_count']}, converged: {result['converged']}")

    def test_1000_probes(self):
        """1,000 probes — high-resolution basins."""
        kernel = SBColonyBlockKernel(num_dims=10, max_probes=1024)
        t0 = time.perf_counter()
        result = kernel.run(SYNTH_DECISION, config={
            "probe_count": 1000, "max_passes": 100,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  1,000 probes × 100 passes (10D):")
        print(f"    Time: {elapsed:.1f}ms")
        print(f"    Strategies: {len(result['strategies'])}")

    def test_10000_probes(self):
        """10,000 probes — massive resolution."""
        kernel = SBColonyBlockKernel(num_dims=10, max_probes=10240)
        t0 = time.perf_counter()
        result = kernel.run(SYNTH_DECISION, config={
            "probe_count": 10000, "max_passes": 100,
            "convergence_check_interval": 10,
            "fused_passes": 10,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  10,000 probes × 100 passes (10D):")
        print(f"    Time: {elapsed:.1f}ms")
        print(f"    Strategies: {len(result['strategies'])}")
        print(f"    Passes: {result['pass_count']}, converged: {result['converged']}")

    def test_10000_probes_10000_passes(self):
        """10,000 probes × 10,000 passes — extreme resolution."""
        kernel = SBColonyBlockKernel(num_dims=10, max_probes=10240)
        t0 = time.perf_counter()
        result = kernel.run(SYNTH_DECISION, config={
            "probe_count": 10000, "max_passes": 10000,
            "convergence_check_interval": 100,
            "fused_passes": 50,
            "convergence_threshold": 0.00001,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  10,000 probes × 10,000 passes (10D):")
        print(f"    Time: {elapsed:.1f}ms ({elapsed/1000:.1f}s)")
        print(f"    Strategies: {len(result['strategies'])}")
        print(f"    Passes: {result['pass_count']}, converged: {result['converged']}")
        print(f"    vs Production (500×150): 169.8ms")

    def test_50_dimensions(self):
        """50 dimensions — high-dimensional factor space."""
        kernel = SBColonyBlockKernel(num_dims=50, max_probes=512, max_attractors=16)
        t0 = time.perf_counter()
        result = kernel.run(SCALE_DECISION_50D, config={
            "probe_count": 200, "max_passes": 100,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        self.assertEqual(result["summary"]["dimensions"], 50)
        print(f"\n  200 probes × 100 passes (50D):")
        print(f"    Time: {elapsed:.1f}ms")
        print(f"    Strategies: {len(result['strategies'])}")

    def test_50_dimensions_1000_probes(self):
        """50D × 1000 probes — high resolution in high dimensions."""
        kernel = SBColonyBlockKernel(num_dims=50, max_probes=1024, max_attractors=16)
        t0 = time.perf_counter()
        result = kernel.run(SCALE_DECISION_50D, config={
            "probe_count": 1000, "max_passes": 200,
            "convergence_check_interval": 20,
            "fused_passes": 10,
        })
        elapsed = (time.perf_counter() - t0) * 1000

        self.assertGreater(len(result["strategies"]), 0)
        print(f"\n  1,000 probes × 200 passes (50D):")
        print(f"    Time: {elapsed:.1f}ms")
        print(f"    Strategies: {len(result['strategies'])}")

    def test_vram_estimate(self):
        """Print VRAM usage estimates for various scales."""
        print("\n  VRAM Usage Estimates:")
        print(f"  {'Probes':>10} {'Dims':>6} {'Passes':>8} {'VRAM':>10} {'Est. Time':>12}")
        print(f"  {'─'*10} {'─'*6} {'─'*8} {'─'*10} {'─'*12}")

        configs = [
            (32, 10, 60, "fast"),
            (500, 10, 150, "production"),
            (1000, 10, 200, "high-res"),
            (10000, 10, 1000, "extreme"),
            (10000, 10, 10000, "ultra"),
            (100000, 10, 100, "massive"),
            (200, 50, 100, "50D"),
            (1000, 50, 200, "50D high"),
            (10000, 100, 500, "100D extreme"),
            (1000000, 10, 50, "1M probes"),
        ]

        for probes, dims, passes, label in configs:
            # 5 buffers × P × D × 4 bytes + P × 4 (deltas) + D² × 4 (matrix)
            vram_bytes = 5 * probes * dims * 4 + probes * 4 + dims * dims * 4
            vram_mb = vram_bytes / (1024 * 1024)
            print(f"  {probes:>10,} {dims:>6} {passes:>8,} {vram_mb:>8.1f} MB   ({label})")


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