"""Tests for SBHomeostasisKernel — GPU Homeostasis Calibrator."""

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_homeostasis_kernel import SBHomeostasisKernel, _WGPU_AVAILABLE


@unittest.skipUnless(_WGPU_AVAILABLE, "wgpu not available")
class TestSBHomeostasisKernel(unittest.TestCase):
    """GPU homeostasis calibration tests."""

    def _make_test_input(self, n_dims=5, n_attractors=3):
        """Create a test influence matrix and attractor positions."""
        matrix = np.zeros((n_dims, n_dims), dtype=np.float32)
        # Add some causal relationships
        for i in range(min(n_dims - 1, 5)):
            matrix[i, i + 1] = 0.5 + 0.1 * i
            if i > 0:
                matrix[i, 0] = -0.3  # tradeoff with first factor

        attractors = np.random.rand(n_attractors, n_dims).astype(np.float32)
        # Spread attractors apart
        attractors[0] = np.full(n_dims, 0.2, dtype=np.float32)
        attractors[1] = np.full(n_dims, 0.8, dtype=np.float32)
        if n_attractors >= 3:
            attractors[2] = np.full(n_dims, 0.5, dtype=np.float32)
            attractors[2, :n_dims // 2] = 0.3
            attractors[2, n_dims // 2:] = 0.7

        return matrix, attractors

    def test_01_kernel_init_3d(self):
        """Kernel initializes for 3D input."""
        k = SBHomeostasisKernel(num_dims=3, max_attractors=8)
        self.assertEqual(k.num_dims, 3)
        self.assertEqual(k._cells_per_config, 32)
        self.assertEqual(k._probe_gens, 10)

    def test_02_kernel_init_10d(self):
        """Kernel initializes for 10D input with higher probe count."""
        k = SBHomeostasisKernel(num_dims=10, max_attractors=8)
        self.assertEqual(k.num_dims, 10)
        self.assertEqual(k._cells_per_config, 48)
        self.assertEqual(k._probe_gens, 20)

    def test_03_analyze_matrix(self):
        """Field analysis computes correct statistics."""
        k = SBHomeostasisKernel(num_dims=5)
        matrix, _ = self._make_test_input(5)
        analysis = k._analyze_matrix(matrix)
        self.assertIn("polarity_ratio", analysis)
        self.assertIn("mean_abs", analysis)
        self.assertGreater(analysis["mean_abs"], 0)

    def test_04_sample_configs_wide(self):
        """Wide scan generates correct number of configs."""
        k = SBHomeostasisKernel(num_dims=5)
        analysis = {"polarity_ratio": 0.3, "mean_abs": 0.4, "spread": 0.5}
        space = k._param_space(analysis)
        cfgs = k._sample_configs(space, None, 1.0, 4)
        self.assertEqual(len(cfgs), 4)
        for c in cfgs:
            self.assertIn("force_scale", c)
            self.assertIn("gravity_strength", c)
            self.assertIn("label", c)

    def test_05_sample_configs_narrow(self):
        """Narrow bisection generates midpoint + perturbations."""
        k = SBHomeostasisKernel(num_dims=5)
        analysis = {"polarity_ratio": 0.3, "mean_abs": 0.4, "spread": 0.5}
        space = k._param_space(analysis)
        centers = k._sample_configs(space, None, 1.0, 4)[:2]
        narrow = k._sample_configs(space, centers, 0.5, 3)
        self.assertEqual(len(narrow), 3)
        self.assertEqual(narrow[0]["label"], "bisect-mid")

    def test_06_run_stage_gpu_3d(self):
        """Stage 1 runs 4 probes on GPU for 3D input."""
        k = SBHomeostasisKernel(num_dims=3, max_attractors=8)
        matrix, attractors = self._make_test_input(3, 3)
        analysis = k._analyze_matrix(matrix)
        space = k._param_space(analysis)
        cfgs = k._sample_configs(space, None, 1.0, 4)
        for c in cfgs:
            c["gravity_range"] = 3.0

        results = k._run_stage_gpu(cfgs, matrix, attractors)
        self.assertEqual(len(results), 4)
        for r in results:
            self.assertIn("score", r)
            self.assertIn("label", r)
            self.assertGreaterEqual(r["score"], 0)

    def test_07_run_stage_gpu_10d(self):
        """Stage 1 runs 4 probes on GPU for 10D input."""
        k = SBHomeostasisKernel(num_dims=10, max_attractors=8)
        matrix, attractors = self._make_test_input(10, 4)
        analysis = k._analyze_matrix(matrix)
        space = k._param_space(analysis)
        cfgs = k._sample_configs(space, None, 1.0, 4)
        for c in cfgs:
            c["gravity_range"] = 5.0

        results = k._run_stage_gpu(cfgs, matrix, attractors)
        self.assertEqual(len(results), 4)
        # At least one should have score > 0
        max_score = max(r["score"] for r in results)
        self.assertGreater(max_score, 0)

    def test_08_full_calibrate_3d(self):
        """Full calibration pipeline — 3D, 7 probes, 2 GPU dispatches."""
        k = SBHomeostasisKernel(num_dims=3, max_attractors=8)
        matrix, attractors = self._make_test_input(3, 3)
        base_cfg = {"gravity_range": 3.0}

        t0 = time.perf_counter()
        result = k.calibrate(matrix, attractors, base_cfg)
        total_ms = (time.perf_counter() - t0) * 1000

        self.assertIn("winning_config", result)
        self.assertIn("calibration", result)
        cal = result["calibration"]
        self.assertEqual(cal["method"], "gpu-adaptive-bisection")
        self.assertEqual(cal["probe_count"], 7)
        self.assertGreater(cal["score"], 0)

        wc = result["winning_config"]
        self.assertIn("force_scale", wc)
        self.assertIn("gravity_strength", wc)

        print(f"\n3D calibration: {total_ms:.1f}ms total "
              f"(stage1={cal['stage1_ms']:.1f}ms, stage2={cal['stage2_ms']:.1f}ms)")
        print(f"  Winner: {cal['label']} (score={cal['score']})")
        print(f"  Config: force={wc['force_scale']:.3f}, "
              f"gravity={wc['gravity_strength']:.2f}, "
              f"noise={wc['thermal_noise']:.3f}")

    def test_09_full_calibrate_10d(self):
        """Full calibration pipeline — 10D, 7 probes, 2 GPU dispatches."""
        k = SBHomeostasisKernel(num_dims=10, max_attractors=8)
        matrix, attractors = self._make_test_input(10, 4)
        base_cfg = {"gravity_range": 5.0}

        t0 = time.perf_counter()
        result = k.calibrate(matrix, attractors, base_cfg)
        total_ms = (time.perf_counter() - t0) * 1000

        cal = result["calibration"]
        self.assertEqual(cal["probe_count"], 7)
        self.assertGreater(cal["score"], 0)

        print(f"\n10D calibration: {total_ms:.1f}ms total "
              f"(stage1={cal['stage1_ms']:.1f}ms, stage2={cal['stage2_ms']:.1f}ms)")
        print(f"  Winner: {cal['label']} (score={cal['score']})")
        scores_str = ', '.join(f"{r['label']}={r['score']}" for r in cal['all_scores'])
        print(f"  All scores: {scores_str}")

    def test_10_calibrate_deterministic_ranking(self):
        """Same input should produce consistent winner ranking across runs."""
        k = SBHomeostasisKernel(num_dims=5, max_attractors=8)
        matrix, attractors = self._make_test_input(5, 3)
        base_cfg = {"gravity_range": 3.5}

        # Run twice
        r1 = k.calibrate(matrix, attractors, base_cfg)
        r2 = k.calibrate(matrix, attractors, base_cfg)

        # Top scorer should be from same stage (wide or narrow)
        # Exact scores may vary due to random probe init, but ranking should be stable
        top1 = r1["calibration"]["all_scores"][0]["label"]
        top2 = r2["calibration"]["all_scores"][0]["label"]
        # Both should have non-zero scores
        self.assertGreater(r1["calibration"]["score"], 0)
        self.assertGreater(r2["calibration"]["score"], 0)
        print(f"\n  Run 1 winner: {top1} ({r1['calibration']['score']})")
        print(f"  Run 2 winner: {top2} ({r2['calibration']['score']})")

    def test_11_calibrate_speed_benchmark(self):
        """Benchmark: 3 runs each of 3D and 10D calibration."""
        for dims, n_att in [(3, 3), (10, 4)]:
            k = SBHomeostasisKernel(num_dims=dims, max_attractors=8)
            matrix, attractors = self._make_test_input(dims, n_att)
            base_cfg = {"gravity_range": max(2.0, min(6.0, 1.5 * np.sqrt(dims)))}

            times = []
            for _ in range(3):
                t0 = time.perf_counter()
                k.calibrate(matrix, attractors, base_cfg)
                times.append((time.perf_counter() - t0) * 1000)

            median = sorted(times)[1]
            print(f"\n  {dims}D GPU homeostasis: median={median:.1f}ms "
                  f"(runs: {', '.join(f'{t:.1f}ms' for t in times)})")


@unittest.skipUnless(_WGPU_AVAILABLE, "wgpu not available")
class TestSBHomeostasisIntegration(unittest.TestCase):
    """Integration: wire GPU homeostasis into Colony Sound Block kernel."""

    def test_12_wire_into_colony_kernel(self):
        """GPU homeostasis produces config that Colony Sound Block accepts."""
        from afs.gpu.sb_colony_block_kernel import SBColonyBlockKernel, gpu_available
        if not gpu_available():
            self.skipTest("GPU not available")

        # Build test input
        n_dims = 5
        factors = [{"name": f"f{i}", "range": [0, 100]} for i in range(n_dims)]
        relationships = [
            {"from": "f0", "to": "f1", "weight": 0.6},
            {"from": "f1", "to": "f2", "weight": 0.5},
            {"from": "f0", "to": "f2", "weight": -0.4},
        ]
        options = [
            {"name": "A", "f0": 80, "f1": 20, "f2": 90, "f3": 40, "f4": 70},
            {"name": "B", "f0": 30, "f1": 70, "f2": 40, "f3": 80, "f4": 30},
        ]

        # Run homeostasis first
        hk = SBHomeostasisKernel(num_dims=n_dims, max_attractors=8)

        # Build matrix + attractors (same as colony kernel does)
        name_to_idx = {f["name"]: i for i, f in enumerate(factors)}
        matrix = np.zeros((n_dims, n_dims), dtype=np.float32)
        for rel in relationships:
            i = name_to_idx.get(rel.get("from"))
            j = name_to_idx.get(rel.get("to"))
            if i is not None and j is not None:
                matrix[i, j] = rel.get("weight", 0.5)

        attractors = np.full((len(options), n_dims), 0.5, dtype=np.float32)
        for oi, opt in enumerate(options):
            if isinstance(opt, str):
                continue
            for di, f in enumerate(factors):
                fn = f["name"]
                if fn in opt and fn != "name":
                    lo, hi = f.get("range", [0, 100])
                    if hi != lo:
                        attractors[oi, di] = np.clip((opt[fn] - lo) / (hi - lo), 0, 1)

        base_cfg = {"gravity_range": max(2.0, min(6.0, 1.5 * np.sqrt(n_dims)))}
        cal_result = hk.calibrate(matrix, attractors, base_cfg)
        wc = cal_result["winning_config"]

        # Now run colony with calibrated config
        ck = SBColonyBlockKernel(num_dims=n_dims, max_probes=64, max_attractors=8)
        colony_result = ck.run(
            {"factors": factors, "relationships": relationships, "options": options},
            config={
                "probe_count": 32,
                "max_passes": 40,
                "force_scale": wc["force_scale"],
                "gravity_strength": wc["gravity_strength"],
                "gravity_range": wc["gravity_range"],
                "gravity_sharpness": wc["gravity_sharpness"],
                "thermal_noise": wc["thermal_noise"],
                "tradeoff_amplifier": wc["tradeoff_amplifier"],
            }
        )

        self.assertGreater(len(colony_result["strategies"]), 0)
        print(f"\n  GPU homeostasis → colony: {len(colony_result['strategies'])} strategies")
        print(f"  Calibrated config: force={wc['force_scale']:.3f}, "
              f"gravity={wc['gravity_strength']:.2f}")
        print(f"  Colony: {colony_result['total_ms']:.1f}ms, "
              f"{colony_result['pass_count']} passes, "
              f"converged={colony_result['converged']}")


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