import gradio as gr
import sympy as sp
from pix2text import Pix2Text
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import re
import io
import logging
from llm_utils import explain_with_llm  # ✅ Added for LLM explanation

# Configure logging for debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define symbolic variables
x, y = sp.symbols('x y')

# ✅ Helper to get safe variable symbol
def get_variable_symbol(varname):
    if varname in {"pi", "e", "I", "i"}:
        return x
    try:
        return sp.Symbol(varname)
    except Exception:
        return x

# Initialize Pix2Text model globally
try:
    p2t_model = Pix2Text.from_config()
    logger.info("Pix2Text model loaded successfully")
except Exception as e:
    logger.error(f"Failed to load Pix2Text model: {e}")
    p2t_model = None

def clean_latex_expression(latex_str):
    """Clean and normalize LaTeX expression for SymPy parsing"""
    if not latex_str:
        return ""

    latex_str = latex_str.strip()
    latex_str = re.sub(r'^\$\$|\$\$$', '', latex_str)  # Remove $$ delimiters
    latex_str = re.sub(r'\\[a-zA-Z]+\{([^}]*)\}', r'\1', latex_str)  # Remove LaTeX commands
    latex_str = re.sub(r'\\{2,}', r'\\', latex_str)  # Fix multiple backslashes
    latex_str = re.sub(r'\s+', ' ', latex_str)  # Normalize whitespace
    latex_str = re.sub(r'\^{([^}]+)}', r'**\1', latex_str)  # Convert x^{n} to x**n
    latex_str = re.sub(r'(\d*\.?\d+)\s*([xy])', r'\1*\2', latex_str)  # Add multiplication: 1.0x -> 1.0*x
    latex_str = re.sub(r'\s*([+\-*/=])\s*', r'\1', latex_str)  # Remove spaces around operators

    if '=' in latex_str:
        left, right = latex_str.split('=')
        latex_str = f"{left} - ({right})"  # Move right-hand side to left

    # ✅ Insert missing multiplication
    latex_str = re.sub(r'(\))([a-zA-Z])', r'\1*\2', latex_str)
    latex_str = re.sub(r'(\d|\w)\(', r'\1*(', latex_str)

    # ✅ Replace LaTeX constants with sympy
    latex_str = latex_str.replace(r'\pi', 'pi')
    latex_str = latex_str.replace(r'\mathrm{e}', 'e')
    latex_str = latex_str.replace(r'\cdot', '*')
    latex_str = latex_str.replace(r'\times', '*')
    latex_str = latex_str.replace(r'\\', '')
    latex_str = re.sub(r'\\sqrt\{([^}]+)\}', r'sqrt(\1)', latex_str)

    return latex_str.strip()

def parse_equation_type(latex_str):
    """Determine if the equation is polynomial (single-variable) or linear system (two-variable)"""
    try:
        cleaned = clean_latex_expression(latex_str)
        if not cleaned:
            return 'polynomial'

        # Check for two-variable system
        if 'y' in cleaned and 'x' in cleaned:
            if '\\\\' in latex_str or '\n' in latex_str or len(re.split(r'\\\\|\n|;', latex_str)) >= 2:
                return 'linear_system'
            return 'linear'  # Single equation with x and y

        # Check for single-variable polynomial
        try:
            expr = sp.sympify(cleaned.split('-')[0] if '-' in cleaned else cleaned)
            if x in expr.free_symbols and y not in expr.free_symbols:
                degree = sp.degree(expr, x)
                return 'polynomial' if degree > 0 else 'linear'
            elif x not in expr.free_symbols and y in expr.free_symbols:
                return 'polynomial'
            else:
                return 'polynomial'
        except:
            if 'x**' in cleaned or '^' in latex_str:
                return 'polynomial'
            return 'polynomial'
    except Exception as e:
        logger.error(f"Error determining equation type: {e}")
        return 'polynomial'

def extract_polynomial_coefficients(latex_str):
    try:
        cleaned = clean_latex_expression(latex_str)
        expr = sp.sympify(cleaned, evaluate=False)
        if x not in expr.free_symbols and y not in expr.free_symbols:
            raise ValueError("No variable (x or y) found in expression")
        variable = next(iter(expr.free_symbols))
        degree = sp.degree(expr, variable)
        if degree < 1 or degree > 8:
            raise ValueError(f"Polynomial degree {degree} is out of supported range (1-8)")
        poly = sp.Poly(expr, variable)
        coeffs = [poly.coeff_monomial(variable**i).evalf() for i in range(degree, -1, -1)]
        return {
            "type": "polynomial",
            "degree": degree,
            "coeffs": " ".join(map(str, coeffs)),
            "latex": latex_str,
            "success": True,
            "variable": str(variable)
        }
    except Exception as e:
        logger.error(f"Error extracting polynomial coefficients: {e}")
        return {
            "type": "polynomial",
            "degree": 2,
            "coeffs": "1 0 0",
            "latex": latex_str,
            "success": False,
            "error": str(e),
            "variable": "x"
        }

