import gradio as gr
import json
import time
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import torch
import warnings
warnings.filterwarnings('ignore')

# Import the correct module based on the repository structure
try:
    from chemistry_llm import ChemistryReactionExtractor
except ImportError:
    # Fallback for different import structure
    try:
        from chemistry_llm.core.extractor import ChemistryReactionExtractor
    except ImportError:
        print("Warning: ChemistryReactionExtractor not found. Using mock implementation.")
        ChemistryReactionExtractor = None

# Global variables
extractor = None
model_loading = False

def load_model():
    """Load the RxNExtract model"""
    global extractor, model_loading
    
    if model_loading:
        return "Model is currently loading, please wait..."
    
    if extractor is not None:
        return "Model already loaded!"
    
    if ChemistryReactionExtractor is None:
        return "❌ ChemistryReactionExtractor module not available. Please check the installation."
    
    model_loading = True
    try:
        # Initialize the extractor with proper configuration based on repository documentation
        extractor = ChemistryReactionExtractor(
            model_path="chemplusx/rxnextract-complete",  # Use the model from repository
            device="cuda" if torch.cuda.is_available() else "cpu",
            config={
                "quantization": {
                    "load_in_4bit": True,
                    "bnb_4bit_quant_type": "nf4",
                    "bnb_4bit_compute_dtype": "float16"
                },
                "model": {
                    "default_temperature": 0.1,
                    "max_new_tokens": 512
                }
            }
        )
        model_loading = False
        return "✅ RxNExtract model loaded successfully!"
    except Exception as e:
        model_loading = False
        return f"❌ Error loading model: {str(e)}"

def analyze_procedure(procedure_text, temperature=0.1):
    """Analyze a chemical procedure using the actual RxNExtract API"""
    global extractor
    
    if extractor is None:
        return "Please load the model first!", "", "", ""
    
    if not procedure_text.strip():
        return "Please enter a chemical procedure to analyze.", "", "", ""
    
    try:
        start_time = time.time()
        
        # Use the correct API method from the repository
        results = extractor.analyze_procedure(
            procedure_text=procedure_text,
            return_raw=False
        )
        
        processing_time = time.time() - start_time
        
        # Add processing time to results
        if isinstance(results, dict):
            results['processing_time'] = processing_time
        
        # Format the results
        formatted_output = format_extraction_results(results, processing_time)
        
        # Create visualizations
        entity_plot = create_entity_visualization(results)
        confidence_plot = create_confidence_visualization(results, processing_time)
        
        # Create summary
        summary = create_summary(results, processing_time)
        
        return summary, formatted_output, entity_plot, confidence_plot
        
    except Exception as e:
        error_msg = f"❌ Error during analysis: {str(e)}"
        return error_msg, "", "", ""

