import sympy as sp
import gradio as gr
import os
import re
import numpy as np
from PIL import Image
import io
import tempfile
import subprocess
from sympy import symbols, diff, integrate, limit, sin, cos, tan, log, sqrt, factorial, Matrix, oo, E, I, pi

print("🚀 Math Solver starting...")

# Install system dependencies
def install_dependencies():
    try:
        print("Installing system dependencies...")
        subprocess.run(["apt-get", "update", "-y"], capture_output=True)
        subprocess.run(["apt-get", "install", "-y", "tesseract-ocr", "libtesseract-dev", "espeak", "espeak-ng"], capture_output=True)
        print("✅ System dependencies installed")
        return True
    except Exception as e:
        print(f"⚠️ Dependency warning: {e}")
        return False

install_dependencies()

# Try to import optional dependencies with fallbacks
try:
    import speech_recognition as sr
    SPEECH_RECOGNITION_AVAILABLE = True
except ImportError:
    SPEECH_RECOGNITION_AVAILABLE = False
    print("Speech recognition not available. Install with: pip install SpeechRecognition")

try:
    from gtts import gTTS
    GTTS_AVAILABLE = True
except ImportError:
    GTTS_AVAILABLE = False
    print("gTTS not available. Install with: pip install gTTS")

try:
    import pyttsx3
    PYTTSX3_AVAILABLE = False # Initialize as False
    try:
        engine = pyttsx3.init()
        engine.setProperty('rate', 150)
        engine.setProperty('volume', 0.9)
        PYTTSX3_AVAILABLE = True # Set to True if initialization succeeds
    except Exception as e:
        print(f"pyttsx3 initialization failed: {e}")
        engine = None
except ImportError:
    PYTTSX3_AVAILABLE = False
    engine = None
    print("pyttsx3 not available. Install with: pip install pyttsx3")


try:
    import pytesseract
    TESSERACT_AVAILABLE = True
    # Point tesseract_cmd to the correct executable if needed
    # pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract' # Uncomment and modify if tesseract is not in PATH
except ImportError:
    TESSERACT_AVAILABLE = False
    print("Tesseract not available. Install with: pip install pytesseract && sudo apt install tesseract-ocr")

try:
    from transformers import pipeline
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False
    print("Transformers not available. Install with: pip install transformers")

class MathSolver:
    def __init__(self):
        self.ai_models_loaded = False
        self.load_ai_models()

    def load_ai_models(self):
        """Load AI models with Hugging Face compatibility"""
        if TRANSFORMERS_AVAILABLE:
            try:
                # Using a simpler model for faster loading in Colab
                self.math_solver = pipeline(
                    "text2text-generation",
                    model="google/flan-t5-small",
                    tokenizer="google/flan-t5-small"
                )
                self.ai_models_loaded = True
                print("✅ AI models loaded successfully")
            except Exception as e:
                print(f"❌ AI model loading failed: {e}")
                self.ai_models_loaded = False
        else:
            print("❌ Transformers not available for AI models")

    def solve_with_ai(self, problem):
        """Solve math problems using AI"""
        if not self.ai_models_loaded:
            return None

        try:
            prompt = f"Solve this math problem: {problem}. Provide the final answer."
            result = self.math_solver(
                prompt,
                max_length=100,
                num_return_sequences=1,
                temperature=0.1
            )
            # Clean up potential conversational text from AI model
            generated_text = result[0]['generated_text']
            # Simple regex to try and isolate the math part if AI adds conversational text
            math_part = re.search(r'([-+]?\d*\.?\d+([eE][-+]?\d+)?|\S+)', generated_text)
            return math_part.group(0) if math_part else generated_text.strip()
        except Exception as e:
            print(f"AI solving error: {e}")
            return None

# Initialize math solver
math_solver = MathSolver()