def solve_polynomial(degree, coeff_string, real_only, variable_name="x"):
    try:
        variable = sp.Symbol(variable_name)
        coeffs = list(map(float, coeff_string.strip().split()))
        if len(coeffs) != degree + 1:
            return f"⚠️ Please enter exactly {degree + 1} coefficients.", None, None

        # Build the polynomial expression
        poly = sum([coeffs[i] * variable**(degree - i) for i in range(degree + 1)])
        simplified = sp.simplify(poly)
        factored = sp.factor(simplified)
        roots = sp.solve(sp.Eq(simplified, 0), variable)

        if real_only:
            roots = [r for r in roots if sp.im(r) == 0]

        # Format roots in LaTeX
        roots_output = "$$\n" + "\\ ".join(
            [f"r_{{{i}}} = {sp.latex(sp.nsimplify(r, rational=True))}" for i, r in enumerate(roots, 1)]
        ) + "\n$$"

        # Format steps in LaTeX
        steps_output = f"""
### Polynomial Expression
$$ {sp.latex(poly)} = 0 $$
### Simplified
$$ {sp.latex(simplified)} = 0 $$
### Factored
$$ {sp.latex(factored)} = 0 $$
### Roots {'(Only Real)' if real_only else '(All Roots)'}
{roots_output}
        """

        # Generate plot using numeric x-axis
        x_vals = np.linspace(-10, 10, 400)
        y_vals = np.polyval(coeffs, x_vals)

        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(x_vals, y_vals, label="Polynomial", color="blue")
        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.grid(True)
        ax.set_title("Graph of the Polynomial")
        ax.set_xlabel(str(variable))
        ax.set_ylabel("f(" + str(variable) + ")")
        ax.legend()

        return steps_output, fig, ""
    except Exception as e:
        return f"❌ Error: {e}", None, ""

def extract_linear_system_coefficients(latex_str):
    try:
        cleaned = clean_latex_expression(latex_str)
        equations = re.split(r'\\\\|\n|;', latex_str)
        if len(equations) < 2:
            equations = re.split(r'(?<=[0-9])\s*(?=[+-]?\s*[0-9]*[xy])', cleaned)
        if len(equations) < 2 or 'y' not in cleaned or 'x' not in cleaned:
            raise ValueError("Could not find two equations or two variables (x, y) in system")

        eq1_str = equations[0].strip()
        eq2_str = equations[1].strip()

        def parse_linear_eq(eq_str):
            if '-' not in eq_str:
                raise ValueError("No equals sign (converted to '-') found")
            left, right = eq_str.split('-')
            expr = sp.sympify(left) - sp.sympify(right or '0')
            a = float(expr.coeff(x, 1)) if expr.coeff(x, 1) else 0
            b = float(expr.coeff(y, 1)) if expr.coeff(y, 1) else 0
            c = float(-expr.as_coefficients_dict()[1]) if 1 in expr.as_coefficients_dict() else 0
            return f"{a} {b} {c}"

        eq1_coeffs = parse_linear_eq(eq1_str)
        eq2_coeffs = parse_linear_eq(eq2_str)
        return {
            "type": "linear",
            "eq1_coeffs": eq1_coeffs,
            "eq2_coeffs": eq2_coeffs,
            "latex": latex_str,
            "success": True
        }
    except Exception as e:
        logger.error(f"Error extracting linear system coefficients: {e}")
        return {
            "type": "linear",
            "eq1_coeffs": "1 1 3",
            "eq2_coeffs": "1 -1 1",
            "latex": latex_str,
            "success": False,
            "error": str(e)
        }