def format_extraction_results(results, processing_time):
    """Format extraction results for display based on actual API structure"""
    if not isinstance(results, dict):
        return "Error: Invalid results format"
    
    # Handle the actual data structure from RxNExtract
    extracted_data = results.get('extracted_data', results)
    confidence = results.get('confidence', 'N/A')
    
    output = []
    output.append("## 📊 Extraction Results\n")
    
    if confidence != 'N/A':
        output.append(f"**🎯 Confidence:** {confidence:.1%}" if isinstance(confidence, float) else f"**🎯 Confidence:** {confidence}")
    output.append(f"**⏱️ Processing Time:** {processing_time:.1f}s\n")
    
    # Handle different possible data structures
    if isinstance(extracted_data, dict):
        # Reactants
        reactants = extracted_data.get('reactants', [])
        if reactants:
            output.append("### 🔵 Reactants")
            for i, reactant in enumerate(reactants, 1):
                if isinstance(reactant, dict):
                    name = reactant.get('name', reactant.get('compound', 'Unknown'))
                    amount = reactant.get('amount', reactant.get('quantity', 'N/A'))
                    output.append(f"{i}. **{name}** - Amount: {amount}")
                else:
                    output.append(f"{i}. **{reactant}**")
            output.append("")
        
        # Reagents
        reagents = extracted_data.get('reagents', [])
        if reagents:
            output.append("### 🟡 Reagents")
            for i, reagent in enumerate(reagents, 1):
                if isinstance(reagent, dict):
                    name = reagent.get('name', reagent.get('compound', 'Unknown'))
                    amount = reagent.get('amount', reagent.get('quantity', 'N/A'))
                    output.append(f"{i}. **{name}** - Amount: {amount}")
                else:
                    output.append(f"{i}. **{reagent}**")
            output.append("")
        
        # Solvents
        solvents = extracted_data.get('solvents', [])
        if solvents:
            output.append("### 🔵 Solvents")
            for i, solvent in enumerate(solvents, 1):
                if isinstance(solvent, dict):
                    name = solvent.get('name', solvent.get('compound', 'Unknown'))
                    amount = solvent.get('amount', solvent.get('quantity', 'N/A'))
                    output.append(f"{i}. **{name}** - Amount: {amount}")
                else:
                    output.append(f"{i}. **{solvent}**")
            output.append("")
        
        # Products
        products = extracted_data.get('products', [])
        if products:
            output.append("### 🟢 Products")
            for i, product in enumerate(products, 1):
                if isinstance(product, dict):
                    name = product.get('name', product.get('compound', 'Unknown'))
                    amount = product.get('amount', product.get('quantity', 'N/A'))
                    yield_val = product.get('yield', 'N/A')
                    output.append(f"{i}. **{name}** - Amount: {amount}, Yield: {yield_val}")
                else:
                    output.append(f"{i}. **{product}**")
            output.append("")
        
        # Conditions
        conditions = extracted_data.get('conditions', {})
        if conditions:
            output.append("### 🌡️ Reaction Conditions")
            if isinstance(conditions, dict):
                for key, value in conditions.items():
                    if value:
                        output.append(f"- **{key.title()}:** {value}")
            else:
                output.append(f"- {conditions}")
            output.append("")
        
        # Workup steps
        workup = extracted_data.get('workup', extracted_data.get('workup_steps', []))
        if workup:
            output.append("### ⚗️ Workup Steps")
            if isinstance(workup, list):
                for i, step in enumerate(workup, 1):
                    output.append(f"{i}. {step}")
            else:
                output.append(f"1. {workup}")
            output.append("")
    
    if len(output) <= 3:  # Only header and processing time
        output.append("No specific reaction data extracted. The model may need adjustment or the procedure text may not contain clear chemical information.")
    
    return "\n".join(output)

def create_entity_visualization(results):
    """Create entity count visualization"""
    if not isinstance(results, dict):
        return None
    
    extracted_data = results.get('extracted_data', results)
    
    if not isinstance(extracted_data, dict):
        return None
    
    # Count entities
    entity_counts = {
        'Reactants': len(extracted_data.get('reactants', [])),
        'Reagents': len(extracted_data.get('reagents', [])),
        'Solvents': len(extracted_data.get('solvents', [])),
        'Products': len(extracted_data.get('products', [])),
        'Conditions': len(extracted_data.get('conditions', {})) if isinstance(extracted_data.get('conditions', {}), dict) else 1 if extracted_data.get('conditions') else 0,
        'Workup Steps': len(extracted_data.get('workup', extracted_data.get('workup_steps', [])))
    }
    
    # Remove zero counts
    entity_counts = {k: v for k, v in entity_counts.items() if v > 0}
    
    if not entity_counts:
        return None
    
    # Create bar chart
    fig = px.bar(
        x=list(entity_counts.keys()),
        y=list(entity_counts.values()),
        title="Extracted Chemical Entities",
        labels={'x': 'Entity Type', 'y': 'Count'},
        color=list(entity_counts.keys()),
        color_discrete_sequence=px.colors.qualitative.Set3
    )
    
    fig.update_layout(
        showlegend=False,
        height=400,
        title_x=0.5,
        xaxis_tickangle=-45
    )
    
    return fig

