"""Tests for afs.core.output_wm — output_to_world_model() pure function.

6 tests covering: basic conversion, relationship extraction, density weighting,
single strategy edge case, missing factor handling, and round-trip compilation.
"""

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

import numpy as np
from afs.core.output_wm import output_to_world_model, _weighted_pearson


# ─── Fixtures ────────────────────────────────────────────────────────────────

def _make_strategy(label, weights, cell_count=50):
    """Build a minimal strategy dict for testing."""
    return {
        "label": label,
        "cell_count": cell_count,
        "values": {k: v * 100 for k, v in weights.items()},
        "concept_weights": weights,
        "is_emergent": False,
        "is_mixed": False,
        "centroid_raw": [0.5] * len(weights),
    }


# ─── Test 1: Basic conversion ─────────────────────────────────────────────────

def test_basic_conversion():
    """3 strategies → factors = union, options = 3, scores correctly scaled."""
    print("test_basic_conversion...", end=" ")
    strategies = [
        _make_strategy("Alpha",   {"speed": 0.9, "cost": 0.2, "quality": 0.7}),
        _make_strategy("Beta",    {"speed": 0.4, "cost": 0.8, "quality": 0.5}),
        _make_strategy("Gamma",   {"speed": 0.6, "cost": 0.5, "quality": 0.9}),
    ]

    wm = output_to_world_model(strategies, source_label="TestRun")

    # Factors: union of all concept_weight keys
    factor_names = [f["name"] for f in wm["factors"]]
    assert set(factor_names) == {"speed", "cost", "quality"}, \
        f"Expected {{speed, cost, quality}}, got {set(factor_names)}"
    assert all(f["range"] == [0, 100] for f in wm["factors"])
    assert all(f["unit"] == "score" for f in wm["factors"])

    # Options: one per strategy
    assert len(wm["options"]) == 3
    opt_names = [o["name"] for o in wm["options"]]
    assert "Alpha" in opt_names and "Beta" in opt_names and "Gamma" in opt_names

    # Scores: concept_weights × 100
    alpha_opt = next(o for o in wm["options"] if o["name"] == "Alpha")
    assert abs(alpha_opt["speed"] - 90.0) < 0.01, f"Expected 90.0, got {alpha_opt['speed']}"
    assert abs(alpha_opt["cost"] - 20.0) < 0.01
    assert abs(alpha_opt["quality"] - 70.0) < 0.01

    # Metadata
    assert wm["_derived"] is True
    assert wm["_strategy_count"] == 3
    assert "TestRun" in wm["question"]

    print(f"PASS (factors={factor_names}, options={opt_names})")


# ─── Test 2: Relationship extraction ─────────────────────────────────────────

def test_relationship_extraction():
    """Correlated pairs (|r|>0.3) are captured; weak pairs are filtered."""
    print("test_relationship_extraction...", end=" ")

    # speed and efficiency are perfectly correlated (both high/low together)
    # speed and cost are strongly anti-correlated (tradeoff)
    # noise_a and noise_b are uncorrelated (random-ish)
    strategies = [
        _make_strategy("S1", {"speed": 0.9, "efficiency": 0.88, "cost": 0.1,  "noise_a": 0.3}),
        _make_strategy("S2", {"speed": 0.2, "efficiency": 0.19, "cost": 0.85, "noise_a": 0.7}),
        _make_strategy("S3", {"speed": 0.6, "efficiency": 0.62, "cost": 0.45, "noise_a": 0.5}),
        _make_strategy("S4", {"speed": 0.8, "efficiency": 0.78, "cost": 0.2,  "noise_a": 0.4}),
    ]

    wm = output_to_world_model(strategies)
    rels = wm["relationships"]

    # Build lookup by sorted pair
    rel_map = {}
    for r in rels:
        pair = tuple(sorted([r["from"], r["to"]]))
        rel_map[pair] = r["strength"]

    # speed–efficiency: should be strongly positive (r ≈ 0.99)
    assert ("efficiency", "speed") in rel_map, \
        f"speed–efficiency correlation missing. Rels: {rels}"
    assert rel_map[("efficiency", "speed")] > 0.3, \
        f"speed–efficiency should be positive, got {rel_map[('efficiency', 'speed')]}"

    # speed–cost: should be strongly negative (tradeoff, r ≈ -0.99)
    assert ("cost", "speed") in rel_map, \
        f"speed–cost tradeoff missing. Rels: {rels}"
    assert rel_map[("cost", "speed")] < -0.3, \
        f"speed–cost should be negative (tradeoff), got {rel_map[('cost', 'speed')]}"

    # All relationship strengths must be in [-1, 1]
    for r in rels:
        assert -1.0 <= r["strength"] <= 1.0, f"strength out of range: {r}"

    print(f"PASS (rels={[(r['from'], r['to'], r['strength']) for r in rels]})")