def solve_linear_system_from_coeffs(eq1_str, eq2_str):
    try:
        coeffs1 = list(map(float, eq1_str.strip().split()))
        coeffs2 = list(map(float, eq2_str.strip().split()))

        if len(coeffs1) != 3 or len(coeffs2) != 3:
            return "⚠️ Please enter exactly 3 coefficients for each equation.", None, None, None

        a1, b1, c1 = coeffs1
        a2, b2, c2 = coeffs2

        eq1 = sp.Eq(a1 * x + b1 * y, c1)
        eq2 = sp.Eq(a2 * x + b2 * y, c2)

        sol = sp.solve([eq1, eq2], (x, y), dict=True)
        if not sol:
            return "❌ No unique solution.", None, None, None

        solution = sol[0]
        eq_latex = f"$$ {sp.latex(eq1)} \\ {sp.latex(eq2)} $$"

        steps = rf"""
### Step-by-step Solution
1. **Original Equations:**
   $$ {sp.latex(eq1)} $$
   $$ {sp.latex(eq2)} $$
2. **Standard Form:** Already provided.
3. **Solve using SymPy `solve`:** Internally applies substitution/elimination.
4. **Solve for `x` and `y`:**
   $$ x = {sp.latex(solution[x])}, \quad y = {sp.latex(solution[y])} $$
5. **Verification:** Substitute back into both equations."""

        x_vals = np.linspace(-10, 10, 400)
        f1 = sp.solve(eq1, y)
        f2 = sp.solve(eq2, y)

        fig, ax = plt.subplots()
        if f1:
            f1_func = sp.lambdify(x, f1[0], modules='numpy')
            ax.plot(x_vals, f1_func(x_vals), label=sp.latex(eq1))
        if f2:
            f2_func = sp.lambdify(x, f2[0], modules='numpy')
            ax.plot(x_vals, f2_func(x_vals), label=sp.latex(eq2))

        ax.plot(solution[x], solution[y], 'ro', label=f"Solution ({solution[x]}, {solution[y]})")
        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.legend()
        ax.set_title("Graph of the Linear System")
        ax.grid(True)

        return eq_latex, steps, fig, ""
    except Exception as e:
        return f"❌ Error: {e}", None, None, None

def extract_equation_from_image(image_file):
    try:
        if p2t_model is None:
            return {
                "type": "error",
                "latex": "Pix2Text model not loaded. Please check installation.",
                "success": False
            }
        if image_file is None:
            return {
                "type": "error",
                "latex": "No image file provided.",
                "success": False
            }

        if isinstance(image_file, str):
            image = Image.open(image_file)
        else:
            image = Image.open(image_file.name)

        if image.mode != 'RGB':
            image = image.convert('RGB')

        logger.info(f"Processing image of size: {image.size}")
        result = p2t_model.recognize_text_formula(image)

        if not result or result.strip() == "":
            return {
                "type": "error",
                "latex": "No text or formulas detected in the image.",
                "success": False
            }

        logger.info(f"Extracted text: {result}")
        eq_type = parse_equation_type(result)

        if eq_type == 'polynomial':
            return extract_polynomial_coefficients(result)
        elif eq_type == 'linear_system':
            return extract_linear_system_coefficients(result)
        else:
            return {
                "type": "error",
                "latex": f"Unsupported equation type detected: {eq_type}",
                "success": False
            }

    except Exception as e:
        logger.error(f"Error processing image: {e}")
        return {
            "type": "error",
            "latex": f"Error processing image: {str(e)}",
            "success": False
        }

def solve_extracted_equation(eq_data, real_only):
    if eq_data["type"] == "polynomial":
        return solve_polynomial(eq_data["degree"], eq_data["coeffs"], real_only, eq_data.get("variable", "x"))
    elif eq_data["type"] == "linear":
        return "❌ Single linear equation not supported. Please upload a system of equations.", None, ""
    elif eq_data["type"] == "linear_system":
        return solve_linear_system_from_coeffs(eq_data["eq1_coeffs"], eq_data["eq2_coeffs"])
    else:
        return "❌ Unknown equation type", None, ""

