agents/math_agent.py

import os
import logging
from typing import List, Dict

import sympy as sp
import numpy as np
import scipy.linalg as la
import scipy.special as special
from llama_index.tools.code_interpreter import CodeInterpreterToolSpec
from scipy.integrate import quad
from scipy.stats import binom, norm, poisson
import numpy.fft as fft

from llama_index.core.agent.workflow import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.tools.wolfram_alpha import WolframAlphaToolSpec

# Setup logging
logger = logging.getLogger(__name__)

# --- Math Tool Functions (with enhanced logging and error handling) ---

# Helper decorator for error handling and logging
def math_tool_handler(func):
    def wrapper(*args, **kwargs):
        func_name = func.__name__
        logger.info(f"Executing math tool: {func_name} with args: {args}, kwargs: {kwargs}")
        try:
            result = func(*args, **kwargs)
            logger.info(f"Tool {func_name} executed successfully. Result: {str(result)[:200]}...")
            # Ensure result is serializable (convert numpy types if necessary)
            if isinstance(result, np.ndarray):
                return result.tolist()
            if isinstance(result, (np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32, np.uint64)):
                 return int(result)
            if isinstance(result, (np.float_, np.float16, np.float32, np.float64)):
                 return float(result)
            if isinstance(result, (np.complex_, np.complex64, np.complex128)):
                 return complex(result)
            if isinstance(result, np.bool_):
                 return bool(result)
            if isinstance(result, dict):
                 return {k: wrapper(v) if isinstance(v, (np.ndarray, np.number, np.bool_)) else v for k, v in result.items()} # Recursively handle dicts
            return result
        except (sp.SympifyError, TypeError, ValueError, np.linalg.LinAlgError, ZeroDivisionError) as e:
            logger.warning(f"Math error in {func_name}: {e}")
            return f"Error in {func_name}: {e}"
        except Exception as e:
            logger.error(f"Unexpected error in {func_name}: {e}", exc_info=True)
            return f"Unexpected error in {func_name}: {e}"
    return wrapper

# --- Symbolic math functions ---
@math_tool_handler
def solve_symbolic_equation(equation: str, variable: str = "x") -> str:
    """Solve a symbolic equation (e.g., 'x**2 - 4') for the given variable."""
    symbol = sp.symbols(variable)
    # Ensure equation is treated as expression == 0 if no equality sign
    if "=" not in equation:
        expr = sp.sympify(equation)
    else:
        lhs, rhs = equation.split("=", 1)
        expr = sp.Eq(sp.sympify(lhs.strip()), sp.sympify(rhs.strip()))
    solutions = sp.solve(expr, symbol)
    return f"Solutions: {solutions}"

@math_tool_handler
def compute_derivative(expression: str, variable: str = "x") -> str:
    """Compute the symbolic derivative of an expression (e.g., 'sin(x)*x**2')."""
    symbol = sp.symbols(variable)
    expr = sp.sympify(expression)
    deriv = sp.diff(expr, symbol)
    return f"Derivative: {deriv}"

@math_tool_handler
def compute_integral(expression: str, variable: str = "x") -> str:
    """Compute the symbolic indefinite integral of an expression (e.g., '1/x')."""
    symbol = sp.symbols(variable)
    expr = sp.sympify(expression)
    integ = sp.integrate(expr, symbol)
    return f"Integral: {integ} + C"

@math_tool_handler
def compute_limit(
    expression: str, variable: str = "x", point: str = "oo"
) -> str:
    """Compute the limit of an expression (e.g., 'sin(x)/x') as variable approaches point (e.g., '0', 'oo')."""
    symbol = sp.symbols(variable)
    expr = sp.sympify(expression)
    # Handle 'oo', '-oo', 'zoo' for infinity, or numerical points
    if point.lower() == "oo":
        pt = sp.oo
    elif point.lower() == "-oo":
        pt = -sp.oo
    elif point.lower() == "zoo":
        pt = sp.zoo # Complex infinity
    else:
        pt = sp.sympify(point)
    lim = sp.limit(expr, symbol, pt)
    return f"Limit at {point}: {lim}"

