#!/usr/bin/env python3
"""
log_prescreen.py — Logarithmic Pre-screening for Symbolic Regression

"Everything Is Logarithms" — multiplicative structure becomes additive under log.

Before searching expression trees, try these transforms:
  1. log-log:  log(y) vs log(x)  → power law  y = a·x^b
  2. log-y:    log(y) vs x       → exponential y = e^(a·x + c)
  3. y-x:      y vs x            → polynomial  y = a·x + b
  4. y-vs-x²:  y vs x²          → quadratic   y = a·x² + b
  5. y-vs-√x:  y vs √x          → sqrt law    y = a·√x + b

If any transform gives R² > threshold, return that as the discovered law.
This collapses O(expression tree search) to O(linear regression) for
the most common physical laws.
"""

from __future__ import annotations

import math
import numpy as np
from dataclasses import dataclass
from typing import Optional, Tuple


@dataclass
class PrescreenResult:
    """Result from log prescreening."""
    name: str               # "power_law", "exponential", "polynomial", etc.
    expression: str         # human-readable: "y = 1.000 * x^1.500"
    r2: float               # R² in original space
    r2_transform: float     # R² in transformed space
    params: Tuple[float, ...]  # (a, b, ...) coefficients
    transform: str          # "log-log", "log-y", "y-x", etc.


def _linear_r2(x: np.ndarray, y: np.ndarray) -> Tuple[float, float, float]:
    """Linear regression y = a·x + b. Returns (a, b, R²)."""
    if len(x) < 2:
        return 0.0, 0.0, 0.0
    coeffs = np.polyfit(x, y, 1)
    a, b = coeffs
    y_pred = a * x + b
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 1e-15 else 0.0
    return float(a), float(b), float(r2)


def _multivar_r2(X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, float]:
    """Multiple linear regression y = X·coeffs. Returns (coeffs, R²)."""
    if X.shape[0] < X.shape[1] + 1:
        return np.zeros(X.shape[1]), 0.0
    # Least squares: coeffs = (X^T X)^-1 X^T y
    try:
        coeffs = np.linalg.lstsq(X, y, rcond=None)[0]
        y_pred = X @ coeffs
        ss_res = np.sum((y - y_pred) ** 2)
        ss_tot = np.sum((y - np.mean(y)) ** 2)
        r2 = 1.0 - ss_res / ss_tot if ss_tot > 1e-15 else 0.0
        return coeffs, float(r2)
    except np.linalg.LinAlgError:
        return np.zeros(X.shape[1]), 0.0


