utils/math_solver.py

import re
import numpy as np
import sympy as sp
from sympy import symbols, Eq, solve, simplify, latex
import math

def preprocess_latex(latex_str):
    """Preprocess LaTeX string to make it compatible with SymPy"""
    # Remove common LaTeX formatting
    latex_str = latex_str.replace('\\left', '').replace('\\right', '')
    latex_str = latex_str.replace('\\ ', '')
    
    # Handle fractions
    latex_str = re.sub(r'\\frac\{([^}]+)\}\{([^}]+)\}', r'(\1)/(\2)', latex_str)
    
    # Handle square roots
    latex_str = re.sub(r'\\sqrt\{([^}]+)\}', r'sqrt(\1)', latex_str)
    
    # Handle exponents
    latex_str = re.sub(r'\^(\{[^}]+\}|\d+|[a-zA-Z])', r'**\1', latex_str)
    latex_str = re.sub(r'\*\*(\{([^}]+)\})', r'**(\2)', latex_str)
    
    # Handle subscripts (for simple cases)
    latex_str = re.sub(r'_\{([^}]+)\}', r'_\1', latex_str)
    
    # Handle Greek letters
    greek_letters = {
        '\\alpha': 'alpha', '\\beta': 'beta', '\\gamma': 'gamma', '\\delta': 'delta',
        '\\epsilon': 'epsilon', '\\zeta': 'zeta', '\\eta': 'eta', '\\theta': 'theta',
        '\\iota': 'iota', '\\kappa': 'kappa', '\\lambda': 'lambda', '\\mu': 'mu',
        '\\nu': 'nu', '\\xi': 'xi', '\\pi': 'pi', '\\rho': 'rho', '\\sigma': 'sigma',
        '\\tau': 'tau', '\\upsilon': 'upsilon', '\\phi': 'phi', '\\chi': 'chi',
        '\\psi': 'psi', '\\omega': 'omega'
    }
    
    for latex_letter, symbol in greek_letters.items():
        latex_str = latex_str.replace(latex_letter, symbol)
    
    return latex_str


def solve_equation(latex_equation):
    """Solve a LaTeX equation using SymPy"""
    try:
        # Preprocess the LaTeX equation
        processed_eq = preprocess_latex(latex_equation)
        
        # Handle equation solving (contains =)
        if '=' in processed_eq:
            # Split equation at equals sign
            left, right = processed_eq.split('=', 1)
            
            # Define symbols (assume x is the variable if not specified)
            x = symbols('x')
            
            # Try to parse and solve the equation
            try:
                left_expr = sp.sympify(left)
                right_expr = sp.sympify(right)
                equation = Eq(left_expr, right_expr)
                solutions = solve(equation, x)
                
                # Format solutions
                solution_str = []
                for sol in solutions:
                    solution_str.append(str(sol))
                
                return {
                    'type': 'equation',
                    'solutions': solution_str,
                    'latex_solutions': [latex(sol) for sol in solutions]
                }
            except Exception as e:
                # If equation solving fails, try simplification
                try:
                    left_expr = sp.sympify(left)
                    right_expr = sp.sympify(right)
                    diff = simplify(left_expr - right_expr)
                    return {
                        'type': 'simplification',
                        'result': str(diff),
                        'latex_result': latex(diff)
                    }
                except:
                    return {
                        'type': 'error',
                        'message': f'Could not solve equation: {str(e)}'
                    }
        
        # Handle expression evaluation
        else:
            try:
                expr = sp.sympify(processed_eq)
                result = simplify(expr)
                return {
                    'type': 'expression',
                    'result': str(result),
                    'latex_result': latex(result)
                }
            except Exception as e:
                return {
                    'type': 'error',
                    'message': f'Could not evaluate expression: {str(e)}'
                }
                
    except Exception as e:
        return {
            'type': 'error',
            'message': f'Error processing equation: {str(e)}'
        }


def solve_system(equations):
    """Solve a system of equations"""
    try:
        # Define symbols
        x, y, z = symbols('x y z')
        symbols_list = [x, y, z]
        
        # Parse equations
        parsed_equations = []
        for eq_str in equations:
            processed_eq = preprocess_latex(eq_str)
            if '=' in processed_eq:
                left, right = processed_eq.split('=', 1)
                left_expr = sp.sympify(left)
                right_expr = sp.sympify(right)
                parsed_equations.append(Eq(left_expr, right_expr))
        
        # Solve system
        solutions = solve(parsed_equations, symbols_list[:len(parsed_equations)])
        
        return {
            'type': 'system',
            'solutions': solutions,
            'latex_solutions': latex(solutions)
        }
    except Exception as e:
        return {
            'type': 'error',
            'message': f'Could not solve system: {str(e)}'
        }


def calculate_numerical(expression):
    """Calculate numerical value of an expression"""
    try:
        processed_expr = preprocess_latex(expression)
        expr = sp.sympify(processed_expr)
        numerical_value = float(expr.evalf())
        return {
            'type': 'numerical',
            'result': numerical_value,
            'latex_result': latex(expr)
        }
    except Exception as e:
        return {
            'type': 'error',
            'message': f'Could not calculate numerical value: {str(e)}'
        }