@math_tool_handler
def simplify_expression(expression: str) -> str:
    """Simplify a symbolic expression (e.g., 'sin(x)**2 + cos(x)**2')."""
    expr = sp.sympify(expression)
    simp = sp.simplify(expr)
    return f"Simplified expression: {simp}"

@math_tool_handler
def expand_expression(expression: str) -> str:
    """Expand a symbolic expression (e.g., '(x+y)**2')."""
    expr = sp.sympify(expression)
    exp = sp.expand(expr)
    return f"Expanded expression: {exp}"

@math_tool_handler
def factor_expression(expression: str) -> str:
    """Factor a symbolic expression (e.g., 'x**2 - y**2')."""
    expr = sp.sympify(expression)
    fact = sp.factor(expr)
    return f"Factored expression: {fact}"


# --- Matrix math functions ---
@math_tool_handler
def matrix_addition(a: List[List[float]], b: List[List[float]]) -> List[List[float]]:
    """Add two matrices element-wise. Input: [[1, 2], [3, 4]], [[5, 6], [7, 8]]."""
    A = np.array(a)
    B = np.array(b)
    if A.shape != B.shape:
        raise ValueError("Matrices must have the same shape for addition.")
    return (A + B)

@math_tool_handler
def matrix_subtraction(a: List[List[float]], b: List[List[float]]) -> List[List[float]]:
    """Subtract matrix B from matrix A element-wise. Input: [[5, 6], [7, 8]], [[1, 2], [3, 4]]."""
    A = np.array(a)
    B = np.array(b)
    if A.shape != B.shape:
        raise ValueError("Matrices must have the same shape for subtraction.")
    return (A - B)

@math_tool_handler
def matrix_multiplication(a: List[List[float]], b: List[List[float]]) -> List[List[float]]:
    """Multiply two matrices. Input: [[1, 2], [3, 4]], [[5, 6], [7, 8]]."""
    A = np.array(a)
    B = np.array(b)
    if A.shape[1] != B.shape[0]:
        raise ValueError("Inner dimensions must match for matrix multiplication.")
    return np.matmul(A, B)

@math_tool_handler
def matrix_inverse(matrix: List[List[float]]) -> List[List[float]]:
    """Compute the inverse of a square matrix. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    if M.shape[0] != M.shape[1]:
        raise ValueError("Matrix must be square to compute inverse.")
    return np.linalg.inv(M)

@math_tool_handler
def matrix_determinant(matrix: List[List[float]]) -> float:
    """Compute the determinant of a square matrix. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    if M.shape[0] != M.shape[1]:
        raise ValueError("Matrix must be square to compute determinant.")
    return np.linalg.det(M)

@math_tool_handler
def matrix_transpose(matrix: List[List[float]]) -> List[List[float]]:
    """Transpose a matrix. Input: [[1, 2, 3], [4, 5, 6]]."""
    M = np.array(matrix)
    return M.T

@math_tool_handler
def matrix_rank(matrix: List[List[float]]) -> int:
    """Compute the rank of a matrix. Input: [[1, 2], [2, 4]]."""
    M = np.array(matrix)
    return np.linalg.matrix_rank(M)

@math_tool_handler
def matrix_trace(matrix: List[List[float]]) -> float:
    """Compute the trace of a square matrix. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    if M.shape[0] != M.shape[1]:
        raise ValueError("Matrix must be square to compute trace.")
    return np.trace(M)

@math_tool_handler
def matrix_norm(matrix: List[List[float]], ord_str: str = "fro") -> float:
    """Compute the norm of a matrix. ord_str can be 'fro' (Frobenius), 'nuc' (nuclear), inf, -inf, 1, -1, 2, -2. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    ord_map = {"fro": "fro", "nuc": "nuc", "inf": np.inf, "-inf": -np.inf, "1": 1, "-1": -1, "2": 2, "-2": -2}
    ord_val = ord_map.get(ord_str)
    if ord_val is None:
        raise ValueError(f"Invalid ord_str: {ord_str}. Must be one of {list(ord_map.keys())}")
    return np.linalg.norm(M, ord=ord_val)

@math_tool_handler
def eigenvalues(matrix: List[List[float]]) -> List[complex]:
    """Compute eigenvalues of a square matrix. Input: [[1, -1], [1, 1]]."""
    M = np.array(matrix)
    if M.shape[0] != M.shape[1]:
        raise ValueError("Matrix must be square to compute eigenvalues.")
    vals = np.linalg.eigvals(M)
    return vals

