"""P9 tests — child colony spawning and merge-back into parent strategies.

Tests verify:
  1. seed_centroid places cells near the requested location
  2. seed_spread overrides default spawn radius
  3. Child colonies converge faster than parent (tighter focus)
  4. _run_child_sims produces sub_strategies on each parent strategy
  5. Sub-strategies are in real units (not normalized)
  6. Tiny parent clusters get empty sub_strategies (no wasted compute)
"""

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

import numpy as np
from afs.core import FieldDefinition, AfsEngine, AfsSimulationConfig
from afs.analysis import detect_clusters
from afs.mcp.server import _run_child_sims, _run_generations, _interpret_clusters


# ─── Helpers ──────────────────────────────────────────────────────────────────

def _make_field_with_attractors(n_dims: int = 6) -> FieldDefinition:
    """Minimal attractor-based field for testing child sims."""
    attractors = []
    dim_names = [f"dim_{i}" for i in range(n_dims)]

    # Two attractors at opposite ends of dim_0
    attractors.append({
        "name": "OptionA",
        "position": {d: (0.2 if d == "dim_0" else 0.5) * 100 for d in dim_names},
        "strength": 0.7,
    })
    attractors.append({
        "name": "OptionB",
        "position": {d: (0.8 if d == "dim_0" else 0.5) * 100 for d in dim_names},
        "strength": 0.7,
    })

    field = FieldDefinition(metadata={
        "worldModel": {
            "attractors": attractors,
        },
        "problem": "test",
    })

    for name in dim_names:
        dim_id = field.add_dimension(
            name=name, category="test",
            display_range=[0, 100], display_unit="score",
        )
        field.promote_dimension(dim_id)

    # Diagonal similarity — mild repulsion between dims
    mat = np.eye(n_dims) + np.random.uniform(-0.3, 0.3, (n_dims, n_dims))
    mat = (mat + mat.T) / 2
    np.fill_diagonal(mat, 1.0)
    field.update_similarity_matrix(mat)
    return field


def _small_config(initial_cells: int = 48, max_pop: int = 100) -> AfsSimulationConfig:
    return AfsSimulationConfig.from_dict({
        "engine": {
            "initial_cell_count": initial_cells,
            "max_population": max_pop,
        },
    })


# ─── Tests ────────────────────────────────────────────────────────────────────

def test_seed_centroid_places_cells_near_target():
    """Cells should start close to the requested centroid."""
    field = _make_field_with_attractors(6)
    config = _small_config()
    n = field.get_dimension_count()

    target = np.full(n, 0.25)  # well away from default center (0.5)
    engine = AfsEngine(field, config, use_gpu=False)
    engine.initialize(seed_centroid=target, seed_spread=0.05)

    positions = engine.get_cell_positions()
    assert positions.shape[0] > 0, "No cells spawned"

    centroid_actual = positions.mean(axis=0)
    dist = float(np.linalg.norm(centroid_actual - target))

    # Should land within 0.1 of the target in N-D space
    assert dist < 0.1, f"Cells spawned too far from target: dist={dist:.3f}"
    print(f"  seed_centroid test: distance to target = {dist:.4f}")


def test_seed_spread_controls_radius():
    """Tight spread should produce a more concentrated initial population."""
    field = _make_field_with_attractors(6)
    config = _small_config(initial_cells=64)
    n = field.get_dimension_count()
    target = np.full(n, 0.5)

    # Wide spread
    engine_wide = AfsEngine(field, config, use_gpu=False)
    engine_wide.initialize(seed_centroid=target, seed_spread=0.25)
    pos_wide = engine_wide.get_cell_positions()
    std_wide = float(np.std(pos_wide))

    # Tight spread
    engine_tight = AfsEngine(field, config, use_gpu=False)
    engine_tight.initialize(seed_centroid=target, seed_spread=0.04)
    pos_tight = engine_tight.get_cell_positions()
    std_tight = float(np.std(pos_tight))

    assert std_tight < std_wide, (
        f"Tight spread should produce lower std: tight={std_tight:.4f}, wide={std_wide:.4f}")
    print(f"  std wide={std_wide:.4f}, tight={std_tight:.4f} (ratio {std_wide/std_tight:.1f}x)")


def test_default_initialize_unchanged():
    """seed_centroid=None should preserve existing initialization behaviour."""
    field = _make_field_with_attractors(6)
    config = _small_config()

    engine_default = AfsEngine(field, config, use_gpu=False)
    engine_default.initialize()

    assert len(engine_default.cells) == config.engine.initial_cell_count
    positions = engine_default.get_cell_positions()
    # Default center is the mean of attractors (~0.5 in all dims)
    mean_pos = positions.mean(axis=0)
    assert np.all(mean_pos > 0.05) and np.all(mean_pos < 0.95), \
        f"Default init placed cells at unexpected location: {mean_pos}"


