app.py

import gradio as gr
import os
import time
import cv2
import numpy as np
from PIL import Image
import tempfile
import json
import re

# Import consolidated modules
from ocr_module import MVM2OCREngine
from reasoning_engine import run_agent_orchestrator
from verification_service import calculate_symbolic_score
from consensus_fusion import evaluate_consensus
from report_module import generate_mvm2_report, export_to_pdf
from image_enhancing import ImageEnhancer
from flow_module import generate_flow_html

# Initialize Engines
ocr_engine = MVM2OCREngine()
enhancer = ImageEnhancer(sigma=1.2)

# Load custom CSS
with open("theme.css", "r") as f:
    css_content = f.read()

def create_gauge(label, value, color="#6366f1"):
    """Generates an animated SVG circular gauge."""
    percentage = max(0, min(100, value * 100))
    dash_offset = 251.2 * (1 - percentage / 100)
    return f"""
    <div class="gauge-container">
        <svg width="100" height="100" viewBox="0 0 100 100">
            <circle class="circle-bg" cx="50" cy="50" r="40" />
            <circle class="circle-progress" cx="50" cy="50" r="40" 
                    stroke="{color}" stroke-dasharray="251.2" 
                    stroke-dashoffset="{dash_offset}" 
                    style="filter: drop-shadow(0 0 5px {color}88);"/>
            <text x="50" y="55" text-anchor="middle" font-size="18" font-weight="bold" fill="white">{int(percentage)}%</text>
        </svg>
        <div style="font-size: 0.8em; color: #94a3b8; font-weight: 500;">{label}</div>
    </div>
    """

def format_step_viewer(consensus_result):
    """Formats the Reasoning Trace with Step-Level Consensus highlights."""
    html = '<div style="display: flex; flex-direction: column; gap: 12px;">'
    
    # We aggregate steps from all agents for a collective view
    agent_data = consensus_result.get("detail_scores", [])
    
    for agent in agent_data:
        # Simulate step-level analysis for UI purposes: 
        # In a real system, we'd have Score_j per step. Here we use the agent's overall score.
        score = agent["Score_j"]
        status_class = "step-valid" if score >= 0.7 else "step-warning"
        icon = "✅" if score >= 0.7 else "⚠️"
        glow_style = "box-shadow: 0 0 10px rgba(16, 185, 129, 0.2);" if score >= 0.7 else "box-shadow: 0 0 10px rgba(245, 158, 11, 0.2);"
        
        # Get matching agent response for trace
        # (This assumes agent_responses were passed in or stored)
        # For the UI, we'll just show the representative trace from valid agents
        if not agent["is_hallucinating"] or score > 0.4:
            verification_msg = f"✅ Ast-Parsed Answer matches consensus group." if score >= 0.7 else f"❌ Answer diverged from consensus."
            html += f"""
            <div class="glass-card reasoning-step {status_class}" style="{glow_style}">
                <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px;">
                    <span class="monospace" style="color: #6366f1; font-weight: 600;">{agent['agent']} Reasoning Path</span>
                    <span style="font-size: 0.8em; background: rgba(0,0,0,0.3); padding: 2px 8px; border-radius: 4px;">Consensus: {score:.2f} {icon}</span>
                </div>
                <div style="font-size: 0.9em; line-height: 1.6; color: #cbd5e1; margin-bottom: 8px;">
                    {"<br>".join([f"• {step}" for step in agent.get('reasoning_trace', ['Processing...'])])}
                </div>
                <div style="font-size: 0.85em; padding: 6px; background: rgba(0,0,0,0.2); border-left: 3px solid {'#10b981' if score >= 0.7 else '#ef4444'}; color: {'#34d399' if score >= 0.7 else '#f87171'};">
                    <strong>math_verify check:</strong> {verification_msg}
                </div>
            </div>
            """
    html += "</div>"
    return html