def create_confidence_visualization(results, processing_time):
    """Create confidence and timing visualization"""
    confidence = results.get('confidence', 0.5) if isinstance(results, dict) else 0.5
    
    # Handle different confidence formats
    if isinstance(confidence, str):
        try:
            confidence = float(confidence)
        except:
            confidence = 0.5
    
    # Create gauge chart for confidence
    fig = go.Figure(go.Indicator(
        mode = "gauge+number+delta",
        value = confidence * 100 if confidence <= 1.0 else confidence,
        domain = {'x': [0, 1], 'y': [0, 1]},
        title = {'text': "Confidence Score (%)"},
        delta = {'reference': 80},
        gauge = {
            'axis': {'range': [None, 100]},
            'bar': {'color': "darkblue"},
            'steps': [
                {'range': [0, 50], 'color': "lightgray"},
                {'range': [50, 80], 'color': "yellow"},
                {'range': [80, 100], 'color': "green"}],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': 90}}))
    
    fig.update_layout(
        height=300,
        title=f"Processing Time: {processing_time:.1f}s"
    )
    
    return fig

def create_summary(results, processing_time):
    """Create a summary of the analysis"""
    if not isinstance(results, dict):
        return "## 📈 Analysis Summary\nError processing results."
    
    extracted_data = results.get('extracted_data', results)
    confidence = results.get('confidence', 'N/A')
    
    # Calculate total entities
    total_entities = 0
    if isinstance(extracted_data, dict):
        total_entities = sum([
            len(extracted_data.get('reactants', [])),
            len(extracted_data.get('reagents', [])),
            len(extracted_data.get('solvents', [])),
            len(extracted_data.get('products', []))
        ])
    
    # Determine confidence level
    confidence_level = "Unknown"
    if isinstance(confidence, (int, float)):
        confidence_val = confidence if confidence <= 1.0 else confidence / 100
        confidence_level = "High" if confidence_val >= 0.8 else "Medium" if confidence_val >= 0.6 else "Low"
    
    # Format confidence display
    conf_display = f"{confidence:.1%}" if isinstance(confidence, float) and confidence <= 1.0 else str(confidence)
    
    summary = f"""
## 📈 Analysis Summary

**🎯 Overall Performance:**
- **Confidence Level:** {confidence_level} ({conf_display})
- **Processing Speed:** {processing_time:.1f} seconds
- **Total Entities Extracted:** {total_entities}

**📊 Extraction Breakdown:**
- **Reactants:** {len(extracted_data.get('reactants', [])) if isinstance(extracted_data, dict) else 0}
- **Products:** {len(extracted_data.get('products', [])) if isinstance(extracted_data, dict) else 0}
- **Reagents:** {len(extracted_data.get('reagents', [])) if isinstance(extracted_data, dict) else 0}
- **Solvents:** {len(extracted_data.get('solvents', [])) if isinstance(extracted_data, dict) else 0}
- **Conditions:** {len(extracted_data.get('conditions', {})) if isinstance(extracted_data, dict) and isinstance(extracted_data.get('conditions', {}), dict) else (1 if isinstance(extracted_data, dict) and extracted_data.get('conditions') else 0)}
- **Workup Steps:** {len(extracted_data.get('workup', extracted_data.get('workup_steps', []))) if isinstance(extracted_data, dict) else 0}

**💡 Quality Assessment:**
{get_quality_assessment(confidence, total_entities)}
"""
    
    return summary

def get_quality_assessment(confidence, total_entities):
    """Get quality assessment based on confidence and entities"""
    # Handle different confidence formats
    if isinstance(confidence, (int, float)):
        conf_val = confidence if confidence <= 1.0 else confidence / 100
    else:
        conf_val = 0.5  # Default for unknown confidence
    
    if conf_val >= 0.8 and total_entities >= 3:
        return "✅ Excellent extraction quality with high confidence and comprehensive entity recognition."
    elif conf_val >= 0.6 and total_entities >= 2:
        return "✅ Good extraction quality with moderate confidence. Results are reliable."
    elif conf_val >= 0.4:
        return "⚠️ Moderate extraction quality. Some information may be missing or uncertain."
    else:
        return "❌ Low extraction quality. Consider reviewing the procedure text for clarity."

def get_example_procedures():
    """Get example procedures for the interface"""
    examples = [
        """Add 2.5 g of benzoic acid to 50 mL of ethanol in a round-bottom flask.
Heat the mixture to reflux for 4 hours while stirring.
Cool the solution to room temperature and filter the precipitate.
Wash the solid with cold ethanol and dry to obtain 2.1 g of product (84% yield).""",
        
        """Dissolve 10.0 g of 4-nitroaniline in 200 mL of concentrated HCl.
Add 15.0 g of tin powder portionwise while maintaining temperature below 10°C.
Stir for 2 hours at room temperature, then heat to 60°C for 1 hour.
Neutralize with NaOH solution and extract with ethyl acetate (3 × 50 mL).
Dry over MgSO4 and concentrate to give 7.2 g of product (78% yield).""",
        
        """In a round-bottom flask, combine 5.0 mmol of styrene, 6.0 mmol of phenylboronic acid,
and 0.1 mmol of Pd(PPh3)4 catalyst in 20 mL of DMF.
Add 15.0 mmol of K2CO3 and heat to 100°C under nitrogen atmosphere.
Stir for 12 hours, then cool and filter through celite.
Purify by column chromatography to obtain 0.85 g of biphenyl derivative (92% yield)."""
    ]
    return examples

# Create the Gradio interface
def create_interface():
    """Create and return the Gradio interface"""
    
    with gr.Blocks(title="RxNExtract - Chemical Reaction Extraction", theme=gr.themes.Soft()) as demo:
        
        # Header
        gr.Markdown("""
        # 🧪 RxNExtract - Chemical Reaction Extraction
        
        Extract chemical entities and reaction information from synthetic procedures using advanced NLP models.
        
        **Professional-grade system for extracting chemical reaction information from procedure texts using fine-tuned LLM with Dynamic prompting and self grounding.**
        """)
        
        # Model loading section
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("### 🤖 Model Management")
                load_btn = gr.Button("Load RxNExtract Model", variant="primary", size="lg")
                model_status = gr.Textbox(
                    label="Model Status",
                    value="Model not loaded. Click 'Load RxNExtract Model' to initialize.",
                    interactive=False
                )
        
        # Main interface
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📝 Input")
                
                procedure_input = gr.Textbox(
                    label="Chemical Procedure",
                    placeholder="Enter your chemical synthesis procedure here...",
                    lines=8,
                    max_lines=15
                )
                
                with gr.Row():
                    temperature_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.1,
                        step=0.1,
                        label="Temperature (Model Creativity)",
                        info="Lower values = more focused, higher values = more creative"
                    )
                
                with gr.Row():
                    analyze_btn = gr.Button("🔍 Analyze Procedure", variant="primary", size="lg")
                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                
                # Example procedures
                gr.Markdown("### 📋 Example Procedures")
                examples = get_example_procedures()
                
                for i, example in enumerate(examples, 1):
                    with gr.Accordion(f"Example {i}: {['Benzoic Acid Synthesis', 'Aniline Reduction', 'Suzuki Coupling'][i-1]}", open=False):
                        example_text = gr.Textbox(
                            value=example,
                            label=f"Example {i}",
                            lines=4,
                            interactive=False
                        )
                        use_example_btn = gr.Button(f"Use Example {i}", size="sm")
                        use_example_btn.click(
                            fn=lambda ex=example: ex,
                            outputs=procedure_input
                        )
            
            with gr.Column(scale=2):
                gr.Markdown("### 📊 Results")
                
                # Summary tab
                with gr.Tabs():
                    with gr.TabItem("📈 Summary"):
                        summary_output = gr.Markdown()
                    
                    with gr.TabItem("📋 Detailed Results"):
                        detailed_output = gr.Markdown()
                    
                    with gr.TabItem("📊 Entity Visualization"):
                        entity_plot = gr.Plot()
                    
                    with gr.TabItem("🎯 Confidence & Timing"):
                        confidence_plot = gr.Plot()
        
        # Event handlers
        load_btn.click(
            fn=load_model,
            outputs=model_status
        )
        
        analyze_btn.click(
            fn=analyze_procedure,
            inputs=[procedure_input, temperature_slider],
            outputs=[summary_output, detailed_output, entity_plot, confidence_plot]
        )
        
        clear_btn.click(
            fn=lambda: ("", "", "", "", ""),
            outputs=[procedure_input, summary_output, detailed_output, entity_plot, confidence_plot]
        )
        
        # Footer
        gr.Markdown("""
        ---
        **About RxNExtract:** This tool uses advanced natural language processing to extract chemical entities, 
        reaction conditions, and procedural information from synthetic chemistry procedures.
        
        **Features:**
        - Modular Architecture with clean, maintainable codebase
        - Dynamic Prompting for better extraction accuracy
        - Memory Efficient 4-bit quantization support
        - Robust XML parsing with structured output
        - Professional logging and error handling
        
        **Powered by:** ChemPlusX Team | [GitHub Repository](https://github.com/chemplusx/RxNExtract)
        """)
    
    return demo

# Main execution
if __name__ == "__main__":
    # Create and launch the interface
    demo = create_interface()
    
    # Launch with appropriate settings for Hugging Face Spaces
    demo.launch(
        server_name="0.0.0.0",  # Required for Hugging Face Spaces
        server_port=7860,       # Default port for Hugging Face Spaces
        share=False,            # Don't create public links
        debug=False,            # Disable debug mode in production
        show_error=True         # Show errors in the interface
    )