# ─── Test 3: Density weighting ────────────────────────────────────────────────

def test_density_weighting():
    """Density weighting changes correlation outcome vs equal-weight baseline.

    Constructed so unweighted Pearson(a, b) ≈ -0.17 (filtered, below 0.3)
    but density-weighted Pearson(a, b) ≈ +0.59 (captured, above 0.3).

    Setup: 4 strategies (100 cells each) where a and b co-vary positively,
    plus 1 outlier strategy (10 cells) where b is very low despite high a.
    Without weighting: outlier drags r negative.
    With density weighting: 400-cell majority dominates → positive r.
    """
    print("test_density_weighting...", end=" ")

    strategies = [
        _make_strategy("S1", {"factor_a": 0.9, "factor_b": 0.9}, cell_count=100),
        _make_strategy("S2", {"factor_a": 0.8, "factor_b": 0.8}, cell_count=100),
        _make_strategy("S3", {"factor_a": 0.7, "factor_b": 0.7}, cell_count=100),
        _make_strategy("S4", {"factor_a": 0.6, "factor_b": 0.6}, cell_count=100),
        _make_strategy("S5", {"factor_a": 0.9, "factor_b": 0.1}, cell_count=10),  # outlier
    ]

    # Verify unweighted r is below threshold (baseline check)
    x = np.array([0.9, 0.8, 0.7, 0.6, 0.9])
    y = np.array([0.9, 0.8, 0.7, 0.6, 0.1])
    equal_w = np.array([0.2, 0.2, 0.2, 0.2, 0.2])
    r_unweighted = _weighted_pearson(x, y, equal_w)
    assert r_unweighted is not None and abs(r_unweighted) < 0.3, \
        f"Test setup: expected |r_unweighted| < 0.3, got {r_unweighted:.4f}"

    # Now run output_to_world_model — should use density weighting
    wm = output_to_world_model(strategies)
    rels = wm["relationships"]
    rel_map = {tuple(sorted([r["from"], r["to"]])): r["strength"] for r in rels}
    pair = ("factor_a", "factor_b")

    assert pair in rel_map, \
        f"factor_a–factor_b relationship missing with density weighting. Rels: {rels}\n" \
        f"Unweighted r was {r_unweighted:.4f} (below 0.3 — density weighting must lift it above)."
    r_weighted = rel_map[pair]
    assert r_weighted > 0.3, \
        f"Expected density-weighted r > 0.3, got {r_weighted:.4f}"

    print(f"PASS (r_unweighted={r_unweighted:.4f} filtered, r_weighted={r_weighted:.4f} captured)")


# ─── Test 4: Single strategy ──────────────────────────────────────────────────

def test_single_strategy():
    """1 strategy → valid WM with 0 relationships (no pairs to correlate)."""
    print("test_single_strategy...", end=" ")

    strategies = [
        _make_strategy("Lone", {"alpha": 0.7, "beta": 0.3, "gamma": 0.9}),
    ]

    wm = output_to_world_model(strategies)

    assert len(wm["options"]) == 1
    assert len(wm["factors"]) == 3
    assert wm["relationships"] == [], \
        f"Single strategy must yield no relationships, got {wm['relationships']}"
    assert wm["_strategy_count"] == 1

    print("PASS (0 relationships, valid structure)")


