"""
Mathematical computation tools for the GAIA Agent.
Includes symbolic math, matrix operations, calculator, and statistics.
"""

import json
from langchain_core.tools import tool


@tool
def symbolic_math(expression: str, operation: str = "simplify", variable: str = "x") -> str:
    """Perform symbolic mathematics operations using SymPy.
    
    Operations available:
    - simplify: Simplify an expression
    - expand: Expand an expression
    - factor: Factor an expression
    - solve: Solve an equation (set equal to 0)
    - differentiate: Compute derivative with respect to variable
    - integrate: Compute indefinite integral with respect to variable
    - limit: Compute limit as variable approaches 0 (use expression like "sin(x)/x")
    - series: Compute Taylor series expansion
    
    Args:
        expression: Mathematical expression (use ** for power, e.g., "x**2 + 2*x + 1")
        operation: One of: simplify, expand, factor, solve, differentiate, integrate, limit, series
        variable: Variable to use for calculus operations (default: "x")
    """
    try:
        import sympy as sp
        from sympy.parsing.sympy_parser import parse_expr, standard_transformations, implicit_multiplication_application
        
        # Define common symbols
        x, y, z, t, n, a, b, c = sp.symbols('x y z t n a b c')
        symbol_map = {'x': x, 'y': y, 'z': z, 't': t, 'n': n, 'a': a, 'b': b, 'c': c}
        
        # Get the main variable
        var = symbol_map.get(variable, sp.Symbol(variable))
        
        # Parse the expression with implicit multiplication support
        transformations = standard_transformations + (implicit_multiplication_application,)
        expr = parse_expr(expression, local_dict=symbol_map, transformations=transformations)
        
        operation = operation.lower().strip()
        
        if operation == "simplify":
            result = sp.simplify(expr)
        elif operation == "expand":
            result = sp.expand(expr)
        elif operation == "factor":
            result = sp.factor(expr)
        elif operation == "solve":
            result = sp.solve(expr, var)
        elif operation in ("differentiate", "diff", "derivative"):
            result = sp.diff(expr, var)
        elif operation in ("integrate", "integral"):
            result = sp.integrate(expr, var)
        elif operation == "limit":
            result = sp.limit(expr, var, 0)
        elif operation == "series":
            result = sp.series(expr, var, 0, 6)
        else:
            return f"Unknown operation: {operation}. Available: simplify, expand, factor, solve, differentiate, integrate, limit, series"
        
        return f"Input: {expression}\nOperation: {operation}\nResult: {result}"
    
    except ImportError:
        return "Error: SymPy is not installed. Please install it with: pip install sympy"
    except Exception as e:
        return f"Symbolic math error: {str(e)}"


