"""PySR formula evaluation for closed-loop CFD validation.

Wraps PySR sympy expressions into v23-structured prediction functions
that plug into the existing predict_v23_deriv validation framework.

Usage:
    from predict_pysr import pysr_predict_v23

    front_func, top_func = build_pysr_functions(front_json_path, top_json_path)
    omega = pysr_predict_v23(obs_slice, a_prev_phys, ..., front_func, top_func, ...)
"""
from __future__ import annotations

import json
from typing import Callable, Dict, List, Optional, Tuple

import numpy as np
import sympy

from SR_analysis.utils.feature_builder import (
    compute_dimensionless, compute_features, build_feature_matrix,
    apply_G_x,
)
from SR_analysis.utils.g_operator import apply_G_raw


def load_pysr_formulas(front_json_path: str, top_json_path: str) -> Tuple[str, str, List[str], List[str]]:
    """Load PySR formulas from saved JSON results.

    Returns
    -------
    front_expr : str
        sympy expression string for front channel.
    top_expr : str
        sympy expression string for top channel.
    front_feat_names : list of str
        Feature names for front (no bias column).
    rear_feat_names : list of str
        Feature names for rear (WITH 'bias' as first element).
    """
    with open(front_json_path) as f:
        front_data = json.load(f)
    with open(top_json_path) as f:
        top_data = json.load(f)

    front_expr = front_data["best_sympy"]
    top_expr = top_data["best_sympy"]
    front_feat_names = front_data["feature_names"]       # no bias
    rear_feat_names = top_data["feature_names"]          # includes 'bias'

    return front_expr, top_expr, front_feat_names, rear_feat_names


def build_pysr_functions(
    front_expr: str,
    top_expr: str,
    front_feat_names: List[str],
    rear_feat_names: List[str],
) -> Tuple[Callable, Callable]:
    """Build numpy callables from PySR sympy expressions.

    Returns
    -------
    front_func : callable
        f(*feature_values) -> float  (no bias, pure feature values in order)
    top_func : callable
        f(*feature_values_including_bias) -> float (bias is first element)
    """
    front_func = sympy.lambdify(front_feat_names, front_expr, modules="numpy")
    # Rear includes 'bias' as first feature in variable_names
    top_func = sympy.lambdify(rear_feat_names, top_expr, modules="numpy")
    return front_func, top_func


def pysr_predict_v23(
    obs_slice: np.ndarray,
    a_prev_phys: np.ndarray,
    a_prev2_phys: np.ndarray,
    mu: float,
    u0: float,
    dt_c: float,
    front_func: Callable,
    top_func: Callable,
    feat_keys_front: List[str],
    feat_keys_rear: List[str],
    target_forces: Optional[np.ndarray] = None,
    sensors_raw: Optional[np.ndarray] = None,
    forces_raw: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Predict action using v23 structure with PySR formulas.

    Parameters (matching predict_v23_deriv signature):
        obs_slice: (12,) current obs [sensors(6) + forces(6)]
        a_prev_phys: (3,) previous physical omega
        a_prev2_phys: (3,) twice-lagged physical omega
        feat_keys_front: feature names for front (no bias)
        feat_keys_rear: feature names for rear (WITHOUT 'bias' — bias added by build_feature_matrix)

    Returns
    -------
        omega_new: (3,) physical omega [front, bottom, top]
    """
    has_mu = any("mu" in k for k in feat_keys_front)

    # Check if augmented features (derivatives or lags) are needed
    needs_aug = any(k.endswith("_dt") or k.endswith("_lag1") for k in feat_keys_front)

    # Build sensor/force arrays
    if needs_aug and sensors_raw is not None and forces_raw is not None:
        # Use provided raw arrays (2 rows: prev + current)
        sensors_curr = obs_slice[0:6].astype(np.float64).reshape(1, 6)
        forces_curr = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    else:
        sensors_raw = obs_slice[0:6].astype(np.float64).reshape(1, 6)
        forces_raw = obs_slice[6:12].astype(np.float64).reshape(1, 6)
        sensors_curr = sensors_raw
        forces_curr = forces_raw

    ap = a_prev_phys.astype(np.float64).reshape(1, 3)
    ap2 = a_prev2_phys.astype(np.float64).reshape(1, 3)

    dim = compute_dimensionless(sensors_curr, forces_curr, u0=u0, d=20.0)
    tf = target_forces.reshape(1, 2) if target_forces is not None else None

    sym = compute_features(
        dim, ap, ap2, mu,
        alpha_mode=False,
        include_mu=has_mu,
        include_cos_sin=False,
        u0=u0,
        target_forces=tf,
        sensors_raw=sensors_raw,
        forces_raw=forces_raw,
    )

    # Front: no bias
    fv_front = build_feature_matrix(sym, feat_keys_front, add_bias=False)[0]
    pred_F = float(front_func(*fv_front))

    # Top: with bias
    fv_top = build_feature_matrix(sym, feat_keys_rear, add_bias=True)[0]
    pred_T = float(top_func(*fv_top))

    # Bottom = -top(Gx) using shared-head
    G_obs, G_a_prev, G_a_prev2 = apply_G_raw(obs_slice, a_prev_phys, a_prev2_phys)
    G_target = np.array([target_forces[0], -target_forces[1]]) if target_forces is not None else None
    G_sensors = G_obs[0:6].astype(np.float64).reshape(1, 6)
    G_forces = G_obs[6:12].astype(np.float64).reshape(1, 6)
    G_dim = compute_dimensionless(G_sensors, G_forces, u0=u0, d=20.0)
    G_sym = compute_features(
        G_dim, G_a_prev.reshape(1, 3), G_a_prev2.reshape(1, 3),
        mu, alpha_mode=False, include_mu=has_mu,
        include_cos_sin=False, u0=u0, target_forces=G_target,
    )
    fv_bot = build_feature_matrix(G_sym, feat_keys_rear, add_bias=True)[0]
    pred_B = -float(top_func(*fv_bot))

    # Convert from alpha (non-dim) to physical omega
    alpha_new = np.array([pred_F, pred_B, pred_T])
    omega_new = alpha_new * u0

    return omega_new.astype(np.float64)
