"""
UNIT TESTS FOR COMPUTATIONAL FREE WILL FRAMEWORK
Validates all mathematical claims and edge cases
"""

import numpy as np
from layers.layer_2_core.free_will_framework import (
    AgentState, FreeWillIndex, CausalEntropyCalculator,
    IntegratedInformationCalculator, CounterfactualDepthCalculator,
    EmergenceProof, QuantumAgencyModel
)

# ============================================================================
# PROPERTY-BASED TESTS
# ============================================================================

def test_fwi_bounded():
    """FWI must always be in [0, 1] regardless of input"""
    for _ in range(100):
        # Random agent configurations
        n_beliefs = np.random.randint(5, 50)
        n_actions = np.random.randint(5, 30)
        
        agent = AgentState(
            belief_state=np.random.randn(n_beliefs) * 10,  # Extreme values
            goal_state=np.random.rand(5),
            meta_belief=np.random.randn(8) * 5,
            action_repertoire=np.random.randn(n_actions, 3) * 100
        )
        
        def dummy_dynamics(s, a):
            a_flat = a.flatten()
            a_proj = a_flat[:len(s)] if len(a_flat) >= len(s) else np.pad(a_flat, (0, len(s) - len(a_flat)))
            return s + a_proj * 0.1
        
        connectivity = np.random.rand(n_beliefs, n_beliefs)
        bounds = np.array([10.0, 10.0, 10.0])
        
        calculator = FreeWillIndex()
        result = calculator.compute(agent, dummy_dynamics, connectivity, bounds)
        
        assert 0 <= result['fwi'] <= 1, f"FWI out of bounds: {result['fwi']}"
        print(f"  Test {_+1}: FWI = {result['fwi']:.4f} ✓")

def test_causal_entropy_monotonicity():
    """More actions should never decrease causal entropy"""
    n_beliefs = 10
    base_state = np.random.randn(n_beliefs)
    
    def dynamics(s, a):
        a_flat = a.flatten()
        a_proj = a_flat[:len(s)] if len(a_flat) >= len(s) else np.pad(a_flat, (0, len(s) - len(a_flat)))
        return s * 0.9 + a_proj * 0.1
    
    calc = CausalEntropyCalculator(time_horizon=10)
    
    entropies = []
    for n_actions in [5, 10, 20, 40]:
        action_space = np.random.randn(n_actions, 3)
        entropy = calc.compute_causal_entropy(base_state, dynamics, action_space)
        entropies.append(entropy)
        print(f"  n_actions={n_actions:2d}: H_causal = {entropy:.4f}")
    
    # Check monotonicity
    for i in range(len(entropies) - 1):
        assert entropies[i] <= entropies[i+1] + 0.5, "Causal entropy decreased with more actions!"
    
    print("  Monotonicity verified ✓")

def test_phi_integration_property():
    """Fully connected graph should have higher Phi than disconnected"""
    n = 10
    
    # Disconnected graph
    disconnected = np.zeros((n, n))
    
    # Fully connected
    connected = np.ones((n, n))
    np.fill_diagonal(connected, 0)
    
    state = np.random.randn(n)
    
    calc = IntegratedInformationCalculator()
    phi_disconnected = calc.compute_phi(disconnected, state)
    phi_connected = calc.compute_phi(connected, state)
    
    print(f"  Φ (disconnected): {phi_disconnected:.4f}")
    print(f"  Φ (connected):    {phi_connected:.4f}")
    
    assert phi_connected > phi_disconnected, "Connected graph should have higher Phi!"
    print("  Integration property verified ✓")

def test_counterfactual_increases_with_diversity():
    """More diverse action space should increase counterfactual depth"""
    n_beliefs = 10
    state = np.random.randn(n_beliefs)
    
    def dynamics(s, a):
        a_flat = a.flatten()
        a_proj = a_flat[:len(s)] if len(a_flat) >= len(s) else np.pad(a_flat, (0, len(s) - len(a_flat)))
        return s + a_proj * 0.5  # Strong influence
    
    calc = CounterfactualDepthCalculator()
    
    # Similar actions
    similar_actions = np.random.randn(20, 3) * 0.1
    n_cf_similar, div_similar = calc.compute_counterfactual_depth(
        state, similar_actions, dynamics
    )
    
    # Diverse actions
    diverse_actions = np.random.randn(20, 3) * 5.0
    n_cf_diverse, div_diverse = calc.compute_counterfactual_depth(
        state, diverse_actions, dynamics
    )
    
    print(f"  Similar actions: {n_cf_similar} futures, div={div_similar:.4f}")
    print(f"  Diverse actions: {n_cf_diverse} futures, div={div_diverse:.4f}")
    
    assert n_cf_diverse >= n_cf_similar, "Diverse actions should create more futures!"
    assert div_diverse > div_similar, "Diverse actions should have higher divergence!"
    print("  Diversity property verified ✓")

