"""Core engine tests — verify AFS Python port produces meaningful results."""

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, compute_concept_weights


def create_coffee_shop_field() -> FieldDefinition:
    """Same 6D coffee shop field as the JS headless demo."""
    field = FieldDefinition(metadata={"problem": "Optimize coffee shop staffing"})

    dims = [
        ("employee_availability", "operational"),
        ("customer_demand", "operational"),
        ("labor_cost", "financial"),
        ("service_quality", "human"),
        ("shift_length", "operational"),
        ("break_schedule", "operational"),
    ]

    for name, cat in dims:
        dim_id = field.add_dimension(name=name, category=cat)
        field.promote_dimension(dim_id)

    matrix = np.array([
        [1.0, 0.65, 0.72, 0.58, 0.68, 0.55],
        [0.65, 1.0, 0.70, 0.66, 0.60, 0.52],
        [0.72, 0.70, 1.0, 0.61, 0.64, 0.56],
        [0.58, 0.66, 0.61, 1.0, 0.63, 0.59],
        [0.68, 0.60, 0.64, 0.63, 1.0, 0.74],
        [0.55, 0.52, 0.56, 0.59, 0.74, 1.0],
    ])
    field.update_similarity_matrix(matrix)
    return field


def create_tradeoff_field() -> FieldDefinition:
    """Field with strong positive AND negative relationships — should differentiate."""
    field = FieldDefinition(metadata={"problem": "Test trade-off dynamics"})

    dims = [
        ("speed", "performance"),
        ("quality", "performance"),
        ("cost", "financial"),
        ("flexibility", "design"),
        ("reliability", "design"),
    ]

    for name, cat in dims:
        dim_id = field.add_dimension(name=name, category=cat)
        field.promote_dimension(dim_id)

    # Strong trade-offs: speed vs quality, cost vs reliability
    matrix = np.array([
        [ 1.0, -0.8,  0.6,  0.5, -0.3],
        [-0.8,  1.0, -0.5, -0.3,  0.7],
        [ 0.6, -0.5,  1.0,  0.4, -0.6],
        [ 0.5, -0.3,  0.4,  1.0,  0.3],
        [-0.3,  0.7, -0.6,  0.3,  1.0],
    ])
    field.update_similarity_matrix(matrix)
    return field


def test_field_definition():
    print("test_field_definition...", end=" ")
    field = create_coffee_shop_field()
    assert field.get_dimension_count() == 6
    assert field.similarity_matrix.shape == (6, 6)
    assert field.similarity_matrix[0, 0] == 1.0
    print("PASS")


def test_field_serialization():
    print("test_field_serialization...", end=" ")
    field = create_coffee_shop_field()
    data = field.to_dict()
    restored = FieldDefinition.from_dict(data)
    assert restored.get_dimension_count() == 6
    assert np.allclose(restored.similarity_matrix, field.similarity_matrix)
    print("PASS")


def test_engine_initialization():
    print("test_engine_initialization...", end=" ")
    field = create_coffee_shop_field()
    config = AfsSimulationConfig.from_dict({"engine": {"initial_cell_count": 32}})
    engine = AfsEngine(field, config)
    engine.initialize()
    assert len(engine.cells) == 32
    assert len(engine.concepts) == 6
    print("PASS")


def test_engine_tick():
    print("test_engine_tick...", end=" ")
    field = create_coffee_shop_field()
    config = AfsSimulationConfig.from_dict({"engine": {"initial_cell_count": 16, "max_population": 50}})
    engine = AfsEngine(field, config)
    engine.initialize()

    initial_positions = engine.get_cell_positions().copy()
    for _ in range(100):
        engine.tick()

    final_positions = engine.get_cell_positions()
    # Cells should have moved
    assert not np.allclose(initial_positions[:len(final_positions)], final_positions[:len(initial_positions)])
    # No NaN
    assert not np.any(np.isnan(final_positions))
    print("PASS")


def test_engine_generations():
    print("test_engine_generations...", end=" ")
    field = create_coffee_shop_field()
    config = AfsSimulationConfig.from_dict({
        "engine": {"initial_cell_count": 32, "max_population": 64, "ticks_per_generation": 100}
    })
    engine = AfsEngine(field, config)
    engine.initialize()

    gen_count = 0
    def on_gen(data):
        nonlocal gen_count
        gen_count += 1

    engine.on("generation", on_gen)

    for _ in range(500):
        engine.tick()

    assert gen_count == 5  # 500 ticks / 100 ticks per gen
    assert engine.generation == 5
    print("PASS")


