"""Tests for WorldModelTranslator, DomainAnalyzer, and FieldCompiler."""

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

import asyncio
import json
import numpy as np
from afs.compiler import WorldModelTranslator, FieldCompiler


# ─── Test World Model (same format as JS engine produces) ─────

SAMPLE_WORLD_MODEL = {
    "domain": "staffing",
    "problem": {
        "statement": "Optimize coffee shop staffing",
        "restatement": "Find optimal staffing levels for a coffee shop",
        "objectives": ["Minimize cost", "Maximize service quality"],
    },
    "concepts": [
        {
            "name": "labor_cost",
            "displayName": "Labor Cost",
            "description": "Total hourly labor expenditure",
            "nature": "resource",
            "category": "financial",
            "controllability": "direct",
            "dataDictionary": {
                "validRange": [500, 5000],
                "displayRange": [500, 5000],
                "displayUnit": "$/week",
                "unitOfMeasure": "$/week",
            },
        },
        {
            "name": "service_quality",
            "displayName": "Service Quality",
            "description": "Customer satisfaction and wait times",
            "nature": "outcome",
            "category": "human",
            "controllability": "indirect",
            "dataDictionary": {
                "validRange": [0, 100],
                "displayRange": [0, 100],
                "displayUnit": "score",
                "unitOfMeasure": "score",
            },
        },
        {
            "name": "staff_count",
            "displayName": "Staff Count",
            "description": "Number of employees on shift",
            "nature": "resource",
            "category": "operational",
            "controllability": "direct",
            "dataDictionary": {
                "validRange": [2, 20],
                "displayRange": [2, 20],
                "displayUnit": "people",
                "unitOfMeasure": "people",
            },
        },
        {
            "name": "customer_demand",
            "displayName": "Customer Demand",
            "description": "Expected customer volume per hour",
            "nature": "force",
            "category": "operational",
            "controllability": "uncontrollable",
            "dataDictionary": {
                "validRange": [10, 200],
                "displayRange": [10, 200],
                "displayUnit": "customers/hr",
                "unitOfMeasure": "customers/hr",
            },
        },
    ],
    "relationships": [
        {
            "fromConcept": "staff_count",
            "toConcept": "labor_cost",
            "direction": 1,
            "strength": 0.9,
            "nature": "enables",
            "mechanism": "More staff directly increases labor expenditure",
            "responseCurve": {
                "type": "linear",
                "truthTable": [
                    {"fromValue": 2, "effectOnTo": 500, "note": "Minimum staffing"},
                    {"fromValue": 10, "effectOnTo": 2500, "note": "Normal staffing"},
                    {"fromValue": 20, "effectOnTo": 5000, "note": "Maximum staffing"},
                ],
            },
            "confidence": "high",
            "symmetry": {"isSymmetric": False},
        },
        {
            "fromConcept": "staff_count",
            "toConcept": "service_quality",
            "direction": 1,
            "strength": 0.7,
            "nature": "enables",
            "mechanism": "More staff reduces wait times and improves service",
            "responseCurve": {"type": "logarithmic", "truthTable": []},
            "confidence": "medium",
            "symmetry": {"isSymmetric": False},
        },
        {
            "fromConcept": "labor_cost",
            "toConcept": "service_quality",
            "direction": -1,
            "strength": 0.6,
            "nature": "trades_off",
            "mechanism": "Budget pressure forces cutting staff, reducing quality",
            "responseCurve": {"type": "linear", "truthTable": []},
            "confidence": "medium",
            "symmetry": {"isSymmetric": False},
        },
        {
            "fromConcept": "customer_demand",
            "toConcept": "service_quality",
            "direction": -1,
            "strength": 0.8,
            "nature": "constrains",
            "mechanism": "Higher demand with same staff degrades service",
            "responseCurve": {"type": "linear", "truthTable": []},
            "confidence": "high",
            "symmetry": {"isSymmetric": False},
        },
    ],
    "tradeoffs": [
        {"concepts": ["labor_cost", "service_quality"], "description": "Cost vs quality", "severity": "severe"},
    ],
    "referencePoints": [
        {
            "name": "Budget Operation",
            "description": "Minimum staffing",
            "position": {"labor_cost": 1000, "service_quality": 30, "staff_count": 3, "customer_demand": 50},
            "source": "industry data",
        },
        {
            "name": "Premium Operation",
            "description": "Full staffing",
            "position": {"labor_cost": 4000, "service_quality": 90, "staff_count": 15, "customer_demand": 150},
            "source": "industry data",
        },
    ],
    "attractors": [
        {
            "name": "Cost Minimizer",
            "description": "Bare minimum staffing",
            "position": {"labor_cost": 800, "service_quality": 25, "staff_count": 2, "customer_demand": 50},
            "strength": 0.7,
        },
        {
            "name": "Quality Leader",
            "description": "Premium service experience",
            "position": {"labor_cost": 4500, "service_quality": 95, "staff_count": 18, "customer_demand": 150},
            "strength": 0.8,
        },
        {
            "name": "Balanced",
            "description": "Mid-range efficiency",
            "position": {"labor_cost": 2500, "service_quality": 70, "staff_count": 10, "customer_demand": 100},
            "strength": 0.9,
        },
    ],
    "confidence": {"overall": "medium", "weakestLink": "customer_demand", "dataGaps": []},
    "validation": {"valid": True, "errors": [], "warnings": []},
}