def test_emergence_proof_consistency():
    """Emergence proof should be consistent across random initializations"""
    
    successes = 0
    trials = 50
    
    for _ in range(trials):
        agent = AgentState(
            belief_state=np.random.randn(10),
            goal_state=np.random.rand(5),
            meta_belief=np.random.randn(8) * 0.3,  # Low variance = high metacognition
            action_repertoire=np.random.randn(20, 3)
        )
        
        # Mock FWI result with good properties
        fwi_result = {
            'fwi': 0.75,
            'counterfactual_count': 15,
            'components': {'integration_phi': 0.8}
        }
        
        proof = EmergenceProof.prove_emergence(
            agent, fwi_result, self_prediction_accuracy=0.85
        )
        
        if proof['emergence_proven']:
            successes += 1
    
    success_rate = successes / trials
    print(f"  Emergence proven in {successes}/{trials} trials ({success_rate:.1%})")
    assert success_rate > 0.7, "Emergence should be common with good parameters!"
    print("  Consistency verified ✓")

def test_quantum_collapse():
    """Quantum measurement should collapse superposition to single state"""
    n_actions = 10
    qagent = QuantumAgencyModel(n_actions)
    
    # Initial superposition
    entropy_before = qagent.get_decision_entropy()
    print(f"  Entropy before measurement: {entropy_before:.4f} nats")
    assert entropy_before > 1.0, "Should start in superposition"
    
    # Measure (collapse)
    action = qagent.measure_action()
    entropy_after = qagent.get_decision_entropy()
    
    print(f"  Collapsed to action: {action}")
    print(f"  Entropy after measurement: {entropy_after:.6f} nats")
    
    assert entropy_after < 0.01, "Should collapse to nearly zero entropy"
    assert 0 <= action < n_actions, "Action out of bounds"
    print("  Collapse verified ✓")

def test_fwi_components_contribute():
    """Each FWI component should actually affect the final score"""
    
    base_agent = AgentState(
        belief_state=np.random.randn(10),
        goal_state=np.random.rand(5),
        meta_belief=np.random.randn(8) * 0.5,
        action_repertoire=np.random.randn(20, 3)
    )
    
    def dynamics(s, a):
        a_flat = a.flatten()
        a_proj = a_flat[:len(s)] if len(a_flat) >= len(s) else np.pad(a_flat, (0, len(s) - len(a_flat)))
        return s * 0.9 + a_proj * 0.1
    
    connectivity = np.random.rand(10, 10)
    connectivity = (connectivity + connectivity.T) / 2
    bounds = np.array([2.0, 2.0, 2.0])
    
    calc = FreeWillIndex()
    
    # Baseline
    result_base = calc.compute(base_agent, dynamics, connectivity, bounds)
    fwi_base = result_base['fwi']
    
    # Test 1: Zero all weights except causal_entropy
    calc_ce_only = FreeWillIndex(weights={
        'causal_entropy': 1.0,
        'integration': 0.0,
        'counterfactual': 0.0,
        'metacognition': 0.0,
        'constraint_penalty': 0.0
    })
    result_ce = calc_ce_only.compute(base_agent, dynamics, connectivity, bounds)
    
    # Test 2: Zero all weights except integration
    calc_phi_only = FreeWillIndex(weights={
        'causal_entropy': 0.0,
        'integration': 1.0,
        'counterfactual': 0.0,
        'metacognition': 0.0,
        'constraint_penalty': 0.0
    })
    result_phi = calc_phi_only.compute(base_agent, dynamics, connectivity, bounds)
    
    print(f"  FWI (balanced):    {fwi_base:.4f}")
    print(f"  FWI (CE only):     {result_ce['fwi']:.4f}")
    print(f"  FWI (Phi only):    {result_phi['fwi']:.4f}")
    
    # They should be different
    assert result_ce['fwi'] != result_phi['fwi'], "Components don't contribute independently!"
    print("  Component independence verified ✓")


# ============================================================================
# REGRESSION TESTS
# ============================================================================