def generate_tts(text, engine_choice="auto"):
    """Generate TTS audio - Hugging Face compatible"""
    temp_path = None
    try:
        # Create temp file for audio
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        temp_path = temp_file.name
        temp_file.close()

        # Clean text for TTS
        clean_text = re.sub(r'[**`]', '', text)
        # Replace common symbols with words for better pronunciation
        clean_text = clean_text.replace('+', ' plus ').replace('-', ' minus ').replace('*', ' times ').replace('/', ' divided by ').replace('**', ' to the power of ')
        clean_text = clean_text.replace('\n', '. ')[:300]  # Limit length and replace newlines

        success = False

        # Try pyttsx3 first if available and preferred
        if engine_choice in ["auto", "pyttsx3"] and PYTTSX3_AVAILABLE and engine:
            try:
                engine.save_to_file(clean_text, temp_path)
                engine.runAndWait()
                success = True
                # print("Generated audio using pyttsx3") # Debug print
            except Exception as e:
                print(f"pyttsx3 failed: {e}")
                success = False # Ensure success is False on failure

        # Fallback to gTTS if pyttsx3 failed or gTTS is preferred
        if not success and (engine_choice in ["auto", "gTTS"] or not PYTTSX3_AVAILABLE) and GTTS_AVAILABLE:
            try:
                tts = gTTS(text=clean_text, lang='en', slow=False)
                tts.save(temp_path)
                success = True
                # print("Generated audio using gTTS") # Debug print
            except Exception as e:
                print(f"gTTS failed: {e}")
                success = False

        if success:
            return temp_path
        else:
            print("Neither pyttsx3 nor gTTS could generate audio.")
            return None

    except Exception as e:
        print(f"TTS generation error: {e}")
        return None
    finally:
        # Clean up temp file if generation failed or was not attempted
        if temp_path and not os.path.exists(temp_path):
             try:
                 os.unlink(temp_path)
             except OSError as e:
                 print(f"Error removing temp file {temp_path}: {e}")


def extract_math_from_image(image_path):
    """Extract math from image using OCR"""
    if not TESSERACT_AVAILABLE:
        return "OCR not available. Please install pytesseract and tesseract-ocr.", ""

    if image_path is None:
        return "No image provided.", ""

    try:
        # Ensure image_path is a string path
        if isinstance(image_path, np.ndarray):
             # Save numpy array to a temp file
             pil_image = Image.fromarray(image_path.astype('uint8')).convert("RGB")
             temp_img_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
             image_path = temp_img_file.name
             pil_image.save(image_path)
             temp_img_file.close()
        elif isinstance(image_path, Image.Image):
             # Save PIL Image to a temp file
             temp_img_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
             image_path = temp_img_file.name
             image_path.convert("RGB").save(image_path)
             temp_img_file.close()
        elif not isinstance(image_path, str):
             return "Invalid image input type.", ""

        # Configure for math symbols (might need tuning)
        # Using --psm 6 for single uniform block of text, --oem 3 for default OCR engine
        custom_config = r'--oem 3 --psm 6'
        text = pytesseract.image_to_string(image_path, config=custom_config)

        # Clean up temp image file if created
        if isinstance(image_path, str) and (image_path.endswith(".png") or image_path.endswith(".jpg")): # Basic check if it's a temp file
             try:
                 os.unlink(image_path)
             except OSError as e:
                 print(f"Error removing temp image file {image_path}: {e}")


        if text.strip():
            # Clean OCR text
            cleaned = clean_ocr_text(text)
            return f"📷 Extracted: {cleaned}", cleaned
        else:
            return "❌ No text found in image", ""

    except Exception as e:
        return f"❌ Image processing error: {str(e)}", ""

def clean_ocr_text(text):
    """Clean OCR-extracted text"""
    corrections = {
        '—': '-', '–': '-', '×': '*', '÷': '/',
        '**': '^', '``': '"', "''": '"',
        'O': '0', 'o': '0', 'l': '1', 'I': '1',
        '=': '==' # For equality checks
    }

    cleaned = text
    for wrong, correct in corrections.items():
        cleaned = cleaned.replace(wrong, correct)

    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    return cleaned