# ─── WorldModelTranslator Tests ───────────────────────────────

def test_influence_matrix_shape():
    print("test_influence_matrix_shape...", end=" ")
    translator = WorldModelTranslator()
    matrix = translator.build_influence_matrix(
        SAMPLE_WORLD_MODEL["concepts"],
        SAMPLE_WORLD_MODEL["relationships"],
    )
    assert matrix.shape == (4, 4), f"Expected (4,4), got {matrix.shape}"
    print("PASS")


def test_influence_matrix_signs():
    """Trade-off relationships should produce negative matrix values."""
    print("test_influence_matrix_signs...", end=" ")
    translator = WorldModelTranslator()
    matrix = translator.build_influence_matrix(
        SAMPLE_WORLD_MODEL["concepts"],
        SAMPLE_WORLD_MODEL["relationships"],
    )
    # staff_count → labor_cost should be positive (direction=1)
    # labor_cost → service_quality should be negative (direction=-1)
    # customer_demand → service_quality should be negative (direction=-1)
    concepts = SAMPLE_WORLD_MODEL["concepts"]
    idx = {c["name"]: i for i, c in enumerate(concepts)}

    assert matrix[idx["staff_count"], idx["labor_cost"]] > 0, "staff→cost should be positive"
    assert matrix[idx["staff_count"], idx["service_quality"]] > 0, "staff→quality should be positive"
    assert matrix[idx["labor_cost"], idx["service_quality"]] < 0, "cost→quality should be negative"
    assert matrix[idx["customer_demand"], idx["service_quality"]] < 0, "demand→quality should be negative"

    has_negative = np.any(matrix < 0)
    assert has_negative, "Matrix should have negative (trade-off) values"
    print(f"PASS (has negatives, shape={matrix.shape})")


def test_influence_matrix_strength():
    print("test_influence_matrix_strength...", end=" ")
    translator = WorldModelTranslator()
    matrix = translator.build_influence_matrix(
        SAMPLE_WORLD_MODEL["concepts"],
        SAMPLE_WORLD_MODEL["relationships"],
    )
    # All values should be in [-1, 1]
    assert np.all(matrix >= -1.0) and np.all(matrix <= 1.0), "Values should be in [-1, 1]"
    print("PASS")


def test_truth_table_interpolation():
    print("test_truth_table_interpolation...", end=" ")
    translator = WorldModelTranslator()

    # With operating point at staff_count=10, the truth table should interpolate
    matrix = translator.build_influence_matrix(
        SAMPLE_WORLD_MODEL["concepts"],
        SAMPLE_WORLD_MODEL["relationships"],
        operating_point={"staff_count": 10},
    )
    idx = {c["name"]: i for i, c in enumerate(SAMPLE_WORLD_MODEL["concepts"])}

    # staff_count → labor_cost has truth table, at staff_count=10 effect is 2500
    val = matrix[idx["staff_count"], idx["labor_cost"]]
    assert val > 0, f"Expected positive value, got {val}"
    print(f"PASS (staff→cost at operating_point=10: {val:.4f})")


def test_reference_point_normalization():
    print("test_reference_point_normalization...", end=" ")
    translator = WorldModelTranslator()
    concepts = SAMPLE_WORLD_MODEL["concepts"]

    # Budget operation: labor_cost=1000 in range [500, 5000] → ~0.11
    pos = translator.normalize_reference_point(
        SAMPLE_WORLD_MODEL["referencePoints"][0],
        concepts,
    )
    assert pos.shape == (4,), f"Expected (4,), got {pos.shape}"
    assert 0 <= pos[0] <= 1, f"Normalized value out of range: {pos[0]}"

    # Premium operation should have higher normalized values for service_quality
    pos_premium = translator.normalize_reference_point(
        SAMPLE_WORLD_MODEL["referencePoints"][1],
        concepts,
    )
    idx = {c["name"]: i for i, c in enumerate(concepts)}
    assert pos_premium[idx["service_quality"]] > pos[idx["service_quality"]], \
        "Premium should have higher service quality"
    print(f"PASS (budget={pos[idx['service_quality']]:.3f}, premium={pos_premium[idx['service_quality']]:.3f})")


def test_cell_position_generation():
    print("test_cell_position_generation...", end=" ")
    translator = WorldModelTranslator()
    concepts = SAMPLE_WORLD_MODEL["concepts"]

    positions = translator.create_cell_positions(
        SAMPLE_WORLD_MODEL["referencePoints"],
        total_cells=20,
        concepts=concepts,
    )
    assert positions.shape == (20, 4), f"Expected (20,4), got {positions.shape}"
    assert np.all(positions >= 0) and np.all(positions <= 1), "Positions should be in [0,1]"
    print(f"PASS (shape={positions.shape})")


