#!/usr/bin/env python3
"""
test_log_prescreen.py — Test log prescreening on known physical laws

Verifies that the "Everything Is Logarithms" pre-filter discovers
common physical laws without expression tree search.
"""

import sys
import numpy as np

sys.path.insert(0, "/home/allaun/SilverSight/python")
from log_prescreen import prescreen


def test_kepler():
    """Kepler's Third Law: T = a^1.5"""
    a = np.array([0.387, 0.723, 1.000, 1.524, 5.203, 9.537, 19.191, 30.069])
    T = np.array([0.241, 0.615, 1.000, 1.881, 11.862, 29.457, 84.011, 164.79])
    result = prescreen(a, T, feature_names=["a"])
    assert result is not None, "Kepler: prescreen failed"
    assert result.r2 > 0.999, f"Kepler: R² = {result.r2:.6f} < 0.999"
    assert "power_law" in result.name, f"Kepler: expected power_law, got {result.name}"
    coeff, exp = result.params
    assert abs(exp - 1.5) < 0.01, f"Kepler: exponent = {exp:.4f}, expected 1.5"
    print(f"  ✓ Kepler: {result.expression}  R²={result.r2:.10f}")
    return True


def test_hooke():
    """Hooke's Law: F = k·x (linear)"""
    x = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8])
    F = np.array([2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0])  # k=20
    result = prescreen(x, F, feature_names=["x"])
    assert result is not None, "Hooke: prescreen failed"
    assert result.r2 > 0.999, f"Hooke: R² = {result.r2:.6f} < 0.999"
    assert "linear" in result.name, f"Hooke: expected linear, got {result.name}"
    print(f"  ✓ Hooke: {result.expression}  R²={result.r2:.10f}")
    return True


def test_newton_gravity():
    """Newton's inverse square law: F = G·m₁·m₂/r² (power law with exp=-2)"""
    r = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
    F = np.array([100.0, 25.0, 11.11, 6.25, 4.0, 2.78, 2.04, 1.5625])  # F = 100/r²
    result = prescreen(r, F, feature_names=["r"])
    assert result is not None, "Newton: prescreen failed"
    assert result.r2 > 0.999, f"Newton: R² = {result.r2:.6f} < 0.999"
    assert "power_law" in result.name, f"Newton: expected power_law, got {result.name}"
    coeff, exp = result.params
    assert abs(exp - (-2.0)) < 0.01, f"Newton: exponent = {exp:.4f}, expected -2.0"
    print(f"  ✓ Newton: {result.expression}  R²={result.r2:.10f}")
    return True


def test_radioactive_decay():
    """Radioactive decay: N = N₀·exp(-λt) (exponential)"""
    t = np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0])
    N = 100.0 * np.exp(-0.3 * t)
    result = prescreen(t, N, feature_names=["t"])
    assert result is not None, "Decay: prescreen failed"
    assert result.r2 > 0.999, f"Decay: R² = {result.r2:.6f} < 0.999"
    assert "exponential" in result.name, f"Decay: expected exponential, got {result.name}"
    lambda_est, N0_log = result.params
    assert abs(lambda_est - (-0.3)) < 0.01, f"Decay: λ = {lambda_est:.4f}, expected -0.3"
    print(f"  ✓ Decay: {result.expression}  R²={result.r2:.10f}")
    return True


def test_free_fall():
    """Free fall: s = ½gt² (quadratic with no linear term)"""
    t = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0])
    g = 9.81
    s = 0.5 * g * t**2
    result = prescreen(t, s, feature_names=["t"])
    assert result is not None, "Free fall: prescreen failed"
    assert result.r2 > 0.999, f"Free fall: R² = {result.r2:.6f} < 0.999"
    print(f"  ✓ Free fall: {result.expression}  R²={result.r2:.10f}")
    return True


def test_sqrt_law():
    """Pendulum period: T = 2π·√(L/g) ≈ a·√L"""
    L = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
    T = 2 * np.pi * np.sqrt(L / 9.81)  # T = 2π√(L/g)
    result = prescreen(L, T, feature_names=["L"])
    assert result is not None, "Pendulum: prescreen failed"
    assert result.r2 > 0.999, f"Pendulum: R² = {result.r2:.6f} < 0.999"
    print(f"  ✓ Pendulum: {result.expression}  R²={result.r2:.10f}")
    return True


def test_stefan_boltzmann():
    """Stefan-Boltzmann: P = σ·T⁴ (power law with exp=4)"""
    T = np.array([100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0, 900.0, 1000.0])
    sigma = 5.67e-8
    P = sigma * T**4
    result = prescreen(T, P, feature_names=["T"])
    assert result is not None, "Stefan-Boltzmann: prescreen failed"
    assert result.r2 > 0.999, f"Stefan-Boltzmann: R² = {result.r2:.6f} < 0.999"
    assert "power_law" in result.name, f"Stefan-Boltzmann: expected power_law, got {result.name}"
    coeff, exp = result.params
    assert abs(exp - 4.0) < 0.01, f"Stefan-Boltzmann: exponent = {exp:.4f}, expected 4.0"
    print(f"  ✓ Stefan-Boltzmann: {result.expression}  R²={result.r2:.10f}")
    return True


def test_no_simple_law():
    """sin(x) — should NOT be caught by prescreen (require expression tree search)"""
    x = np.linspace(0, 2 * np.pi, 20)
    y = np.sin(x)
    result = prescreen(x, y, feature_names=["x"], r2_threshold=0.999)
    if result is None:
        print(f"  ✓ sin(x): correctly not detected (requires expression tree search)")
    else:
        print(f"  ~ sin(x): detected as {result.name} (R²={result.r2:.6f})")
    return True


def main():
    print("=" * 60)
    print("SilverSight Log Prescreen — Physical Law Tests")
    print("=" * 60)
    print()

    tests = [
        ("Kepler's Third Law (T = a^1.5)", test_kepler),
        ("Hooke's Law (F = kx)", test_hooke),
        ("Newton Gravity (F = 1/r²)", test_newton_gravity),
        ("Radioactive Decay (N = N₀e^(-λt))", test_radioactive_decay),
        ("Free Fall (s = ½gt²)", test_free_fall),
        ("Pendulum Period (T ∝ √L)", test_sqrt_law),
        ("Stefan-Boltzmann (P ∝ T⁴)", test_stefan_boltzmann),
        ("sin(x) — NOT a simple law", test_no_simple_law),
    ]

    passed = 0
    failed = 0
    for name, test_fn in tests:
        try:
            result = test_fn()
            if result:
                passed += 1
            else:
                failed += 1
                print(f"  ✗ {name}")
        except Exception as e:
            failed += 1
            print(f"  ✗ {name}: {e}")

    print()
    print("=" * 60)
    print(f"Results: {passed} passed, {failed} failed out of {len(tests)}")
    print("=" * 60)

    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())