def test_known_configuration():
    """Test against a known configuration with expected output"""
    
    np.random.seed(42)  # Reproducibility
    
    agent = AgentState(
        belief_state=np.array([1.0, -0.5, 0.3, 0.8, -0.2, 0.1, 0.9, -0.3, 0.4, -0.1]),
        goal_state=np.array([0.5, 0.7, 0.3, 0.9, 0.1]),
        meta_belief=np.array([0.2, -0.1, 0.3, 0.0, 0.1, -0.2, 0.15, 0.05]),
        action_repertoire=np.random.randn(20, 3)
    )
    
    def dynamics(s, a):
        a_flat = a.flatten()
        a_proj = a_flat[:len(s)] if len(a_flat) >= len(s) else np.pad(a_flat, (0, len(s) - len(a_flat)))
        return s * 0.9 + a_proj * 0.1
    
    connectivity = np.eye(10) * 0.5 + np.random.rand(10, 10) * 0.1
    connectivity = (connectivity + connectivity.T) / 2
    bounds = np.array([2.0, 2.0, 2.0])
    
    calc = FreeWillIndex()
    result = calc.compute(agent, dynamics, connectivity, bounds)
    
    print(f"  Regression FWI: {result['fwi']:.4f}")
    
    # Expected range (not exact due to randomness in calculations)
    assert 0.5 < result['fwi'] < 0.9, f"FWI out of expected range: {result['fwi']}"
    assert result['counterfactual_count'] >= 1, "Should have counterfactuals"
    print("  Regression test passed ✓")


# ============================================================================
# EDGE CASES
# ============================================================================

def test_zero_actions():
    """System should handle zero available actions gracefully"""
    
    agent = AgentState(
        belief_state=np.random.randn(10),
        goal_state=np.random.rand(5),
        meta_belief=np.random.randn(8),
        action_repertoire=np.array([]).reshape(0, 3)  # No actions
    )
    
    def dynamics(s, a):
        return s  # No change if no actions
    
    connectivity = np.random.rand(10, 10)
    bounds = np.array([1.0, 1.0, 1.0])
    
    calc = FreeWillIndex()
    
    try:
        result = calc.compute(agent, dynamics, connectivity, bounds)
        print(f"  Zero-action FWI: {result['fwi']:.4f}")
        assert result['fwi'] < 0.3, "Should have very low FWI with no actions"
        print("  Edge case handled ✓")
    except Exception as e:
        print(f"  Caught expected error: {e}")
        print("  Edge case handled ✓")

def test_perfect_self_prediction():
    """Perfect self-prediction should fail Gödel limit"""
    
    agent = AgentState(
        belief_state=np.random.randn(10),
        goal_state=np.random.rand(5),
        meta_belief=np.zeros(8),  # Perfect certainty
        action_repertoire=np.random.randn(20, 3)
    )
    
    fwi_result = {
        'fwi': 0.9,
        'counterfactual_count': 10,
        'components': {'integration_phi': 0.8}
    }
    
    # Perfect prediction
    proof = EmergenceProof.prove_emergence(agent, fwi_result, self_prediction_accuracy=1.0)
    
    print(f"  Gödel limit with perfect prediction: {proof['godel_limit']}")
    assert not proof['godel_limit'], "Should fail Gödel limit with perfect prediction!"
    assert not proof['emergence_proven'], "Emergence should not be proven!"
    print("  Edge case handled ✓")


# ============================================================================
# RUN ALL TESTS
# ============================================================================

if __name__ == "__main__":
    print("\n" + "="*80)
    print("COMPREHENSIVE UNIT TEST SUITE")
    print("="*80)
    
    tests = [
        ("FWI Bounded [0,1]", test_fwi_bounded),
        ("Causal Entropy Monotonicity", test_causal_entropy_monotonicity),
        ("Phi Integration Property", test_phi_integration_property),
        ("Counterfactual Diversity", test_counterfactual_increases_with_diversity),
        ("Emergence Proof Consistency", test_emergence_proof_consistency),
        ("Quantum Collapse", test_quantum_collapse),
        ("Component Independence", test_fwi_components_contribute),
        ("Known Configuration", test_known_configuration),
        ("Zero Actions Edge Case", test_zero_actions),
        ("Perfect Prediction Edge Case", test_perfect_self_prediction)
    ]
    
    passed = 0
    failed = 0
    
    for name, test_func in tests:
        print(f"\n[TEST] {name}")
        try:
            test_func()
            passed += 1
            print(f"  ✓ PASSED")
        except AssertionError as e:
            failed += 1
            print(f"  ✗ FAILED: {e}")
        except Exception as e:
            failed += 1
            print(f"  ✗ ERROR: {e}")
    
    print("\n" + "="*80)
    print(f"RESULTS: {passed} passed, {failed} failed")
    print("="*80 + "\n")