def test_child_colony_runs_and_produces_clusters():
    """A child engine seeded at a cluster centroid should find sub-clusters."""
    field = _make_field_with_attractors(6)
    config = _small_config(initial_cells=48, max_pop=100)
    n = field.get_dimension_count()

    centroid = np.array([0.2, 0.5, 0.5, 0.5, 0.5, 0.5])  # near OptionA
    child = AfsEngine(field, config, use_gpu=False)
    child.initialize(seed_centroid=centroid, seed_spread=0.07)
    child.total_ticks = 20 * 100
    _run_generations(child, generations=20)

    positions = child.get_cell_positions()
    assert positions.shape[0] > 0, "Child colony died"

    clusters = detect_clusters(positions)
    print(f"  child colony → {len(clusters)} sub-clusters, {len(child.cells)} cells alive")
    # May find 0-3 sub-clusters depending on field tension; just verify it ran cleanly
    assert isinstance(clusters, list)


def test_run_child_sims_attaches_sub_strategies():
    """_run_child_sims should attach sub_strategies list to every strategy."""
    field = _make_field_with_attractors(6)
    config = _small_config(initial_cells=48, max_pop=100)
    factor_names = [f"dim_{i}" for i in range(6)]
    concepts = [
        {
            "name": f"dim_{i}",
            "dataDictionary": {"validRange": [0, 100]},
        }
        for i in range(6)
    ]

    # Simulate two parent strategies (one substantial, one tiny)
    substantial = {
        "label": "OptionA",
        "cell_count": 40,
        "values": {f"dim_{i}": 20.0 for i in range(6)},
        "concept_weights": {f"dim_{i}": 0.5 for i in range(6)},
        "is_emergent": False,
        "is_mixed": False,
        "centroid_raw": [0.2, 0.5, 0.5, 0.5, 0.5, 0.5],
    }
    tiny = {
        "label": "OptionB",
        "cell_count": 2,   # below _CHILD_MIN_CELLS threshold → skipped
        "values": {f"dim_{i}": 80.0 for i in range(6)},
        "concept_weights": {f"dim_{i}": 0.5 for i in range(6)},
        "is_emergent": False,
        "is_mixed": False,
        "centroid_raw": [0.8, 0.5, 0.5, 0.5, 0.5, 0.5],
    }

    strategies = [substantial, tiny]
    enriched = _run_child_sims(
        field, config, strategies, factor_names, concepts, child_generations=15)

    assert len(enriched) == 2, f"Expected 2 enriched strategies, got {len(enriched)}"

    # Substantial cluster should have sub_strategies key (may be empty list if no clusters found)
    assert "sub_strategies" in enriched[0], "sub_strategies missing from substantial strategy"
    assert isinstance(enriched[0]["sub_strategies"], list)

    # Tiny cluster should be skipped → empty sub_strategies
    assert "sub_strategies" in enriched[1], "sub_strategies missing from tiny strategy"
    assert enriched[1]["sub_strategies"] == [], \
        f"Tiny cluster should have empty sub_strategies, got {enriched[1]['sub_strategies']}"

    print(f"  substantial sub_strategies: {len(enriched[0]['sub_strategies'])}")
    print(f"  tiny sub_strategies: {len(enriched[1]['sub_strategies'])} (expected 0)")


def test_sub_strategies_values_in_real_units():
    """Sub-strategy values should be in real units (0-100 range), not 0-1 normalized."""
    field = _make_field_with_attractors(6)
    config = _small_config(initial_cells=48, max_pop=100)
    factor_names = [f"dim_{i}" for i in range(6)]
    concepts = [
        {
            "name": f"dim_{i}",
            "dataDictionary": {"validRange": [0, 100]},
        }
        for i in range(6)
    ]

    strategy = {
        "label": "OptionA",
        "cell_count": 50,
        "values": {},
        "concept_weights": {},
        "is_emergent": False,
        "is_mixed": False,
        "centroid_raw": [0.2, 0.5, 0.5, 0.5, 0.5, 0.5],
    }

    enriched = _run_child_sims(
        field, config, [strategy], factor_names, concepts, child_generations=20)

    sub = enriched[0]["sub_strategies"]
    for s in sub:
        for dim, val in s["values"].items():
            assert 0.0 <= val <= 100.0, \
                f"Sub-strategy value out of real-unit range: {dim}={val}"
    print(f"  {len(sub)} sub-strategies, all values in [0, 100] range")


# ─── Run ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    tests = [
        test_seed_centroid_places_cells_near_target,
        test_seed_spread_controls_radius,
        test_default_initialize_unchanged,
        test_child_colony_runs_and_produces_clusters,
        test_run_child_sims_attaches_sub_strategies,
        test_sub_strategies_values_in_real_units,
    ]

    passed = 0
    failed = 0
    for t in tests:
        name = t.__name__
        try:
            print(f"\n[RUN] {name}")
            t()
            print(f"[PASS] {name}")
            passed += 1
        except Exception as e:
            import traceback
            print(f"[FAIL] {name}: {e}")
            traceback.print_exc()
            failed += 1

    print(f"\n{'='*50}")
    print(f"Results: {passed} passed, {failed} failed")
    sys.exit(0 if failed == 0 else 1)