def image_tab():
    """Create the Image Upload Solver tab"""
    with gr.Tab("Image Upload Solver"):
        gr.Markdown("## Solve Equations from Image")

        with gr.Row():
            image_input = gr.File(
                label="Upload Question Image",
                file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                file_count="single"
            )
            image_upload_btn = gr.Button("Process Image")

        gr.Markdown("**Supported Formats:** .pdf, .png, .jpg, .jpeg")

        with gr.Row():
            real_image_checkbox = gr.Checkbox(label="Show Only Real Roots (for Polynomials)", value=False)
            preview_image_btn = gr.Button("Preview Equation")

        image_equation_display = gr.Markdown()

        with gr.Row():
            confirm_image_btn = gr.Button("Display Solution", visible=False)
            edit_image_btn = gr.Button("Make Changes Manually", visible=False)

        edit_latex_input = gr.Textbox(label="Edit LaTeX Equation", visible=False, lines=3)
        save_edit_btn = gr.Button("Save Changes", visible=False)

        image_steps_md = gr.Markdown()
        image_plot_output = gr.Plot()
        extracted_eq_state = gr.State()

        llm_url_input = gr.Textbox(label="LLM Microservice URL (optional)", placeholder="https://your-llm.ngrok.app")
        explain_image_btn = gr.Button("Explain with LLM")
        image_solution_txt = gr.Textbox(visible=False)

        def handle_image_upload(image_file):
            if image_file is None:
                return "", None, "", None, None
            try:
                eq_data = extract_equation_from_image(image_file)
                return "", eq_data, "", None, None
            except Exception:
                return "", None, "", None, None

        image_upload_btn.click(
            fn=handle_image_upload,
            inputs=[image_input],
            outputs=[image_equation_display, extracted_eq_state, image_steps_md,
                     image_plot_output, edit_latex_input]
        )

        def preview_image_equation(eq_data, real_only):
            if not eq_data:
                return "⚠️ No equation data available.", gr.update(visible=False), gr.update(visible=False), "", None
            if eq_data["type"] == "error":
                return eq_data["latex"], gr.update(visible=False), gr.update(visible=False), "", None

            preview_text = f"""
### ✅ Confirm {'Polynomial' if eq_data['type'] == 'polynomial' else 'Linear System'}
**Extracted LaTeX:** {eq_data['latex']}
            """
            return preview_text, gr.update(visible=True), gr.update(visible=True), "", None

        preview_image_btn.click(
            fn=preview_image_equation,
            inputs=[extracted_eq_state, real_image_checkbox],
            outputs=[image_equation_display, confirm_image_btn, edit_image_btn,
                     image_steps_md, image_plot_output]
        )

        def confirm_image_solution(eq_data, real_only):
            if not eq_data or eq_data["type"] == "error":
                return "⚠️ No valid equation to solve.", None, ""
            try:
                steps, plot, _ = solve_extracted_equation(eq_data, real_only)
                return steps, plot, gr.update(value=steps)
            except Exception as e:
                return f"❌ Error solving equation: {str(e)}", None, ""

        confirm_image_btn.click(
            fn=confirm_image_solution,
            inputs=[extracted_eq_state, real_image_checkbox],
            outputs=[image_steps_md, image_plot_output, image_solution_txt]
        )

        def enable_manual_edit(eq_data):
            latex_value = eq_data.get("latex", "") if eq_data and eq_data["type"] != "error" else "Error in extraction."
            return (
                gr.update(visible=True, value=latex_value),
                gr.update(visible=True),
                gr.update(visible=False),
                gr.update(visible=False)
            )

        edit_image_btn.click(
            fn=enable_manual_edit,
            inputs=[extracted_eq_state],
            outputs=[edit_latex_input, save_edit_btn, confirm_image_btn, edit_image_btn]
        )

        def save_manual_changes(latex_input, real_only):
            try:
                if not latex_input.strip():
                    return "⚠️ Please enter a valid equation.", None, ""
                eq_type = parse_equation_type(latex_input)
                if eq_type == 'polynomial':
                    eq_data = extract_polynomial_coefficients(latex_input)
                    steps, plot, _ = solve_polynomial(eq_data["degree"], eq_data["coeffs"], real_only, eq_data.get("variable", "x"))
                elif eq_type == 'linear_system':
                    eq_data = extract_linear_system_coefficients(latex_input)
                    _, steps, plot, _ = solve_linear_system_from_coeffs(eq_data["eq1_coeffs"], eq_data["eq2_coeffs"])
                else:
                    return "❌ Unsupported equation type", None, ""
                return steps, plot, gr.update(value=steps)
            except Exception as e:
                return f"❌ Error parsing manual input: {str(e)}", None, ""

        save_edit_btn.click(
            fn=save_manual_changes,
            inputs=[edit_latex_input, real_image_checkbox],
            outputs=[image_steps_md, image_plot_output, image_solution_txt]
        )

        explain_image_btn.click(
            fn=lambda sol, url: explain_with_llm(sol, "image", url),
            inputs=[image_solution_txt, llm_url_input],
            outputs=[image_steps_md]
        )

    return (
        image_input, image_upload_btn, real_image_checkbox, preview_image_btn,
        image_equation_display, confirm_image_btn, edit_image_btn, edit_latex_input,
        save_edit_btn, image_steps_md, image_plot_output, extracted_eq_state,
        llm_url_input, explain_image_btn, image_solution_txt
    )