def test_cluster_detection():
    print("test_cluster_detection...", end=" ")
    # Create two clearly separated clusters
    n_dims = 3
    cluster_a = np.random.uniform(0.1, 0.3, (20, n_dims))
    cluster_b = np.random.uniform(0.7, 0.9, (20, n_dims))
    positions = np.vstack([cluster_a, cluster_b])

    clusters = detect_clusters(positions)
    assert len(clusters) >= 2, f"Expected >= 2 clusters, got {len(clusters)}"
    total = sum(c.cell_count for c in clusters)
    assert total > 30  # most cells assigned
    print(f"PASS ({len(clusters)} clusters, {total} cells assigned)")


def test_tradeoff_differentiation():
    """Run a field with trade-offs and verify cells spread out."""
    print("test_tradeoff_differentiation...", end=" ")
    field = create_tradeoff_field()
    config = AfsSimulationConfig.from_dict({
        "engine": {
            "initial_cell_count": 48,
            "max_population": 100,
            "concept_gravity_strength": 3.0,
            "force_scale": 0.05,
            "thermal_noise": 0.02,
        }
    })
    engine = AfsEngine(field, config)
    engine.initialize()

    for _ in range(3000):  # 30 generations
        engine.tick()

    positions = engine.get_cell_positions()
    assert not np.any(np.isnan(positions)), "NaN in cell positions"

    # Check spread: variance across dimensions should be non-trivial
    variance = np.var(positions, axis=0)
    mean_var = float(np.mean(variance))
    assert mean_var > 0.001, f"Cells not spreading: mean variance = {mean_var}"

    # Try cluster detection
    clusters = detect_clusters(positions)
    print(f"PASS (pop={len(engine.cells)}, spread={mean_var:.4f}, clusters={len(clusters)})")


def test_concept_weights():
    print("test_concept_weights...", end=" ")
    centroid = np.array([0.3, 0.3, 0.3])
    concept_positions = np.array([
        [0.3, 0.3, 0.3],  # right on the centroid
        [0.8, 0.8, 0.8],  # far away
    ])
    names = ["close_concept", "far_concept"]

    cluster = type("C", (), {"centroid": centroid})()
    weights = compute_concept_weights(cluster, concept_positions, names, 3)
    assert weights["close_concept"] > weights["far_concept"]
    print("PASS")


def test_no_nan_stability():
    """Run for many generations and verify no NaN contamination."""
    print("test_no_nan_stability...", end=" ")
    field = create_tradeoff_field()
    config = AfsSimulationConfig.from_dict({
        "engine": {
            "initial_cell_count": 64,
            "max_population": 128,
            "concept_gravity_strength": 5.0,
            "force_scale": 0.1,
            "thermal_noise": 0.04,
        }
    })
    engine = AfsEngine(field, config)
    engine.initialize()

    for tick in range(10000):  # 100 generations
        engine.tick()
        if tick % 1000 == 0:
            positions = engine.get_cell_positions()
            if np.any(np.isnan(positions)):
                raise AssertionError(f"NaN at tick {tick}")

    positions = engine.get_cell_positions()
    energies = np.array([c.energy for c in engine.cells.values()])
    assert not np.any(np.isnan(positions)), "NaN in final positions"
    assert not np.any(np.isnan(energies)), "NaN in final energies"
    print(f"PASS (100 gen, pop={len(engine.cells)}, avg_E={np.mean(energies):.1f})")


if __name__ == "__main__":
    print("=" * 60)
    print("AFS Python Engine — Core Tests")
    print("=" * 60)
    print()

    tests = [
        test_field_definition,
        test_field_serialization,
        test_engine_initialization,
        test_engine_tick,
        test_engine_generations,
        test_cluster_detection,
        test_concept_weights,
        test_tradeoff_differentiation,
        test_no_nan_stability,
    ]

    passed = 0
    failed = 0
    for test in tests:
        try:
            test()
            passed += 1
        except Exception as e:
            print(f"FAIL: {e}")
            failed += 1

    print()
    print(f"Results: {passed} passed, {failed} failed")