def voice_to_text(audio_path):
    """Convert voice to text"""
    if not SPEECH_RECOGNITION_AVAILABLE:
        return "Speech recognition not available. Please type your problem."

    if audio_path is None:
        return "No audio provided."

    recognizer = sr.Recognizer()
    try:
        with sr.AudioFile(audio_path) as source:
            audio_data = recognizer.record(source)
            text = recognizer.recognize_google(audio_data)
            return text
    except sr.UnknownValueError:
        return "Could not understand audio"
    except sr.RequestError:
        return "Speech service unavailable"
    except Exception as e:
        return f"Audio error: {str(e)}"


def convert_speech_to_math(text):
    """Convert natural language to math expressions - SIMPLE & RELIABLE"""
    if not text or text.strip() == "":
        return "0"
    
    text = text.lower().strip()
    print(f"Original voice input: '{text}'")  # Debug
    
    # Remove common question phrases
    question_phrases = ["what is", "calculate", "compute", "solve", "what's", "how much is"]
    for phrase in question_phrases:
        text = text.replace(phrase, "").strip()
    
    # Handle simple arithmetic directly
    if any(op in text for op in ['+', '-', '*', '/', 'x', '×']):
        # Replace word operators with symbols
        text = text.replace('x', '*').replace('×', '*')
        return text
    
    # Handle spoken arithmetic patterns
    if "plus" in text:
        text = text.replace("plus", "+")
    if "minus" in text:
        text = text.replace("minus", "-")
    if "times" in text or "multiplied by" in text:
        text = text.replace("times", "*").replace("multiplied by", "*")
    if "divided by" in text:
        text = text.replace("divided by", "/")
    
    # Handle number words
    number_words = {
        'zero': '0', 'one': '1', 'two': '2', 'three': '3', 'four': '4',
        'five': '5', 'six': '6', 'seven': '7', 'eight': '8', 'nine': '9',
        'ten': '10', 'eleven': '11', 'twelve': '12', 'thirteen': '13', 
        'fourteen': '14', 'fifteen': '15', 'sixteen': '16', 'seventeen': '17',
        'eighteen': '18', 'nineteen': '19', 'twenty': '20', 'thirty': '30',
        'forty': '40', 'fifty': '50', 'sixty': '60', 'seventy': '70',
        'eighty': '80', 'ninety': '90'
    }
    
    for word, num in number_words.items():
        text = text.replace(word, num)
    
    # Clean up spaces around operators
    import re
    text = re.sub(r'\s*([+\-*/])\s*', r'\1', text)
    
    print(f"Converted to math: '{text}'")  # Debug
    return text

def evaluate_math_safely(expression):
    """Safely evaluate mathematical expressions"""
    try:
        # Handle simple arithmetic directly first
        if any(op in expression for op in ['+', '-', '*', '/']):
            try:
                # Use sympy for evaluation
                result = sp.sympify(expression)
                return result
            except:
                pass
        
        # Try direct evaluation
        result = sp.sympify(expression)
        return result
        
    except Exception as e:
        print(f"Evaluation error for '{expression}': {e}")
        return None

