#!/usr/bin/env python3
"""
test_kepler.py — Test SilverSight symbolic regression on Kepler's Third Law

Given 8 planets' distance (a) and orbital period (T):
  T = a^1.5  (Kepler's Third Law, 1618)

This is the same test GP-ELITE uses. We must rediscover a^1.5 from raw data.
"""

import sys
import math
import numpy as np
from itertools import product

sys.path.insert(0, "/home/allaun/SilverSight/python")

from expr_tree import (
    ExprNode, var_x, const, unary, binary,
    x_squared, x_cubed, sqrt_x, x_sqrt_x, sin_x, cos_x,
)
from linear_scaling import linear_scale, compute_r2, compute_mse


# ── Kepler data (NASA) ─────────────────────────────────────────────────

planets = ["Mercury","Venus","Earth","Mars","Jupiter","Saturn","Uranus","Neptune"]
a = np.array([0.387, 0.723, 1.000, 1.524, 5.203, 9.537, 19.191, 30.069])   # AU
T = np.array([0.241, 0.615, 1.000, 1.881, 11.862, 29.457, 84.011, 164.79]) # years


# ── BIC fitness ─────────────────────────────────────────────────────────

def bic_fitness(expr: ExprNode, xs: np.ndarray, ys: np.ndarray) -> float:
    """BIC = n·ln(MSE) + k·ln(n)

    Lower is better. Includes linear scaling.
    """
    n = len(ys)
    try:
        raw = expr.evaluate(xs)
        if not np.all(np.isfinite(raw)):
            return 1e15
        _, _, scaled = linear_scale(raw, ys)
        mse = compute_mse(scaled, ys)
        if mse < 1e-15:
            mse = 1e-15
        k = expr.complexity()
        return n * math.log(mse) + k * math.log(n)
    except Exception:
        return 1e15


# ── Expression candidates ──────────────────────────────────────────────

def build_candidates() -> list[tuple[str, ExprNode]]:
    """Build a library of candidate expressions."""
    x = var_x()
    candidates = []

    # Simple powers
    candidates.append(("x", x))
    candidates.append(("x^2", x_squared()))
    candidates.append(("x^3", x_cubed()))
    candidates.append(("sqrt(x)", sqrt_x()))
    candidates.append(("x*sqrt(x)", x_sqrt_x()))
    candidates.append(("x^1.5 (explicit)", binary("pow", x, const(1.5))))

    # Trig
    candidates.append(("sin(x)", sin_x()))
    candidates.append(("cos(x)", cos_x()))

    # Combinations
    candidates.append(("x^2 + x", binary("+", x_squared(), x)))
    candidates.append(("x^2 * x", binary("*", x_squared(), x)))
    candidates.append(("x * x * x", binary("*", binary("*", x, x), x)))
    candidates.append(("sqrt(x) * x", binary("*", sqrt_x(), x)))
    candidates.append(("x * x^0.5", binary("*", x, binary("pow", x, const(0.5)))))
    candidates.append(("x^0.5 * x", binary("*", binary("pow", x, const(0.5)), x)))

    # Log/exp
    candidates.append(("log(x)", unary("log", x)))
    candidates.append(("exp(x)", unary("exp", x)))
    candidates.append(("exp(-x)", unary("exp", unary("neg", x))))

    # Nested
    candidates.append(("sqrt(x^2)", unary("sqrt", x_squared())))
    candidates.append(("x^2^0.5", binary("pow", x_squared(), const(0.5))))

    # Constants
    candidates.append(("1", const(1.0)))
    candidates.append(("2", const(2.0)))
    candidates.append(("0.5", const(0.5)))

    # Linear scaling will handle a,b — test shape candidates
    candidates.append(("a*x + b", x))  # linear scaling makes this a*x+b
    candidates.append(("a*x^1.5 + b", x_sqrt_x()))  # this is the target

    return candidates


# ── Chaos game search (simplified) ─────────────────────────────────────

def chaos_game_search(
    xs: np.ndarray,
    ys: np.ndarray,
    n_iterations: int = 100,
    seed: int = 42,
) -> tuple[ExprNode, float, str]:
    """Simplified chaos game search over expression space.

    Uses IFS contraction: at each step, try to improve the best expression
    by applying transformations (add constant, change operator, nest).
    """
    rng = np.random.RandomState(seed)

    # Start with best candidate from library
    candidates = build_candidates()
    best_expr = None
    best_bic = 1e15
    best_name = ""

    for name, expr in candidates:
        bic = bic_fitness(expr, xs, ys)
        if bic < best_bic:
            best_bic = bic
            best_expr = expr
            best_name = name

    # Iterative refinement via chaos game
    for iteration in range(n_iterations):
        # Generate neighborhood: modify best expression
        neighbors = _generate_neighbors(best_expr, rng, xs, ys)

        for neighbor_expr, neighbor_name in neighbors:
            bic = bic_fitness(neighbor_expr, xs, ys)
            if bic < best_bic:
                best_bic = bic
                best_expr = neighbor_expr
                best_name = neighbor_name

    return best_expr, best_bic, best_name


def _generate_neighbors(
    expr: ExprNode,
    rng: np.random.RandomState,
    xs: np.ndarray,
    ys: np.ndarray,
) -> list[tuple[ExprNode, str]]:
    """Generate neighboring expressions by small modifications."""
    neighbors = []

    # Try replacing constants with nearby values
    if expr.op == "const":
        for delta in [-0.5, -0.1, 0.1, 0.5]:
            new_val = expr.value + delta
            neighbors.append((const(new_val), f"c={new_val:.2f}"))

    # Try wrapping in unary ops
    for op in ["sqrt", "log", "exp", "sin", "cos"]:
        neighbors.append((unary(op, expr.copy()), f"{op}({expr.to_string()})"))

    # Try multiplying by x
    neighbors.append((binary("*", expr.copy(), var_x()), f"({expr.to_string()})*x"))

    # Try raising to powers
    for p in [0.5, 1.5, 2.0, 3.0]:
        neighbors.append((
            binary("pow", expr.copy(), const(p)),
            f"({expr.to_string()})^{p}"
        ))

    return neighbors