def process_mvm2_pipeline(image, auto_enhance):
    if image is None:
        yield None, "Please upload an image.", None, "", None, None, generate_flow_html("idle")
        return
    
    # 1. Preprocessing & Preview
    yield None, "Enhancing image...", None, "", None, None, generate_flow_html("enhance")
    enhanced_img_np, meta = enhancer.enhance(image)
    temp_img_path = os.path.join(tempfile.gettempdir(), 'input_processed.png')
    cv2.imwrite(temp_img_path, enhanced_img_np)
    
    preview_img = Image.fromarray(cv2.cvtColor(enhanced_img_np, cv2.COLOR_BGR2RGB))
    yield preview_img, "Extracting LaTeX...", None, "", None, None, generate_flow_html("ocr")

    # 2. OCR Extraction
    ocr_results = ocr_engine.process_image(temp_img_path)
    latex_text = ocr_results['latex_output']
    ocr_conf = ocr_results['weighted_confidence']
    
    yield preview_img, latex_text, None, "", None, None, generate_flow_html("reasoning")

    # 3. Multi-Agent Reasoning
    agent_responses = run_agent_orchestrator(latex_text)

    # 4. Advanced Heuristics Refinement stage
    yield preview_img, latex_text, None, "", None, None, generate_flow_html("heuristics")
    time.sleep(0.5)

    # 5. Consensus Fusion
    yield preview_img, latex_text, None, "", None, None, generate_flow_html("consensus")
    consensus_result = evaluate_consensus(agent_responses, ocr_confidence=ocr_conf)
    
    # Attach traces back to detail_scores for UI formatting
    for i, score_data in enumerate(consensus_result["detail_scores"]):
        for res in agent_responses:
            if res["agent"] == score_data["agent"]:
                consensus_result["detail_scores"][i]["reasoning_trace"] = res["response"].get("Reasoning Trace", [])
                break

    # 6. Gauges & UI Elements
    avg_v_sym = np.mean([s["V_sym"] for s in consensus_result["detail_scores"]])
    avg_l_logic = np.mean([s["L_logic"] for s in consensus_result["detail_scores"]])
    avg_c_clf = np.mean([s["C_clf"] for s in consensus_result["detail_scores"]])
    
    gauges_html = f"""
    <div class="signal-panel">
        {create_gauge("Symbolic", avg_v_sym, "#10b981")}
        {create_gauge("Logic", avg_l_logic, "#6366f1")}
        {create_gauge("Classifier", avg_c_clf, "#8b5cf6")}
    </div>
    """
    
    winner = consensus_result["winning_score"]
    calibrated_conf = winner * (0.9 + 0.1 * ocr_conf)
    conf_bar = f"""
    <div style="margin-top: 20px;">
        <div style="display: flex; justify-content: space-between; margin-bottom: 8px;">
            <span style="font-weight: 600; color: #94a3b8;">Final Confidence Calibration</span>
            <span style="color: #10b981; font-weight: bold;">{calibrated_conf:.3f}</span>
        </div>
        <div style="width: 100%; bg: rgba(255,255,255,0.05); height: 8px; border-radius: 4px; overflow: hidden;">
            <div style="width: {min(100, calibrated_conf*100)}%; background: linear-gradient(90deg, #6366f1 0%, #10b981 100%); height: 100%; transition: width 1s ease;"></div>
        </div>
    </div>
    """

    # 7. Report & PDF
    reports = generate_mvm2_report(consensus_result, latex_text, ocr_conf)
    md_report = format_step_viewer(consensus_result)
    
    pdf_path = os.path.join(tempfile.gettempdir(), f'MVM2_Report_{reports["report_id"]}.pdf')
    export_to_pdf(json.loads(reports['json']), pdf_path)
    
    final_flow = generate_flow_html("success")
    
    yield preview_img, latex_text, gauges_html, conf_bar, md_report, pdf_path, final_flow

# Build Interface
with gr.Blocks(css=css_content, title="MVM²: Senior UI AI Dashboard") as demo:
    with gr.Row(elem_id="header-row"):
        gr.Markdown(
            """
            <div style="text-align: center; padding: 20px 0;">
                <h1 style="font-size: 2.5em; margin-bottom: 0;">MVM² <span style="color: #6366f1;">Neuro-Symbolic</span></h1>
                <p style="color: #94a3b8; font-size: 1.1em; margin-top: 8px;">High-Fidelity Mathematical Verification & Consensus Dashboard</p>
            </div>
            """
        )
    
    with gr.Row():
        # --- LEFT PANEL: Upload & Preview ---
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### 📤 Input Intelligence")
            input_img = gr.Image(type="pil", label="Capture Solution", elem_classes="glass-card")
            enhance_toggle = gr.Checkbox(label="Enable Opti-Scan Preprocessing", value=True)
            run_btn = gr.Button("INITIALIZE VERIFICATION", variant="primary", elem_classes="download-btn")
            
            gr.Markdown("#### 🔍 Preprocessing Preview")
            preview_output = gr.Image(label="Enhanced Signal", interactive=False, elem_classes="preview-img")

        # --- CENTER STAGE: Canvas ---
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Solver Visualization"):
                    gr.Markdown("### 🎨 MVM² Verification Canvas")
                    with gr.Group(elem_classes="glass-card"):
                        canvas_latex = gr.Textbox(label="Canonical LaTeX Transcription", lines=2, interactive=False, elem_classes="monospace")
                        calib_bar_html = gr.HTML()
                    
                    gr.Markdown("### 🪜 Dynamic Reasoning Trace")
                    trace_html = gr.HTML()

                with gr.TabItem("How It Works (Architecture Flow)"):
                    gr.Markdown("### 🚀 Real-Time Pipeline Visualization")
                    flow_view = gr.HTML(generate_flow_html("idle"))
                    gr.Markdown(
                        """
                        **Pipeline Phases:**
                        1. **Enhance:** CLAHE & Gaussian Blur noise reduction.
                        2. **OCR:** Pix2Text LaTeX structure reconstruction.
                        3. **Reasoning:** Quad-agent parallel logic processing.
                        4. **Verification:** SymPy deterministic symbolic check.
                        5. **Consensus:** Multi-signal weighted confidence fusion.
                        """
                    )

        # --- RIGHT PANEL: Signal Intel ---
        with gr.Column(scale=1, variant="panel"):
            gr.Markdown("### ⚡ Signal Intelligence")
            with gr.Group(elem_classes="glass-card"):
                signal_gauges = gr.HTML()
            
            gr.Markdown("### 📄 Educational Assessment")
            download_btn = gr.File(label="Download Diagnostic PDF", elem_classes="download-btn")
            
            with gr.Group(elem_classes="glass-card status-box"):
                gr.Markdown(
                    """
                    **System Status**
                    - Pix2Text VLM: `Online`
                    - SymPy Core: `1.12.0`
                    - Consensus: `4-Agent parallel`
                    """
                )

    run_btn.click(
        fn=process_mvm2_pipeline,
        inputs=[input_img, enhance_toggle],
        outputs=[preview_output, canvas_latex, signal_gauges, calib_bar_html, trace_html, download_btn, flow_view]
    )

if __name__ == "__main__":
    demo.launch()