def process_math_complete(query):
    """Complete math processing with proper error handling"""
    try:
        print(f"Processing query: '{query}'")
        
        # Convert natural language to math
        math_expr = convert_speech_to_math(query)
        print(f"Converted expression: '{math_expr}'")
        
        # Evaluate the math
        result = evaluate_math_safely(math_expr)
        print(f"Raw result: {result}, Type: {type(result)}")
        
        if result is None:
            return "❌ Could not evaluate the expression. Please try rephrasing.", None
        
        # Handle boolean results (like comparisons)
        if isinstance(result, bool):
            if "=" in math_expr or "==" in math_expr or "!=" in math_expr:
                return f"**Expression**: `{math_expr}`\n**Result**: `{result}`", None
            else:
                return "❌ Unexpected boolean result. Please check your expression.", None
        
        # Format successful result
        if hasattr(result, 'evalf'):
            numerical = result.evalf()
            result_text = f"""
**Input**: {query}
**Expression**: `{math_expr}`
**Result**: `{result}`
**Numerical Value**: `{numerical}`
**Method**: Symbolic Math
"""
        else:
            result_text = f"""
**Input**: {query}
**Expression**: `{math_expr}`
**Result**: `{result}`
**Method**: Direct Evaluation
"""
        
        return result_text, None
        
    except Exception as e:
        error_msg = f"❌ Processing error: {str(e)}"
        print(error_msg)
        return error_msg, None

def main_handler(input_type, user_input, calculus_op="differentiate", variable="x", image=None):
    """Main handler for all math operations - UPDATED"""
    try:
        # Handle voice/text input
        if input_type in ["basic", "equation"] and user_input:
            return process_math_complete(user_input)
        
        # Handle calculus
        elif input_type == "calculus" and user_input:
            return process_calculus(calculus_op, user_input, variable), None
        
        # Handle image input
        elif input_type == "image" and image is not None:
            ocr_result, extracted_text = extract_text_from_image(image)
            if extracted_text:
                return process_math_complete(extracted_text)
            return ocr_result, None
        else:
            return "❌ Please provide valid input", None
            
    except Exception as e:
        return f"❌ Unexpected error: {str(e)}", None



def evaluate_advanced_math(expression):
    """Evaluate mathematical expressions using SymPy"""
    x, y, z = symbols('x y z') # Define symbols

    try:
        # Handle various operations
        expr_lower = expression.lower()

        if 'diff(' in expr_lower:
            match = re.search(r'diff\((.*?),\s*(\w+)\)', expression)
            if match:
                expr_str, var = match.groups()
                # Ensure variable is a symbol
                return diff(sp.sympify(expr_str), symbols(var))

        elif 'integrate(' in expr_lower or 'int(' in expr_lower:
            match = re.search(r'(?:integrate|int)\((.*?),\s*(\w+)\)', expression)
            if match:
                expr_str, var = match.groups()
                 # Ensure variable is a symbol
                return integrate(sp.sympify(expr_str), symbols(var))

        elif 'limit(' in expr_lower:
            match = re.search(r'limit\((.*?),\s*(\w+)\s*->\s*([^)]+)\)', expression)
            if match:
                expr_str, var, point = match.groups()
                 # Ensure variable is a symbol
                return limit(sp.sympify(expr_str), symbols(var), sp.sympify(point))

        elif 'factorial(' in expr_lower:
            match = re.search(r'factorial\((\d+)\)', expression)
            if match:
                return factorial(int(match.group(1)))

        # Default evaluation using sympify
        return sp.sympify(expression)

    except Exception as e:
        raise ValueError(f"Could not evaluate: {expression}. Error: {str(e)}")