@tool
def matrix_operations(matrix_a: str, operation: str = "determinant", matrix_b: str = "") -> str:
    """Perform matrix operations using NumPy.
    
    Operations available:
    - determinant: Compute determinant of matrix_a
    - inverse: Compute inverse of matrix_a
    - transpose: Compute transpose of matrix_a
    - eigenvalues: Compute eigenvalues of matrix_a
    - eigenvectors: Compute eigenvectors of matrix_a
    - rank: Compute rank of matrix_a
    - trace: Compute trace of matrix_a
    - multiply: Matrix multiplication of matrix_a @ matrix_b
    - add: Element-wise addition of matrix_a + matrix_b
    - solve: Solve linear system Ax = b (matrix_a is A, matrix_b is b vector)
    
    Args:
        matrix_a: Matrix as JSON array, e.g., "[[1,2],[3,4]]"
        operation: One of: determinant, inverse, transpose, eigenvalues, eigenvectors, rank, trace, multiply, add, solve
        matrix_b: Second matrix for binary operations (as JSON array)
    """
    try:
        import numpy as np
        
        # Parse matrix_a
        try:
            a = np.array(json.loads(matrix_a), dtype=float)
        except json.JSONDecodeError:
            return f"Error parsing matrix_a: {matrix_a}. Use JSON format like [[1,2],[3,4]]"
        
        operation = operation.lower().strip()
        
        # Parse matrix_b if needed
        b = None
        if matrix_b and operation in ("multiply", "add", "solve"):
            try:
                b = np.array(json.loads(matrix_b), dtype=float)
            except json.JSONDecodeError:
                return f"Error parsing matrix_b: {matrix_b}. Use JSON format like [[1,2],[3,4]]"
        
        if operation == "determinant":
            if a.shape[0] != a.shape[1]:
                return "Error: Determinant requires a square matrix"
            result = np.linalg.det(a)
        elif operation == "inverse":
            if a.shape[0] != a.shape[1]:
                return "Error: Inverse requires a square matrix"
            result = np.linalg.inv(a)
        elif operation == "transpose":
            result = a.T
        elif operation == "eigenvalues":
            if a.shape[0] != a.shape[1]:
                return "Error: Eigenvalues require a square matrix"
            result = np.linalg.eigvals(a)
        elif operation == "eigenvectors":
            if a.shape[0] != a.shape[1]:
                return "Error: Eigenvectors require a square matrix"
            eigenvalues, eigenvectors = np.linalg.eig(a)
            result = {"eigenvalues": eigenvalues.tolist(), "eigenvectors": eigenvectors.tolist()}
        elif operation == "rank":
            result = np.linalg.matrix_rank(a)
        elif operation == "trace":
            result = np.trace(a)
        elif operation == "multiply":
            if b is None:
                return "Error: multiply operation requires matrix_b"
            result = a @ b
        elif operation == "add":
            if b is None:
                return "Error: add operation requires matrix_b"
            result = a + b
        elif operation == "solve":
            if b is None:
                return "Error: solve operation requires matrix_b (the b vector)"
            result = np.linalg.solve(a, b)
        else:
            return f"Unknown operation: {operation}. Available: determinant, inverse, transpose, eigenvalues, eigenvectors, rank, trace, multiply, add, solve"
        
        # Format result
        if isinstance(result, np.ndarray):
            result_str = np.array2string(result, precision=6, suppress_small=True)
        elif isinstance(result, dict):
            result_str = json.dumps(result, indent=2)
        else:
            result_str = str(result)
        
        return f"Matrix A:\n{np.array2string(a)}\nOperation: {operation}\nResult:\n{result_str}"
    
    except ImportError:
        return "Error: NumPy is not installed. Please install it with: pip install numpy"
    except np.linalg.LinAlgError as e:
        return f"Linear algebra error: {str(e)}"
    except Exception as e:
        return f"Matrix operation error: {str(e)}"


