tools/calculator.py

import sympy as sp
import numpy as np
import matplotlib.pyplot as plt
import io
import base64
from typing import Any, Dict, List, Optional, Union
import logging
import re
import math

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class CalculatorTool:
    def __init__(self):
        self.variables = {}
        self.last_result = None
    
    def evaluate_expression(self, expression: str) -> Dict[str, Any]:
        """
        Safely evaluate mathematical expressions
        """
        try:
            # Clean the expression
            expression = self._clean_expression(expression)
            
            # Try sympy first for symbolic computation
            try:
                result = sp.sympify(expression).evalf()
                self.last_result = float(result)
                return {
                    'result': float(result),
                    'expression': expression,
                    'type': 'symbolic',
                    'formatted': str(result)
                }
            except:
                # Fall back to basic evaluation
                result = eval(expression, {"__builtins__": {}}, self._get_safe_namespace())
                self.last_result = result
                return {
                    'result': result,
                    'expression': expression,
                    'type': 'numeric',
                    'formatted': str(result)
                }
                
        except Exception as e:
            logger.error(f"Error evaluating expression: {e}")
            return {
                'error': str(e),
                'expression': expression,
                'result': None
            }
    
    def _clean_expression(self, expression: str) -> str:
        """
        Clean and prepare expression for evaluation
        """
        # Replace common math notation
        replacements = {
            '^': '**',
            '×': '*',
            '÷': '/',
            'π': 'pi',
            'e': 'E'
        }
        
        for old, new in replacements.items():
            expression = expression.replace(old, new)
        
        return expression
    
    def _get_safe_namespace(self) -> Dict[str, Any]:
        """
        Get safe namespace for expression evaluation
        """
        safe_dict = {
            'abs': abs, 'round': round, 'min': min, 'max': max,
            'sum': sum, 'pow': pow, 'divmod': divmod,
            'sin': math.sin, 'cos': math.cos, 'tan': math.tan,
            'asin': math.asin, 'acos': math.acos, 'atan': math.atan,
            'sinh': math.sinh, 'cosh': math.cosh, 'tanh': math.tanh,
            'log': math.log, 'log10': math.log10, 'log2': math.log2,
            'exp': math.exp, 'sqrt': math.sqrt, 'factorial': math.factorial,
            'pi': math.pi, 'e': math.e, 'inf': math.inf, 'nan': math.nan,
            'degrees': math.degrees, 'radians': math.radians,
            'ceil': math.ceil, 'floor': math.floor,
        }
        safe_dict.update(self.variables)
        return safe_dict
    
    def solve_equation(self, equation: str, variable: str = 'x') -> Dict[str, Any]:
        """
        Solve equations symbolically
        """
        try:
            # Parse equation
            if '=' in equation:
                left, right = equation.split('=', 1)
                eq = sp.Eq(sp.sympify(left), sp.sympify(right))
            else:
                eq = sp.sympify(equation)
            
            # Solve
            var = sp.Symbol(variable)
            solutions = sp.solve(eq, var)
            
            return {
                'equation': equation,
                'variable': variable,
                'solutions': [str(sol) for sol in solutions],
                'numeric_solutions': [float(sol.evalf()) if sol.is_real else complex(sol.evalf()) for sol in solutions]
            }
            
        except Exception as e:
            logger.error(f"Error solving equation: {e}")
            return {
                'error': str(e),
                'equation': equation,
                'solutions': []
            }
    
    def plot_function(self, expression: str, x_range: tuple = (-10, 10), 
                     points: int = 1000) -> str:
        """
        Plot a mathematical function and return base64 encoded image
        """
        try:
            x = sp.Symbol('x')
            expr = sp.sympify(expression)
            
            # Convert to numpy function
            f = sp.lambdify(x, expr, 'numpy')
            
            # Generate points
            x_vals = np.linspace(x_range[0], x_range[1], points)
            y_vals = f(x_vals)
            
            # Create plot
            plt.figure(figsize=(10, 6))
            plt.plot(x_vals, y_vals, 'b-', linewidth=2)
            plt.grid(True, alpha=0.3)
            plt.xlabel('x')
            plt.ylabel('f(x)')
            plt.title(f'Plot of f(x) = {expression}')
            
            # Convert to base64
            buffer = io.BytesIO()
            plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
            buffer.seek(0)
            plot_data = base64.b64encode(buffer.getvalue()).decode()
            plt.close()
            
            return plot_data
            
        except Exception as e:
            logger.error(f"Error plotting function: {e}")
            return ""
    
    def calculate_derivative(self, expression: str, variable: str = 'x', 
                           order: int = 1) -> Dict[str, Any]:
        """
        Calculate derivative of an expression
        """
        try:
            var = sp.Symbol(variable)
            expr = sp.sympify(expression)
            
            derivative = sp.diff(expr, var, order)
            
            return {
                'original': expression,
                'derivative': str(derivative),
                'order': order,
                'variable': variable,
                'simplified': str(sp.simplify(derivative))
            }
            
        except Exception as e:
            logger.error(f"Error calculating derivative: {e}")
            return {
                'error': str(e),
                'original': expression
            }
    
    def calculate_integral(self, expression: str, variable: str = 'x', 
                          limits: Optional[tuple] = None) -> Dict[str, Any]:
        """
        Calculate integral of an expression
        """
        try:
            var = sp.Symbol(variable)
            expr = sp.sympify(expression)
            
            if limits:
                # Definite integral
                result = sp.integrate(expr, (var, limits[0], limits[1]))
                integral_type = 'definite'
            else:
                # Indefinite integral
                result = sp.integrate(expr, var)
                integral_type = 'indefinite'
            
            return {
                'original': expression,
                'integral': str(result),
                'type': integral_type,
                'variable': variable,
                'limits': limits,
                'numeric_value': float(result.evalf()) if result.is_number else None
            }
            
        except Exception as e:
            logger.error(f"Error calculating integral: {e}")
            return {
                'error': str(e),
                'original': expression
            }
    
    def matrix_operations(self, operation: str, *matrices) -> Dict[str, Any]:
        """
        Perform matrix operations
        """
        try:
            # Convert input to sympy matrices
            sp_matrices = []
            for matrix in matrices:
                if isinstance(matrix, list):
                    sp_matrices.append(sp.Matrix(matrix))
                else:
                    sp_matrices.append(sp.sympify(matrix))
            
            result = None
            
            if operation == 'add' and len(sp_matrices) >= 2:
                result = sp_matrices[0] + sp_matrices[1]
            elif operation == 'multiply' and len(sp_matrices) >= 2:
                result = sp_matrices[0] * sp_matrices[1]
            elif operation == 'inverse' and len(sp_matrices) >= 1:
                result = sp_matrices[0].inv()
            elif operation == 'determinant' and len(sp_matrices) >= 1:
                result = sp_matrices[0].det()
            elif operation == 'transpose' and len(sp_matrices) >= 1:
                result = sp_matrices[0].T
            elif operation == 'eigenvalues' and len(sp_matrices) >= 1:
                result = sp_matrices[0].eigenvals()
            
            return {
                'operation': operation,
                'result': str(result) if result is not None else None,
                'matrices_count': len(sp_matrices)
            }
            
        except Exception as e:
            logger.error(f"Error in matrix operation: {e}")
            return {
                'error': str(e),
                'operation': operation
            }
    
    def statistics_calculations(self, data: List[float], operation: str) -> Dict[str, Any]:
        """
        Perform statistical calculations
        """
        try:
            data = np.array(data)
            result = None
            
            if operation == 'mean':
                result = np.mean(data)
            elif operation == 'median':
                result = np.median(data)
            elif operation == 'std':
                result = np.std(data)
            elif operation == 'var':
                result = np.var(data)
            elif operation == 'min':
                result = np.min(data)
            elif operation == 'max':
                result = np.max(data)
            elif operation == 'sum':
                result = np.sum(data)
            elif operation == 'range':
                result = np.max(data) - np.min(data)
            
            return {
                'operation': operation,
                'result': float(result) if result is not None else None,
                'data_size': len(data),
                'data_preview': data[:5].tolist() if len(data) > 5 else data.tolist()
            }
            
        except Exception as e:
            logger.error(f"Error in statistics calculation: {e}")
            return {
                'error': str(e),
                'operation': operation
            }
    
    def unit_conversion(self, value: float, from_unit: str, to_unit: str) -> Dict[str, Any]:
        """
        Convert between different units
        """
        # Basic unit conversion factors (could be expanded)
        conversions = {
            # Length
            ('m', 'cm'): 100,
            ('m', 'mm'): 1000,
            ('m', 'km'): 0.001,
            ('cm', 'm'): 0.01,
            ('mm', 'm'): 0.001,
            ('km', 'm'): 1000,
            ('ft', 'm'): 0.3048,
            ('in', 'cm'): 2.54,
            
            # Weight
            ('kg', 'g'): 1000,
            ('g', 'kg'): 0.001,
            ('lb', 'kg'): 0.453592,
            ('kg', 'lb'): 2.20462,
            
            # Temperature (special handling needed)
            # Time
            ('h', 'min'): 60,
            ('min', 's'): 60,
            ('h', 's'): 3600,
            ('day', 'h'): 24,
        }
        
        try:
            if (from_unit, to_unit) in conversions:
                result = value * conversions[(from_unit, to_unit)]
            elif (to_unit, from_unit) in conversions:
                result = value / conversions[(to_unit, from_unit)]
            else:
                return {
                    'error': f"Conversion from {from_unit} to {to_unit} not supported",
                    'value': value
                }
            
            return {
                'original_value': value,
                'original_unit': from_unit,
                'converted_value': result,
                'converted_unit': to_unit,
                'conversion_factor': result / value if value != 0 else None
            }
            
        except Exception as e:
            logger.error(f"Error in unit conversion: {e}")
            return {
                'error': str(e),
                'value': value
            }
    
    def set_variable(self, name: str, value: Any) -> bool:
        """
        Set a variable for use in calculations
        """
        try:
            self.variables[name] = value
            logger.info(f"Set variable {name} = {value}")
            return True
        except Exception as e:
            logger.error(f"Error setting variable: {e}")
            return False
    
    def get_variables(self) -> Dict[str, Any]:
        """
        Get all stored variables
        """
        return self.variables.copy()
    
    def clear_variables(self) -> bool:
        """
        Clear all stored variables
        """
        try:
            self.variables.clear()
            logger.info("Cleared all variables")
            return True
        except Exception as e:
            logger.error(f"Error clearing variables: {e}")
            return False
    
    def format_result_for_llm(self, result: Dict[str, Any]) -> str:
        """
        Format calculation results for LLM consumption
        """
        if 'error' in result:
            return f"Error: {result['error']}"
        
        if 'result' in result:
            return f"Result: {result['result']}\nExpression: {result.get('expression', 'N/A')}"
        
        # Handle other result types
        formatted_parts = []
        for key, value in result.items():
            if key not in ['error'] and value is not None:
                formatted_parts.append(f"{key.title()}: {value}")
        
        return "\n".join(formatted_parts) if formatted_parts else "No result to display"