import gradio as gr
import os
from simple_classifier import SimpleRockClassifier


classifier = SimpleRockClassifier()

def analyze_core(image):
    """Analyze a drill core image"""
    # Save uploaded image temporarily
    temp_path = "temp_upload.jpg"
    image.save(temp_path)
    
    # Get prediction
    try:
        result = classifier.predict(temp_path)
        
        # Format response
        response = f"""
        ##  Drill Core Analysis Results
        
        ### Primary Prediction
        **Rock Type:** `{result['rock_type']}`  
        **Confidence:** `{result['confidence']:.2f}`
        
        ### Analysis Details
        {result['explanation']}
        """
        
    except Exception as e:
        response = f"##  Error\nAn error occurred during analysis: {str(e)}"
    
    
    if os.path.exists(temp_path):
        os.remove(temp_path)
    
    return response

# Create Gradio interface
with gr.Blocks(title="Geologist_AI - Core Logger") as demo:
    gr.Markdown("#  Geologist_AI - Core Logger")
    gr.Markdown("Upload a drill core image to identify the rock type")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="📷 Drill Core Image")
            submit_btn = gr.Button("🔍 Analyze Core Sample", variant="primary")
        with gr.Column():
            output_text = gr.Markdown(label="📊 Analysis Results")
    
    submit_btn.click(
        fn=analyze_core,
        inputs=image_input,
        outputs=output_text
    )
    
    gr.Markdown("---")
    gr.Markdown("### About this Tool")
    gr.Markdown("""
    This AI-powered geologist identifies rock types based on:
    - **Visual color analysis**
    - **Deep learning feature extraction**
    
    **Supported rock types:**
    - Gold-bearing rock 
    - Iron-rich rock 
    - Lithium-rich rock 
    - Copper-bearing rock 
    - Quartz-rich rock 
    - Waste rock 
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()