@tool
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression with high precision.
    
    Supports standard math operations: +, -, *, /, ** (power), % (modulo)
    Also supports functions: sqrt, sin, cos, tan, log, log10, exp, abs, ceil, floor, round
    Constants: pi, e
    
    Args:
        expression: Mathematical expression to evaluate, e.g., "sqrt(2) * pi" or "2**10 + 5"
    """
    try:
        import math
        from decimal import Decimal, getcontext
        
        # Set high precision
        getcontext().prec = 50
        
        # Safe evaluation context with math functions
        safe_dict = {
            'sqrt': math.sqrt,
            'sin': math.sin,
            'cos': math.cos,
            'tan': math.tan,
            'asin': math.asin,
            'acos': math.acos,
            'atan': math.atan,
            'atan2': math.atan2,
            'log': math.log,
            'log10': math.log10,
            'log2': math.log2,
            'exp': math.exp,
            'pow': pow,
            'abs': abs,
            'ceil': math.ceil,
            'floor': math.floor,
            'round': round,
            'pi': math.pi,
            'e': math.e,
            'inf': math.inf,
            'factorial': math.factorial,
            'gcd': math.gcd,
            'lcm': math.lcm,
            'degrees': math.degrees,
            'radians': math.radians,
        }
        
        # Basic security check - only allow safe characters
        allowed_chars = set('0123456789+-*/.()%, abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ_')
        if not all(c in allowed_chars for c in expression):
            return f"Error: Expression contains invalid characters. Only math operations and functions are allowed."
        
        # Evaluate the expression
        result = eval(expression, {"__builtins__": {}}, safe_dict)
        
        # Format the result
        if isinstance(result, float):
            # Check if it's close to an integer
            if result == int(result):
                result_str = str(int(result))
            else:
                result_str = f"{result:.15g}"  # High precision but remove trailing zeros
        else:
            result_str = str(result)
        
        return f"Expression: {expression}\nResult: {result_str}"
    
    except ZeroDivisionError:
        return "Error: Division by zero"
    except ValueError as e:
        return f"Math error: {str(e)}"
    except Exception as e:
        return f"Calculator error: {str(e)}"


@tool
def statistical_analysis(data: str, operation: str = "describe") -> str:
    """Perform statistical analysis on numerical data.
    
    Operations available:
    - describe: Full statistical summary (mean, median, std, min, max, quartiles)
    - mean: Arithmetic mean
    - median: Median value
    - mode: Most frequent value
    - std: Standard deviation
    - var: Variance
    - correlation: Correlation coefficient (requires 2D data like [[x1,y1],[x2,y2],...])
    - regression: Linear regression (requires 2D data)
    - percentile: Compute 25th, 50th, 75th percentiles
    - zscore: Compute z-scores for each value
    
    Args:
        data: Numerical data as JSON array, e.g., "[1, 2, 3, 4, 5]" or "[[1,2],[3,4]]" for 2D
        operation: One of: describe, mean, median, mode, std, var, correlation, regression, percentile, zscore
    """
    try:
        import numpy as np
        from scipy import stats as sp_stats
        
        # Parse data
        try:
            arr = np.array(json.loads(data), dtype=float)
        except json.JSONDecodeError:
            return f"Error parsing data: {data}. Use JSON format like [1, 2, 3, 4, 5]"
        
        operation = operation.lower().strip()
        
        if operation == "describe":
            if arr.ndim == 1:
                result = {
                    "count": len(arr),
                    "mean": float(np.mean(arr)),
                    "median": float(np.median(arr)),
                    "std": float(np.std(arr)),
                    "variance": float(np.var(arr)),
                    "min": float(np.min(arr)),
                    "max": float(np.max(arr)),
                    "25th_percentile": float(np.percentile(arr, 25)),
                    "50th_percentile": float(np.percentile(arr, 50)),
                    "75th_percentile": float(np.percentile(arr, 75)),
                    "sum": float(np.sum(arr)),
                }
            else:
                result = {
                    "shape": arr.shape,
                    "mean_per_column": np.mean(arr, axis=0).tolist(),
                    "std_per_column": np.std(arr, axis=0).tolist(),
                }
        elif operation == "mean":
            result = float(np.mean(arr))
        elif operation == "median":
            result = float(np.median(arr))
        elif operation == "mode":
            mode_result = sp_stats.mode(arr.flatten(), keepdims=False)
            result = {"mode": float(mode_result.mode), "count": int(mode_result.count)}
        elif operation == "std":
            result = float(np.std(arr))
        elif operation == "var":
            result = float(np.var(arr))
        elif operation == "correlation":
            if arr.ndim != 2 or arr.shape[1] != 2:
                return "Error: correlation requires 2D data with 2 columns, e.g., [[x1,y1],[x2,y2],...]"
            result = float(np.corrcoef(arr[:, 0], arr[:, 1])[0, 1])
        elif operation == "regression":
            if arr.ndim != 2 or arr.shape[1] != 2:
                return "Error: regression requires 2D data with 2 columns, e.g., [[x1,y1],[x2,y2],...]"
            slope, intercept, r_value, p_value, std_err = sp_stats.linregress(arr[:, 0], arr[:, 1])
            result = {
                "slope": float(slope),
                "intercept": float(intercept),
                "r_squared": float(r_value**2),
                "p_value": float(p_value),
                "std_error": float(std_err),
                "equation": f"y = {slope:.6f}x + {intercept:.6f}"
            }
        elif operation == "percentile":
            result = {
                "25th": float(np.percentile(arr, 25)),
                "50th": float(np.percentile(arr, 50)),
                "75th": float(np.percentile(arr, 75)),
                "90th": float(np.percentile(arr, 90)),
                "95th": float(np.percentile(arr, 95)),
                "99th": float(np.percentile(arr, 99)),
            }
        elif operation == "zscore":
            zscores = sp_stats.zscore(arr.flatten())
            result = zscores.tolist()
        else:
            return f"Unknown operation: {operation}. Available: describe, mean, median, mode, std, var, correlation, regression, percentile, zscore"
        
        # Format result
        if isinstance(result, dict):
            result_str = json.dumps(result, indent=2)
        elif isinstance(result, list):
            result_str = json.dumps(result)
        else:
            result_str = str(result)
        
        return f"Data: {len(arr.flatten())} values\nOperation: {operation}\nResult:\n{result_str}"
    
    except ImportError as e:
        missing = "scipy" if "scipy" in str(e) else "numpy"
        return f"Error: {missing} is not installed. Please install it with: pip install {missing}"
    except Exception as e:
        return f"Statistical analysis error: {str(e)}"