def process_math(query, use_ai=True, auto_play=True, tts_engine_choice="auto"):
    """Process math query and return result"""
    try:
        # Convert natural language
        math_expr = convert_speech_to_math(query)

        result = None
        method_used = "Symbolic Math"

        # Try symbolic math first
        try:
            result = evaluate_advanced_math(math_expr)
        except ValueError:
             # If symbolic math failed, try AI if enabled
            if use_ai and math_solver.ai_models_loaded:
                ai_result = math_solver.solve_with_ai(query)
                if ai_result:
                    result = ai_result
                    method_used = "AI Model"
                else:
                     result = f"❌ Unable to solve '{query}' using AI. Trying basic evaluation."
                     method_used = "Fallback Evaluation"

            # Final fallback to basic evaluation if AI also failed or not used
            if result is None or "Unable to solve" in str(result):
                 try:
                     # Attempt a very basic evaluation, might fail on complex expressions
                     result = eval(math_expr)
                     method_used = "Basic Evaluation (eval)"
                 except:
                      result = f"❌ Unable to solve '{query}'. Try rephrasing or check syntax."
                      method_used = "Failed"


        # Format result
        if isinstance(result, sp.Basic): # Check if it's a SymPy object
            try:
                numerical = result.evalf()
                result_text = f"""**Input**: `{query}`
**Symbolic Result**: `{result}`
**Numerical Result**: `{numerical}`
**Method**: {method_used}"""
            except Exception as e:
                 # Handle cases where evalf might fail
                 result_text = f"""**Input**: `{query}`
**Symbolic Result**: `{result}`
**Numerical Result**: Could not evaluate numerically ({e})
**Method**: {method_used}"""
        else: # For results from AI or basic eval
            result_text = f"""**Input**: `{query}`
**Result**: `{result}`
**Method**: {method_used}"""

        # Generate audio
        audio_path = None
        if auto_play and "Unable to solve" not in result_text:
            speak_text = f"Result is {result}"
            audio_path = generate_tts(speak_text, engine_choice=tts_engine_choice)

        return result_text, audio_path

    except Exception as e:
        error_msg = f"❌ An unexpected error occurred: {str(e)}"
        audio_path = generate_tts("Sorry, an error occurred while processing that problem.", engine_choice=tts_engine_choice) if auto_play else None
        return error_msg, audio_path

def process_all_inputs(audio=None, text_input=None, image=None, use_ai=True, auto_play=True, tts_engine_choice="auto"):
    """Process all input types"""
    query = ""
    output_message = ""

    # Priority: Image > Audio > Text
    if image is not None:
        extraction_result, extracted_text = extract_math_from_image(image)
        output_message = extraction_result
        if extracted_text:
            query = extracted_text
        else:
            # If image processing failed or found no text, return the message and None for audio
            audio_path = generate_tts(output_message, engine_choice=tts_engine_choice) if auto_play and "No text found" not in output_message else None
            return output_message, audio_path

    if not query and audio is not None:
        voice_text = voice_to_text(audio)
        if any(msg in voice_text for msg in ["not available", "not understand", "unavailable", "error"]):
            return voice_text, None # Return error message and None for audio directly

        query = voice_text
        output_message = f"🎤 Transcribed: {query}"

    if not query and text_input:
        query = text_input
        output_message = f"📝 Input: {query}"


    if not query:
        msg = "Please provide input via voice, text, or image."
        audio_path = generate_tts(msg, engine_choice=tts_engine_choice) if auto_play else None
        return msg, audio_path

    # Process the math query
    result_text, audio_path = process_math(query, use_ai, auto_play, tts_engine_choice)

    # Combine initial message with the result
    final_output_text = f"{output_message}\n\n{result_text}" if output_message and "Extracted:" not in output_message else result_text

    # Return the output text and audio path
    # Ensure audio_path is None if no audio was generated to satisfy Gradio's expected output format
    return final_output_text, audio_path if audio_path and os.path.exists(audio_path) else None