# ── Main test ──────────────────────────────────────────────────────────

def main():
    print("=" * 60)
    print("SilverSight Symbolic Regression — Kepler Test")
    print("=" * 60)
    print()
    print("Data: 8 planets (NASA)")
    for p, ai, Ti in zip(planets, a, T):
        print(f"  {p:8s}  a = {ai:7.3f} AU   T = {Ti:8.3f} yr")
    print()
    print("Target: T = a^1.5  (Kepler's Third Law)")
    print()

    # Test 1: Direct candidates
    print("─" * 60)
    print("TEST 1: Candidate Library")
    print("─" * 60)

    candidates = build_candidates()
    results = []

    for name, expr in candidates:
        try:
            raw = expr.evaluate(a)
            if not np.all(np.isfinite(raw)):
                results.append((name, expr, 1e15, 0.0, "NaN/inf"))
                continue
            a_coeff, b_coeff, scaled = linear_scale(raw, T)
            r2 = compute_r2(scaled, T)
            mse = compute_mse(scaled, T)
            bic = bic_fitness(expr, a, T)
            results.append((name, expr, bic, r2, f"a={a_coeff:.4f}, b={b_coeff:.4f}"))
        except Exception as exc:
            results.append((name, expr, 1e15, 0.0, str(exc)[:40]))

    # Sort by BIC (lower is better)
    results.sort(key=lambda x: x[2])

    print(f"\n{'Expression':<25} {'BIC':>10} {'R²':>8} {'Params':<25}")
    print("-" * 70)
    for name, expr, bic, r2, params in results[:10]:
        if bic > 1e14:
            print(f"{name:<25} {'FAIL':>10} {r2:>8.6f} {params:<25}")
        else:
            print(f"{name:<25} {bic:>10.2f} {r2:>8.6f} {params:<25}")

    # Test 2: Best candidate with linear scaling
    print()
    print("─" * 60)
    print("TEST 2: Best Candidate Details")
    print("─" * 60)

    best_name, best_expr, best_bic, best_r2, best_params = results[0]
    raw = best_expr.evaluate(a)
    a_coeff, b_coeff, scaled = linear_scale(raw, T)

    print(f"\nBest: {best_name}")
    print(f"  Expression: {best_expr.to_string()}")
    print(f"  BIC: {best_bic:.2f}")
    print(f"  R²: {best_r2:.6f}")
    print(f"  Scaling: T = {a_coeff:.4f} · g(a) + {b_coeff:.4f}")
    print(f"  Where g(a) = {best_expr.to_string()}")
    print()

    # Test 3: Verify a^1.5 specifically
    print("─" * 60)
    print("TEST 3: Verify a^1.5 (Kepler's Law)")
    print("─" * 60)

    kepler_expr = x_sqrt_x()  # x * sqrt(x) = x^1.5
    raw = kepler_expr.evaluate(a)
    a_coeff, b_coeff, scaled = linear_scale(raw, T)
    r2 = compute_r2(scaled, T)
    mse = compute_mse(scaled, T)
    bic = bic_fitness(kepler_expr, a, T)

    print(f"\n  Expression: {kepler_expr.to_string()}")
    print(f"  BIC: {bic:.2f}")
    print(f"  R²: {r2:.6f}")
    print(f"  Scaling: T = {a_coeff:.6f} · a^1.5 + {b_coeff:.6f}")
    print(f"  Expected:  T ≈ 1.0 · a^1.5 + 0.0")
    print()

    # Show predictions vs actual
    print(f"  {'Planet':<10} {'Actual':>10} {'Predicted':>10} {'Error':>10}")
    print("  " + "-" * 42)
    for p, Ti, Si in zip(planets, T, scaled):
        print(f"  {p:<10} {Ti:>10.3f} {Si:>10.3f} {abs(Ti-Si):>10.4f}")

    # Test 4: Chaos game search
    print()
    print("─" * 60)
    print("TEST 4: Chaos Game Search")
    print("─" * 60)

    best_chaos, bic_chaos, name_chaos = chaos_game_search(a, T, n_iterations=50)
    raw_chaos = best_chaos.evaluate(a)
    a_chaos, b_chaos, scaled_chaos = linear_scale(raw_chaos, T)
    r2_chaos = compute_r2(scaled_chaos, T)

    print(f"\n  Best from chaos game: {name_chaos}")
    print(f"  Expression: {best_chaos.to_string()}")
    print(f"  BIC: {bic_chaos:.2f}")
    print(f"  R²: {r2_chaos:.6f}")
    print(f"  Scaling: T = {a_chaos:.6f} · g(a) + {b_chaos:.6f}")

    # Summary
    print()
    print("=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(f"  Target: T = a^1.5 (Kepler's Third Law)")
    print(f"  Data: 8 planets")
    print(f"  Best candidate: {best_name}")
    print(f"  Best R²: {best_r2:.6f}")
    print(f"  a^1.5 R²: {r2:.6f}")

    if best_r2 > 0.999:
        print()
        print("  ✓ PASS: SilverSight rediscovers Kepler's Third Law")
    elif r2 > 0.999:
        print()
        print("  ✓ PASS: a^1.5 fits with R² > 0.999")
    else:
        print()
        print(f"  ✗ FAIL: Best R² = {best_r2:.6f} < 0.999")

    return 0


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