@math_tool_handler
def eigenvectors(matrix: List[List[float]]) -> List[List[complex]]:
    """Compute eigenvectors of a square matrix. Returns list of eigenvectors. Input: [[1, -1], [1, 1]]."""
    M = np.array(matrix)
    if M.shape[0] != M.shape[1]:
        raise ValueError("Matrix must be square to compute eigenvectors.")
    vals, vecs = np.linalg.eig(M)
    # Return eigenvectors as rows or columns? Let's return as list of column vectors
    return vecs.T # Transpose to get eigenvectors as list items

@math_tool_handler
def svd_decompose(matrix: List[List[float]]) -> Dict[str, List]:
    """Compute the singular value decomposition (U, S, Vh) of a matrix. Input: [[1, 2], [3, 4], [5, 6]]."""
    M = np.array(matrix)
    U, S, Vh = np.linalg.svd(M)
    return {"U": U, "S": S, "Vh": Vh}

@math_tool_handler
def lu_decompose(matrix: List[List[float]]) -> Dict[str, List]:
    """Compute the LU decomposition (P, L, U) of a matrix. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    P, L, U = la.lu(M)
    return {"P": P, "L": L, "U": U}

@math_tool_handler
def qr_decompose(matrix: List[List[float]]) -> Dict[str, List]:
    """Compute the QR decomposition (Q, R) of a matrix. Input: [[1, 2], [3, 4]]."""
    M = np.array(matrix)
    Q, R = np.linalg.qr(M)
    return {"Q": Q, "R": R}

# --- Statistics functions ---
@math_tool_handler
def mean(values: List[float]) -> float:
    """Compute the mean of a list of numbers. Input: [1, 2, 3, 4, 5]."""
    if not values:
        raise ValueError("Input list cannot be empty for mean calculation.")
    return np.mean(np.array(values))

@math_tool_handler
def median(values: List[float]) -> float:
    """Compute the median of a list of numbers. Input: [1, 3, 2, 4, 5]."""
    if not values:
        raise ValueError("Input list cannot be empty for median calculation.")
    return np.median(np.array(values))

@math_tool_handler
def std_dev(values: List[float], ddof: int = 1) -> float:
    """Compute the sample standard deviation (ddof=1) or population (ddof=0) of a list. Input: [1, 2, 3, 4, 5]."""
    if not values or len(values) < ddof:
         raise ValueError(f"Input list must have at least {ddof} elements for std dev with ddof={ddof}.")
    return np.std(np.array(values), ddof=ddof)

@math_tool_handler
def variance(values: List[float], ddof: int = 1) -> float:
    """Compute the sample variance (ddof=1) or population (ddof=0) of a list. Input: [1, 2, 3, 4, 5]."""
    if not values or len(values) < ddof:
         raise ValueError(f"Input list must have at least {ddof} elements for variance with ddof={ddof}.")
    return np.var(np.array(values), ddof=ddof)

@math_tool_handler
def percentile(values: List[float], percent: float) -> float:
    """Compute the q-th percentile (0<=q<=100) of a list. Input: [1, 2, 3, 4, 5], 75."""
    if not values:
        raise ValueError("Input list cannot be empty for percentile calculation.")
    if not (0 <= percent <= 100):
        raise ValueError("Percent must be between 0 and 100.")
    return np.percentile(np.array(values), percent)

@math_tool_handler
def covariance(x: List[float], y: List[float], ddof: int = 1) -> float:
    """Compute sample covariance (ddof=1) or population (ddof=0) between two lists. Input: [1, 2, 3], [4, 5, 6]."""
    X = np.array(x)
    Y = np.array(y)
    if X.size != Y.size:
        raise ValueError("Input lists must have the same length for covariance.")
    if X.size == 0 or X.size < ddof:
         raise ValueError(f"Input lists must have at least {ddof} elements for covariance with ddof={ddof}.")
    # np.cov returns the covariance matrix, we want the off-diagonal element
    return np.cov(X, Y, ddof=ddof)[0, 1]

@math_tool_handler
def correlation(x: List[float], y: List[float]) -> float:
    """Compute Pearson correlation coefficient between two lists. Input: [1, 2, 3], [1, 2, 3.1]."""
    X = np.array(x)
    Y = np.array(y)
    if X.size != Y.size:
        raise ValueError("Input lists must have the same length for correlation.")
    if X.size < 2:
        raise ValueError("Need at least 2 data points for correlation.")
    # np.corrcoef returns the correlation matrix
    corr_matrix = np.corrcoef(X, Y)
    # Handle case where std dev is zero (results in nan)
    if np.isnan(corr_matrix[0, 1]):
        logger.warning("Correlation resulted in NaN, likely due to zero standard deviation in one or both inputs.")
        # Return 0 or raise error? Let's return 0 for now.
        return 0.0
    return corr_matrix[0, 1]

@math_tool_handler
def linear_regression(x: List[float], y: List[float]) -> Dict[str, float]:
    """Perform simple linear regression (y = mx + c). Returns slope (m) and intercept (c). Input: [1, 2, 3], [2, 4.1, 5.9]."""
    X = np.array(x)
    Y = np.array(y)
    if X.size != Y.size:
        raise ValueError("Input lists must have the same length for linear regression.")
    if X.size < 2:
        raise ValueError("Need at least 2 data points for linear regression.")
    slope, intercept = np.polyfit(X, Y, 1)
    return {"slope": slope, "intercept": intercept}

# --- Numerical functions ---
@math_tool_handler
def find_polynomial_roots(coefficients: List[float]) -> List[complex]:
    """Find roots of a polynomial given coefficients [a_n, a_n-1, ..., a_0]. Input: [1, -3, 2] for x^2-3x+2."""
    if not coefficients:
        raise ValueError("Coefficient list cannot be empty.")
    return np.roots(coefficients)

@math_tool_handler
def interpolate_value(x_vals: List[float], y_vals: List[float], x: float) -> float:
    """Linear interpolate a value at x given data points (x_vals, y_vals). Input: [0, 1, 2], [0, 1, 4], 1.5."""
    if len(x_vals) != len(y_vals):
        raise ValueError("x_vals and y_vals must have the same length.")
    if len(x_vals) < 2:
        raise ValueError("Need at least 2 data points for interpolation.")
    # Ensure x_vals are sorted for np.interp
    sorted_indices = np.argsort(x_vals)
    x_sorted = np.array(x_vals)[sorted_indices]
    y_sorted = np.array(y_vals)[sorted_indices]
    return np.interp(x, x_sorted, y_sorted)

@math_tool_handler
def numerical_integration(
    func_str: str, a: float, b: float, variable: str = "x"
) -> float:
    """Numerically integrate func_str (e.g., 'x**2 * sin(x)') from a to b. Input: 'x**2', 0, 1."""
    symbol = sp.symbols(variable)
    # Security Note: Using sympify/lambdify can be risky if func_str is untrusted.
    # Consider using a safer evaluation method if input is external.
    try:
        func = sp.sympify(func_str)
        f_lambdified = sp.lambdify(symbol, func, modules=["numpy"])
    except (sp.SympifyError, SyntaxError) as sym_err:
        raise ValueError(f"Invalid function string: {func_str}. Error: {sym_err}")
    
    result, abserr = quad(f_lambdified, a, b)
    logger.info(f"Numerical integration estimated absolute error: {abserr}")
    return result

@math_tool_handler
def solve_ode(
    func_str: str, y0: float, t_eval: List[float], args: tuple = ()
) -> List[float]:
    """Solve a first-order ODE dy/dt = f(t, y) using scipy.integrate.solve_ivp.
       func_str should define f(t, y), e.g., '-y + sin(t)'.
       y0 is the initial condition y(t_eval[0]).
       t_eval is the list of time points to evaluate the solution at.
       args are optional additional arguments passed to f(t, y, *args).
       Input: func_str='-y', y0=1, t_eval=[0, 1, 2, 3, 4]."""
    from scipy.integrate import solve_ivp
    import math # Make math functions available
    
    # Security Note: Using eval is dangerous with untrusted input.
    # A safer approach would parse the expression or use a restricted environment.
    def ode_func(t, y, *args):
        try:
            # Provide t, y, args, and safe math functions in the eval context
            local_vars = {"t": t, "y": y, "math": math, "np": np}
            # Add args if provided
            if args:
                # Assuming args correspond to p1, p2, ... in the func_str
                for i, arg_val in enumerate(args):
                    local_vars[f"p{i+1}"] = arg_val 
            return eval(func_str, {"__builtins__": {}}, local_vars)
        except Exception as e:
            # Log the error and raise it to be caught by the handler
            logger.error(f"Error evaluating ODE function {func_str} at t={t}, y={y}: {e}")
            raise ValueError(f"Error in ODE function definition: {e}")

    if not t_eval:
        raise ValueError("t_eval list cannot be empty.")
    t_span = (min(t_eval), max(t_eval))
    
    sol = solve_ivp(ode_func, t_span, [y0], t_eval=t_eval, args=args)
    
    if not sol.success:
        raise RuntimeError(f"ODE solver failed: {sol.message}")
        
    return sol.y[0] # Return the solution for y

# --- Vector functions ---
@math_tool_handler
def dot_product(a: List[float], b: List[float]) -> float:
    """Compute dot product of two vectors. Input: [1, 2, 3], [4, 5, 6]."""
    A = np.array(a)
    B = np.array(b)
    if A.shape != B.shape:
        raise ValueError("Vectors must have the same dimension for dot product.")
    return np.dot(A, B)

@math_tool_handler
def cross_product(a: List[float], b: List[float]) -> List[float]:
    """Compute cross product of two 3D vectors. Input: [1, 0, 0], [0, 1, 0]."""
    A = np.array(a)
    B = np.array(b)
    if A.size != 3 or B.size != 3:
        raise ValueError("Cross product is only defined for 3D vectors.")
    return np.cross(A, B)

@math_tool_handler
def vector_magnitude(a: List[float]) -> float:
    """Compute magnitude (Euclidean norm) of a vector. Input: [3, 4]."""
    if not a:
        raise ValueError("Input vector cannot be empty.")
    return np.linalg.norm(np.array(a))

@math_tool_handler
def vector_normalize(a: List[float]) -> List[float]:
    """Normalize a vector to unit length. Input: [3, 4]."""
    A = np.array(a)
    norm = np.linalg.norm(A)
    if norm == 0:
        raise ValueError("Cannot normalize a zero vector.")
    return (A / norm)

@math_tool_handler
def vector_angle(a: List[float], b: List[float], degrees: bool = False) -> float:
    """Compute the angle (in radians or degrees) between two vectors. Input: [1, 0], [0, 1]."""
    dot = dot_product(a, b) # Use our handled dot_product
    norm_a = vector_magnitude(a)
    norm_b = vector_magnitude(b)
    if norm_a == 0 or norm_b == 0:
        raise ValueError("Cannot compute angle with zero vector(s).")
    # Clip argument to arccos to avoid domain errors due to floating point inaccuracies
    cos_theta = np.clip(dot / (norm_a * norm_b), -1.0, 1.0)
    angle_rad = np.arccos(cos_theta)
    return np.degrees(angle_rad) if degrees else angle_rad

# --- Probability functions ---
@math_tool_handler
def binomial_pmf(k: int, n: int, p: float) -> float:
    """Compute binomial probability mass function P(X=k | n, p). Input: k=2, n=5, p=0.5."""
    if not (0 <= p <= 1):
        raise ValueError("Probability p must be between 0 and 1.")
    if not (0 <= k <= n):
        raise ValueError("k must be between 0 and n (inclusive).")
    return binom.pmf(k, n, p)

@math_tool_handler
def normal_pdf(x: float, mu: float = 0, sigma: float = 1) -> float:
    """Compute normal distribution probability density function N(x | mu, sigma). Input: x=0, mu=0, sigma=1."""
    if sigma <= 0:
        raise ValueError("Standard deviation sigma must be positive.")
    return norm.pdf(x, mu, sigma)

@math_tool_handler
def normal_cdf(x: float, mu: float = 0, sigma: float = 1) -> float:
    """Compute normal distribution cumulative distribution function P(X<=x | mu, sigma). Input: x=0, mu=0, sigma=1."""
    if sigma <= 0:
        raise ValueError("Standard deviation sigma must be positive.")
    return norm.cdf(x, mu, sigma)

@math_tool_handler
def poisson_pmf(k: int, lam: float) -> float:
    """Compute Poisson probability mass function P(X=k | lambda). Input: k=2, lam=3."""
    if lam < 0:
        raise ValueError("Rate parameter lambda must be non-negative.")
    if k < 0 or not isinstance(k, int):
        raise ValueError("k must be a non-negative integer.")
    return poisson.pmf(k, lam)

# --- Special functions ---
@math_tool_handler
def gamma_function(x: float) -> float:
    """Compute the gamma function Gamma(x). Input: 5."""
    return special.gamma(x)

@math_tool_handler
def beta_function(x: float, y: float) -> float:
    """Compute the beta function B(x, y). Input: 2, 3."""
    return special.beta(x, y)

@math_tool_handler
def erf_function(x: float) -> float:
    """Compute the error function erf(x). Input: 1."""
    return special.erf(x)

# --- Fourier Transform functions ---
@math_tool_handler
def fft_transform(y: List[float]) -> List[complex]:
    """Compute the Fast Fourier Transform (FFT) of a real sequence y. Input: [0, 1, 0, -1]."""
    if not y:
        raise ValueError("Input list cannot be empty for FFT.")
    return fft.fft(np.array(y))

@math_tool_handler
def ifft_transform(y_complex: List[complex]) -> List[complex]:
    """Compute the inverse Fast Fourier Transform (IFFT) of a complex sequence. Input: result from fft_transform."""
    if not y_complex:
        raise ValueError("Input list cannot be empty for IFFT.")
    return fft.ifft(np.array(y_complex))

# --- Tool List Creation ---

def get_python_math_tools() -> List[FunctionTool]:
    """Returns a list of FunctionTools for the Python math functions."""
    py_tools = [
        # Symbolic
        FunctionTool.from_defaults(fn=solve_symbolic_equation),
        FunctionTool.from_defaults(fn=compute_derivative),
        FunctionTool.from_defaults(fn=compute_integral),
        FunctionTool.from_defaults(fn=compute_limit),
        FunctionTool.from_defaults(fn=simplify_expression),
        FunctionTool.from_defaults(fn=expand_expression),
        FunctionTool.from_defaults(fn=factor_expression),
        # Matrix
        FunctionTool.from_defaults(fn=matrix_addition),
        FunctionTool.from_defaults(fn=matrix_subtraction),
        FunctionTool.from_defaults(fn=matrix_multiplication),
        FunctionTool.from_defaults(fn=matrix_inverse),
        FunctionTool.from_defaults(fn=matrix_determinant),
        FunctionTool.from_defaults(fn=matrix_transpose),
        FunctionTool.from_defaults(fn=matrix_rank),
        FunctionTool.from_defaults(fn=matrix_trace),
        FunctionTool.from_defaults(fn=matrix_norm),
        FunctionTool.from_defaults(fn=eigenvalues),
        FunctionTool.from_defaults(fn=eigenvectors),
        FunctionTool.from_defaults(fn=svd_decompose),
        FunctionTool.from_defaults(fn=lu_decompose),
        FunctionTool.from_defaults(fn=qr_decompose),
        # Statistics
        FunctionTool.from_defaults(fn=mean),
        FunctionTool.from_defaults(fn=median),
        FunctionTool.from_defaults(fn=std_dev),
        FunctionTool.from_defaults(fn=variance),
        FunctionTool.from_defaults(fn=percentile),
        FunctionTool.from_defaults(fn=covariance),
        FunctionTool.from_defaults(fn=correlation),
        FunctionTool.from_defaults(fn=linear_regression),
        # Numerical
        FunctionTool.from_defaults(fn=find_polynomial_roots),
        FunctionTool.from_defaults(fn=interpolate_value),
        FunctionTool.from_defaults(fn=numerical_integration),
        FunctionTool.from_defaults(fn=solve_ode),
        # Vector
        FunctionTool.from_defaults(fn=dot_product),
        FunctionTool.from_defaults(fn=cross_product),
        FunctionTool.from_defaults(fn=vector_magnitude),
        FunctionTool.from_defaults(fn=vector_normalize),
        FunctionTool.from_defaults(fn=vector_angle),
        # Probability
        FunctionTool.from_defaults(fn=binomial_pmf),
        FunctionTool.from_defaults(fn=normal_pdf),
        FunctionTool.from_defaults(fn=normal_cdf),
        FunctionTool.from_defaults(fn=poisson_pmf),
        # Special Functions
        FunctionTool.from_defaults(fn=gamma_function),
        FunctionTool.from_defaults(fn=beta_function),
        FunctionTool.from_defaults(fn=erf_function),
        # Fourier
        FunctionTool.from_defaults(fn=fft_transform),
        FunctionTool.from_defaults(fn=ifft_transform),
    ]
    # Update descriptions for clarity if needed (optional)
    for tool in py_tools:
        tool.metadata.description = f"(Python) {tool.metadata.description}"
    logger.info(f"Created {len(py_tools)} Python math tools.")
    return py_tools

# --- Wolfram Alpha Tool --- 
_wolfram_alpha_tools = None

def get_wolfram_alpha_tools() -> List[FunctionTool]:
    """Initializes and returns Wolfram Alpha tools (singleton)."""
    global _wolfram_alpha_tools
    if _wolfram_alpha_tools is None:
        logger.info("Initializing WolframAlphaToolSpec...")
        wolfram_alpha_app_id = os.getenv("WOLFRAM_ALPHA_APP_ID")
        if not wolfram_alpha_app_id:
            logger.warning("WOLFRAM_ALPHA_APP_ID not set. Wolfram Alpha tools will be unavailable.")
            _wolfram_alpha_tools = []
        else:
            try:
                spec = WolframAlphaToolSpec(app_id=wolfram_alpha_app_id)
                _wolfram_alpha_tools = spec.to_tool_list()
                # Add prefix to description for clarity
                for tool in _wolfram_alpha_tools:
                     tool.metadata.description = f"(WolframAlpha) {tool.metadata.description}"
                logger.info(f"WolframAlpha tools initialized: {len(_wolfram_alpha_tools)} tools.")
            except Exception as e:
                logger.error(f"Failed to initialize WolframAlpha tools: {e}", exc_info=True)
                _wolfram_alpha_tools = []
    return _wolfram_alpha_tools


# Use LlamaIndex's built-in Code Interpreter Tool Spec for safe execution
# This assumes the necessary environment (e.g., docker) for the spec is available
try:
    code_interpreter_spec = CodeInterpreterToolSpec()
    # Get the tool(s) from the spec. It might return multiple tools.
    code_interpreter_tools = code_interpreter_spec.to_tool_list()
    if not code_interpreter_tools:
        raise RuntimeError("CodeInterpreterToolSpec did not return any tools.")
    # Assuming the primary tool is the first one, or find by name if necessary
    code_interpreter_tool = next((t for t in code_interpreter_tools if t.metadata.name == "code_interpreter"), None)
    if code_interpreter_tool is None:
         raise RuntimeError("Could not find 'code_interpreter' tool in CodeInterpreterToolSpec results.")
    logger.info("CodeInterpreterToolSpec initialized successfully.")
except Exception as e:
    logger.error(f"Failed to initialize CodeInterpreterToolSpec: {e}", exc_info=True)
    # Fallback: Define a dummy tool or raise error to prevent agent start?
    # For now, let initialization fail if the safe interpreter isn't available.
    raise RuntimeError("CodeInterpreterToolSpec failed to initialize. Cannot create code_agent.") from e

# --- Agent Initialization ---

def initialize_math_agent() -> ReActAgent:
    """Initializes the Math Agent with Python and Wolfram Alpha tools."""
    logger.info("Initializing MathAgent...")

    # Configuration
    agent_llm_model = os.getenv("MATH_AGENT_LLM_MODEL", "gemini-2.5-pro-preview-03-25")
    gemini_api_key = os.getenv("GEMINI_API_KEY")

    if not gemini_api_key:
        logger.error("GEMINI_API_KEY not found in environment variables for MathAgent.")
        raise ValueError("GEMINI_API_KEY must be set for MathAgent")

    try:
        llm = GoogleGenAI(
            api_key=gemini_api_key,
            model=agent_llm_model,
            temperature=0.05
        )
        logger.info(f"Using agent LLM: {agent_llm_model}")

        # Combine Python tools and Wolfram Alpha tools
        all_tools = get_python_math_tools() + get_wolfram_alpha_tools() + [code_interpreter_tool]
        if not all_tools:
             logger.warning("No math tools available (Python or WolframAlpha). MathAgent may be ineffective.")

        # System prompt (consider loading from file)
        system_prompt = """\
        You are MathAgent, a powerful mathematical problem solver. Your goal is to accurately answer mathematical questions using the available tools.
        
        Available Tools:
        - Python Tools: A comprehensive suite for symbolic math (SymPy), numerical computation (NumPy/SciPy), statistics, linear algebra, calculus, ODEs, and transforms. Prefixed with '(Python)'. Use these for precise calculations when the method is clear.
        - WolframAlpha Tool: Accesses Wolfram Alpha for complex queries, natural language math questions, data, and real-world facts. Prefixed with '(WolframAlpha)'. Use this for broader questions, knowledge-based math, or when Python tools are insufficient.
        
        Workflow:
        1. **Thought**: Analyze the question. Determine the mathematical concepts involved. Decide the best tool or sequence of tools to use. Prefer Python tools for specific, well-defined calculations. Use WolframAlpha for complex, ambiguous, or knowledge-based queries.
        2. **Action**: Call the chosen tool with the correct arguments. Ensure inputs match the tool's requirements (e.g., list of lists for matrices, strings for symbolic expressions).
        3. **Observation**: Examine the tool's output. Check for errors or unexpected results.
        4. **Iteration**: If the result is incorrect or incomplete, rethink the approach. Try a different tool, adjust parameters, or break the problem down further. If a Python tool fails, consider rephrasing for WolframAlpha.
        5. **Final Answer**: Once the correct answer is obtained, state it clearly and concisely. Provide the numerical result, symbolic expression, or explanation as requested.
        6. **Hand-Off**: Pass the final mathematical result or analysis to **planner_agent** for integration into the overall response.
        
        Constraints:
        - Always use a tool for calculations; do not perform calculations yourself.
        - Clearly state which tool you are using and why.
        - Handle potential errors gracefully and report them if they prevent finding a solution.
        - Pay close attention to input formats required by each tool (e.g., lists for vectors/matrices, strings for symbolic expressions).
        
        If your response exceeds the maximum token limit and cannot be completed in a single reply, please conclude your output with the marker [CONTINUE]. In subsequent interactions, I will prompt you with “continue” to receive the next portion of the response.
        """

        agent = ReActAgent(
            name="math_agent",
            description=(
                "MathAgent solves mathematical problems using a suite of Python tools (SymPy, NumPy, SciPy) and WolframAlpha. "
                "It handles symbolic math, numerical computation, statistics, linear algebra, calculus, and more."
            ),
            tools=all_tools,
            llm=llm,
            system_prompt=system_prompt,
            can_handoff_to=["planner_agent", "reasoning_agent"],
        )
        logger.info("MathAgent initialized successfully.")
        return agent

    except Exception as e:
        logger.error(f"Error during MathAgent initialization: {e}", exc_info=True)
        raise

# Example usage (for testing if run directly)
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logger.info("Running math_agent.py directly for testing...")

    # Ensure API keys are set for testing
    required_keys = ["GEMINI_API_KEY"] # WOLFRAM_ALPHA_APP_ID is optional
    missing_keys = [key for key in required_keys if not os.getenv(key)]
    if missing_keys:
        print(f"Error: Required environment variable(s) not set: {', '.join(missing_keys)}. Cannot run test.")
    else:
        if not os.getenv("WOLFRAM_ALPHA_APP_ID"):
            print("Warning: WOLFRAM_ALPHA_APP_ID not set. WolframAlpha tools will be unavailable for testing.")
        try:
            test_agent = initialize_math_agent()
            print("Math Agent initialized successfully for testing.")
            # Example test
            # result = test_agent.chat("What is the integral of x**2 from 0 to 1?")
            # print(f"Test query result: {result}")
            # result2 = test_agent.chat("what is the population of france?") # Test WolframAlpha
            # print(f"Test query 2 result: {result2}")
        except Exception as e:
            print(f"Error during testing: {e}")