def prescreen(
    x: np.ndarray,
    y: np.ndarray,
    feature_names: list[str] = None,
    r2_threshold: float = 0.999,
) -> Optional[PrescreenResult]:
    """Try log transforms before expression tree search.

    Returns PrescreenResult if any transform gives R² > threshold.
    Returns None if no simple law found (caller should use full search).

    Args:
        x: input array (n,) or (n, d) for multivariate
        y: target array (n,)
        feature_names: optional names for features
        r2_threshold: minimum R² to accept (default 0.999)

    Returns:
        PrescreenResult or None
    """
    x = np.asarray(x, dtype=np.float64)
    y = np.asarray(y, dtype=np.float64)

    if feature_names is None:
        if x.ndim == 1:
            feature_names = ["x"]
        else:
            feature_names = [f"x{i}" for i in range(x.shape[1])]

    # Ensure 2D
    if x.ndim == 1:
        x = x.reshape(-1, 1)

    n = len(y)
    if n < 3:
        return None

    results = []

    # ── 0. Linear: y = a·x + b (CHECK FIRST — power law with exp=1 is also linear) ──
    for j, name in enumerate(feature_names):
        a, b, r2 = _linear_r2(x[:, j], y)
        if r2 > r2_threshold:
            results.append(PrescreenResult(
                name="linear",
                expression=f"{a:.6f}*{name} + {b:.6f}",
                r2=r2,
                r2_transform=r2,
                params=(a, b),
                transform="y-x",
            ))

    # ── 1. Power law: y = a·x^b  (log-log: log(y) = b·log(x) + log(a)) ──
    if np.all(x > 0) and np.all(y > 0):
        log_x = np.log(x)
        log_y = np.log(y)
        for j, name in enumerate(feature_names):
            a, b, r2 = _linear_r2(log_x[:, j], log_y)
            if r2 > r2_threshold:
                # y = e^b · x^a  (a is slope, b is intercept)
                coeff = np.exp(b)
                exponent = a
                y_pred = coeff * x[:, j] ** exponent
                r2_orig = 1 - np.sum((y - y_pred)**2) / np.sum((y - np.mean(y))**2)
                results.append(PrescreenResult(
                    name="power_law",
                    expression=f"{coeff:.6f} * {name}^{exponent:.6f}",
                    r2=float(r2_orig),
                    r2_transform=r2,
                    params=(coeff, exponent),
                    transform="log-log",
                ))

    # ── 2. Exponential: y = e^(a·x + c)  (log-y: log(y) = a·x + c) ──
    if np.all(y > 0):
        log_y = np.log(y)
        for j, name in enumerate(feature_names):
            a, c, r2 = _linear_r2(x[:, j], log_y)
            if r2 > r2_threshold:
                y_pred = np.exp(a * x[:, j] + c)
                r2_orig = 1 - np.sum((y - y_pred)**2) / np.sum((y - np.mean(y))**2)
                results.append(PrescreenResult(
                    name="exponential",
                    expression=f"exp({a:.6f}*{name} + {c:.6f})",
                    r2=float(r2_orig),
                    r2_transform=r2,
                    params=(a, c),
                    transform="log-y",
                ))

    # ── 4. Quadratic: y = a·x² + b·x + c ──
    for j, name in enumerate(feature_names):
        X_quad = np.column_stack([x[:, j]**2, x[:, j], np.ones(n)])
        coeffs, r2 = _multivar_r2(X_quad, y)
        if r2 > r2_threshold:
            a, b, c = coeffs
            results.append(PrescreenResult(
                name="quadratic",
                expression=f"{a:.6f}*{name}² + {b:.6f}*{name} + {c:.6f}",
                r2=r2,
                r2_transform=r2,
                params=(a, b, c),
                transform="y-x²",
            ))

    # ── 5. Sqrt: y = a·√x + b ──
    if np.all(x >= 0):
        for j, name in enumerate(feature_names):
            sqrt_x = np.sqrt(x[:, j])
            a, b, r2 = _linear_r2(sqrt_x, y)
            if r2 > r2_threshold:
                results.append(PrescreenResult(
                    name="sqrt_law",
                    expression=f"{a:.6f}*sqrt({name}) + {b:.6f}",
                    r2=r2,
                    r2_transform=r2,
                    params=(a, b),
                    transform="y-√x",
                ))

    # ── 6. Multivariate power law: y = a·x₁^b₁·x₂^b₂·... ──
    if x.shape[1] > 1 and np.all(x > 0) and np.all(y > 0):
        log_x_all = np.log(x)
        log_y = np.log(y)
        X_log = np.column_stack([log_x_all, np.ones(n)])
        coeffs, r2 = _multivar_r2(X_log, log_y)
        if r2 > r2_threshold:
            exponents = coeffs[:-1]
            intercept = coeffs[-1]
            coeff = np.exp(intercept)
            terms = [f"{name}^{e:.4f}" for name, e in zip(feature_names, exponents)]
            expression = f"{coeff:.6f} * " + " * ".join(terms)
            results.append(PrescreenResult(
                name="multivariate_power_law",
                expression=expression,
                r2=r2,
                r2_transform=r2,
                params=tuple(coeffs),
                transform="log-log-multivar",
            ))

    # Return best result
    if not results:
        return None

    best = max(results, key=lambda r: r.r2)
    if best.r2 >= r2_threshold:
        return best
    return None


def prescreen_or_search(
    x: np.ndarray,
    y: np.ndarray,
    feature_names: list[str] = None,
    r2_threshold: float = 0.999,
) -> Tuple[str, float, str]:
    """Try prescreen first, return (expression, r2, method).

    If prescreen fails, returns ("", 0.0, "none") — caller should use full search.
    """
    result = prescreen(x, y, feature_names, r2_threshold)
    if result is not None:
        return result.expression, result.r2, result.name
    return "", 0.0, "none"