def test_cell_positions_no_refs():
    """Without reference points, should generate Gaussian positions."""
    print("test_cell_positions_no_refs...", end=" ")
    translator = WorldModelTranslator()
    positions = translator.create_cell_positions(
        [],
        total_cells=32,
        concepts=SAMPLE_WORLD_MODEL["concepts"],
    )
    assert positions.shape == (32, 4)
    assert np.all(positions >= 0) and np.all(positions <= 1)
    print("PASS")


def test_influence_vector():
    print("test_influence_vector...", end=" ")
    translator = WorldModelTranslator()
    matrix = translator.build_influence_matrix(
        SAMPLE_WORLD_MODEL["concepts"],
        SAMPLE_WORLD_MODEL["relationships"],
    )
    vec = translator.build_influence_vector(0, matrix)
    assert vec.shape == (4,)
    assert np.array_equal(vec, matrix[0])
    print("PASS")


def test_relinearize():
    print("test_relinearize...", end=" ")
    translator = WorldModelTranslator()
    concepts = SAMPLE_WORLD_MODEL["concepts"]
    relationships = SAMPLE_WORLD_MODEL["relationships"]

    # Create some cell positions
    positions = np.random.uniform(0.2, 0.8, (16, 4))
    matrix = translator.relinearize(positions, concepts, relationships)
    assert matrix.shape == (4, 4)
    assert not np.any(np.isnan(matrix))
    print("PASS")


# ─── FieldCompiler Tests (from world model, no LLM) ───────────

def test_compile_from_world_model():
    print("test_compile_from_world_model...", end=" ")
    compiler = FieldCompiler.__new__(FieldCompiler)
    compiler.translator = WorldModelTranslator()
    compiler.analyzer = None  # Not needed for this path

    result = compiler.compile_from_world_model(SAMPLE_WORLD_MODEL)

    field = result["field"]
    assert field.get_dimension_count() == 4
    assert field.similarity_matrix is not None
    assert field.similarity_matrix.shape == (4, 4)
    assert np.any(field.similarity_matrix < 0), "Should have trade-off values"

    validation = result["validation"]
    assert validation["valid"]

    # Check world model metadata preserved
    assert field.metadata.get("worldModel") is not None
    assert field.metadata["worldModel"]["relationshipCount"] == 4
    print(f"PASS (dims={field.get_dimension_count()}, rels={field.metadata['worldModel']['relationshipCount']})")


def test_compile_from_world_model_to_engine():
    """Full pipeline: world model → FieldCompiler → AfsEngine → clusters."""
    print("test_compile_from_world_model_to_engine...", end=" ")
    from afs.core import AfsEngine, AfsSimulationConfig
    from afs.calibrator import HomeostasisCalibrator
    from afs.analysis import detect_clusters

    compiler = FieldCompiler.__new__(FieldCompiler)
    compiler.translator = WorldModelTranslator()

    result = compiler.compile_from_world_model(SAMPLE_WORLD_MODEL)
    field = result["field"]

    # Quick calibration
    cal = HomeostasisCalibrator(field, probe_generations=10, probe_cells=16, target_cells=48, target_max_pop=96)
    config = cal.calibrate(verbose=False)

    engine = AfsEngine(field, config)
    engine.initialize()

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

    positions = engine.get_cell_positions()
    assert not np.any(np.isnan(positions))
    clusters = detect_clusters(positions)
    print(f"PASS (pop={len(engine.cells)}, clusters={len(clusters)}, winner={config._calibration['label']})")


def test_world_model_json_interchange():
    """Verify world model can round-trip through JSON (JS ↔ Python interchange)."""
    print("test_world_model_json_interchange...", end=" ")
    json_str = json.dumps(SAMPLE_WORLD_MODEL)
    restored = json.loads(json_str)

    compiler = FieldCompiler.__new__(FieldCompiler)
    compiler.translator = WorldModelTranslator()

    result = compiler.compile_from_world_model(restored)
    field = result["field"]
    assert field.get_dimension_count() == 4
    assert np.any(field.similarity_matrix < 0)

    # Field can also round-trip
    field_json = field.to_json()
    from afs.core import FieldDefinition
    restored_field = FieldDefinition.from_json(field_json)
    assert restored_field.get_dimension_count() == 4
    assert np.allclose(restored_field.similarity_matrix, field.similarity_matrix)
    print("PASS")


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

    tests = [
        test_influence_matrix_shape,
        test_influence_matrix_signs,
        test_influence_matrix_strength,
        test_truth_table_interpolation,
        test_reference_point_normalization,
        test_cell_position_generation,
        test_cell_positions_no_refs,
        test_influence_vector,
        test_relinearize,
        test_compile_from_world_model,
        test_compile_from_world_model_to_engine,
        test_world_model_json_interchange,
    ]

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

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