# Create the interface
def create_interface():
    global PYTTSX3_AVAILABLE, GTTS_AVAILABLE, SPEECH_RECOGNITION_AVAILABLE, TESSERACT_AVAILABLE, math_solver

    with gr.Blocks(theme=gr.themes.Soft(), title="Math Solver Pro") as demo:
        gr.Markdown("""
        # 🧮 Math Solver Pro
        **Solve math problems using Voice, Text, or Images with Audio Responses**

        *Powered by SymPy • Hugging Face • Advanced Math Engine*
        """)

        with gr.Row():
            with gr.Column():
                # Input Methods
                gr.Markdown("### 📥 Input Methods")

                with gr.Tab("🎤 Voice"):
                    audio_input = gr.Audio(
                        sources=["microphone", "upload"],
                        type="filepath", # Changed to filepath
                        label="Speak Math Problem"
                    )

                with gr.Tab("📝 Text"):
                    text_input = gr.Textbox(
                        label="Type Math Problem",
                        placeholder="Examples: 2+2, derivative of x^2, integrate sin(x)",
                        lines=3
                    )

                with gr.Tab("📷 Image"):
                    image_input = gr.Image(
                        label="Upload Math Image",
                        type="filepath", # Changed to filepath
                        show_download_button=False
                    )

                # Settings
                with gr.Accordion("⚙️ Settings", open=False):
                    with gr.Row():
                        use_ai = gr.Checkbox(
                            value=math_solver.ai_models_loaded, # Reflect actual AI load status
                            label="Use AI Models",
                             interactive=math_solver.ai_models_loaded # Only interactive if loaded
                        )
                        auto_play = gr.Checkbox(
                            value=True,
                            label="Auto-Play Audio"
                        )
                    with gr.Row():
                         tts_engine_choice = gr.Radio(
                            ["auto", "pyttsx3", "gTTS", "None"],
                            label="TTS Engine",
                            value="auto",
                            info="auto: prefers pyttsx3 if available, then gTTS. None: no audio."
                         )

                # Action Buttons
                with gr.Row():
                    solve_btn = gr.Button("🧠 Solve", variant="primary")
                    clear_btn = gr.Button("🔄 Clear")

            with gr.Column():
                # Results
                gr.Markdown("### 📊 Results")
                output_text = gr.Markdown(
                    label="Solution",
                    value="Your solution will appear here..."
                )

                audio_output = gr.Audio(
                    label="🔊 Audio Result",
                    autoplay=True,
                    visible=True,
                    value=None # Initialize with None
                )

        # System Status
        with gr.Accordion("🤖 System Status", open=False):
            status_text = f"""
            **Available Features:**
            - ✅ Advanced Math Engine (SymPy)
            - {'✅' if SPEECH_RECOGNITION_AVAILABLE else '❌'} Voice Input (Requires `SpeechRecognition`)
            - {'✅' if TESSERACT_AVAILABLE else '❌'} Image OCR (Requires `pytesseract` and `tesseract-ocr`)
            - {'✅' if GTTS_AVAILABLE else '❌'} Online TTS (Requires `gTTS`)
            - {'✅' if PYTTSX3_AVAILABLE else '❌'} Offline TTS (Requires `pyttsx3`)
            - {'✅' if math_solver.ai_models_loaded else '❌'} AI Models (Requires `transformers`)
            """
            gr.Markdown(status_text)

        # Examples
        with gr.Accordion("📚 Examples", open=True):
            gr.Markdown("""
            **Try these examples:**
            - **Voice**: "What is 15 times 27?"
            - **Text**: `integrate x^2 + 3x + 1 from 0 to 1`
            - **Image**: Upload equation photo (e.g., `sqrt(16)`)
            - **Text**: `diff(sin(x) + cos(x), x)`
            - **Voice**: "Calculate factorial of 7"
            """)

        # Event handlers
        solve_btn.click(
            fn=process_all_inputs,
            inputs=[audio_input, text_input, image_input, use_ai, auto_play, tts_engine_choice],
            outputs=[output_text, audio_output]
        )

        def clear_all():
            # Return None for inputs and initial values for outputs to clear the interface
            # The temporary file will be managed by Gradio itself when the component value changes
            return None, "", None, "Your solution will appear here...", None

        clear_btn.click(
            fn=clear_all,
            inputs=[], # Clear button doesn't need inputs
            outputs=[audio_input, text_input, image_input, output_text, audio_output]
        )

        text_input.submit(
            fn=process_all_inputs,
            inputs=[gr.State(None), text_input, gr.State(None), use_ai, auto_play, tts_engine_choice],
            outputs=[output_text, audio_output]
        )

    return demo

# Hugging Face Spaces entry point
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)