# ─── Test 5: Missing factor handling ─────────────────────────────────────────

def test_missing_factor_handling():
    """Strategies with non-overlapping factors → zero-fill, valid output."""
    print("test_missing_factor_handling...", end=" ")

    strategies = [
        _make_strategy("S1", {"shared": 0.8, "only_in_s1": 0.9}),
        _make_strategy("S2", {"shared": 0.3, "only_in_s2": 0.7}),
        _make_strategy("S3", {"shared": 0.6}),
    ]

    wm = output_to_world_model(strategies)

    factor_names = {f["name"] for f in wm["factors"]}
    assert "shared" in factor_names
    assert "only_in_s1" in factor_names
    assert "only_in_s2" in factor_names

    # S2 should have only_in_s1 = 0 (zero-filled)
    s2_opt = next(o for o in wm["options"] if o["name"] == "S2")
    assert s2_opt["only_in_s1"] == 0.0, \
        f"Missing factor should zero-fill, got {s2_opt['only_in_s1']}"

    # S1 should have only_in_s2 = 0 (zero-filled)
    s1_opt = next(o for o in wm["options"] if o["name"] == "S1")
    assert s1_opt["only_in_s2"] == 0.0, \
        f"Missing factor should zero-fill, got {s1_opt['only_in_s2']}"

    # All option dicts must have all factor keys
    for opt in wm["options"]:
        for fname in factor_names:
            assert fname in opt, f"Option {opt['name']} missing factor {fname}"

    print(f"PASS (factors={sorted(factor_names)}, zero-fill verified)")


# ─── Test 6: Round-trip compilation ──────────────────────────────────────────

def test_round_trip_compiles():
    """Derived WM → _build_world_model → compile_from_world_model: no errors."""
    print("test_round_trip_compiles...", end=" ")

    from afs.mcp.server import _build_world_model
    from afs.compiler import WorldModelTranslator, FieldCompiler

    strategies = [
        _make_strategy("Explorer", {"reach": 0.8, "depth": 0.3, "speed": 0.9}),
        _make_strategy("Digger",   {"reach": 0.2, "depth": 0.9, "speed": 0.4}),
        _make_strategy("Balanced", {"reach": 0.6, "depth": 0.5, "speed": 0.6}),
    ]

    derived_wm = output_to_world_model(strategies, source_label="round_trip_test")
    assert derived_wm is not None

    # Convert input format → full world model
    full_wm = _build_world_model(
        question=derived_wm["question"],
        options=derived_wm["options"],
        factors=derived_wm["factors"],
        relationships=derived_wm["relationships"],
    )

    # Compile: no LLM needed
    compiler = FieldCompiler.__new__(FieldCompiler)
    compiler.translator = WorldModelTranslator()
    compiler.analyzer = None

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

    assert field.get_dimension_count() > 0, "Compiled field has no dimensions"
    assert field.similarity_matrix is not None, "Influence matrix missing"
    n = field.get_dimension_count()
    assert field.similarity_matrix.shape == (n, n), \
        f"Matrix shape mismatch: {field.similarity_matrix.shape} != ({n}, {n})"

    print(f"PASS (dims={field.get_dimension_count()}, matrix={field.similarity_matrix.shape})")


# ─── Runner ───────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    tests = [
        test_basic_conversion,
        test_relationship_extraction,
        test_density_weighting,
        test_single_strategy,
        test_missing_factor_handling,
        test_round_trip_compiles,
    ]
    passed = 0
    failed = 0
    for t in tests:
        try:
            t()
            passed += 1
        except Exception as e:
            print(f"FAIL: {e}")
            import traceback
            traceback.print_exc()
            failed += 1
    print(f"\n{'='*50}")
    print(f"Results: {passed} passed, {failed} failed")
    if failed:
        sys.exit(1)
