"""Gradio app for Maritime Intelligence Classifier."""
import gradio as gr
from setfit import SetFitModel
from pathlib import Path
import os

# Try to load model from Hugging Face Hub first, then fall back to local
# Set MODEL_PATH environment variable or update this line with your Hugging Face repo ID
MODEL_PATH = os.getenv("MODEL_PATH", "gamaly/maritime-intelligence-classifier")
LOCAL_MODEL_PATH = "./maritime_classifier"

# Load model
print("Loading model...")
try:
    # Try Hugging Face Hub first
    if "/" in MODEL_PATH and not Path(MODEL_PATH).exists():
        model = SetFitModel.from_pretrained(MODEL_PATH)
        print(f"✓ Loaded model from Hugging Face: {MODEL_PATH}")
    else:
        # Try local path
        if Path(LOCAL_MODEL_PATH).exists():
            model = SetFitModel.from_pretrained(LOCAL_MODEL_PATH)
            print(f"✓ Loaded model from local path: {LOCAL_MODEL_PATH}")
        else:
            raise FileNotFoundError(f"Model not found at {MODEL_PATH} or {LOCAL_MODEL_PATH}")
except Exception as e:
    print(f"⚠️  Error loading model: {e}")
    print("Make sure the model is trained or uploaded to Hugging Face")
    model = None

def predict_text(text):
    """Predict whether text is actionable (YES) or not (NO)."""
    if model is None:
        return "Error: Model not loaded. Please train the model first.", 0.0, "error"
    
    if not text or not text.strip():
        return "Please enter some text to classify.", 0.0, "neutral"
    
    try:
        # Make prediction
        prediction = model.predict([text])[0]
        probabilities = model.predict_proba([text])[0]
        
        # Get confidence
        confidence = probabilities[prediction] * 100
        
        # Convert to labels
        label = "YES (Actionable)" if prediction == 1 else "NO (Not Actionable)"
        
        # Determine status for styling
        status = "actionable" if prediction == 1 else "not_actionable"
        
        return label, confidence, status
    except Exception as e:
        return f"Error during prediction: {str(e)}", 0.0, "error"

def get_explanation(status):
    """Get explanation based on prediction status."""
    explanations = {
        "actionable": "✓ This text contains actionable vessel-specific evidence (e.g., specific vessel names, crimes, incidents).",
        "not_actionable": "✗ This text does not contain actionable vessel-specific evidence (e.g., general maritime news, non-specific information).",
        "error": "⚠️ An error occurred. Please check the model is properly loaded.",
        "neutral": ""
    }
    return explanations.get(status, "")

# Create Gradio interface
with gr.Blocks(title="Maritime Intelligence Classifier", theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        # 🚢 Maritime Intelligence Classifier
        
        Classify maritime news articles as containing **actionable vessel-specific evidence** (YES) or not (NO).
        
        **Actionable articles** typically include:
        - Specific vessel names
        - Specific crimes or incidents
        - Evidence that can be used for investigation
        
        **Non-actionable articles** are general maritime news without specific vessel details.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Article Text",
                placeholder="Paste or type the maritime news article text here...",
                lines=10,
                max_lines=20
            )
            
            submit_btn = gr.Button("Classify", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            prediction_output = gr.Label(
                label="Prediction",
                value={"YES (Actionable)": 0.0, "NO (Not Actionable)": 0.0}
            )
            
            confidence_output = gr.Number(
                label="Confidence",
                value=0.0,
                precision=1
            )
            
            explanation_output = gr.Markdown()
    
    # Example texts
    gr.Markdown("### 📝 Example Texts")
    with gr.Row():
        example_yes = gr.Examples(
            examples=[
                ["The fishing vessel Marine 707 was involved in the disappearance of fisheries observer Samuel Abayateye in Ghanaian waters. The observer's decapitated body was found weeks later."],
                ["Authorities detained the Meng Xin 15 after discovering evidence of illegal saiko transshipment and threats against fisheries observers."],
            ],
            inputs=text_input,
            label="YES Examples (Actionable)"
        )
        
        example_no = gr.Examples(
            examples=[
                ["A new maritime museum opened in the port city, showcasing historical ships and ocean exploration artifacts."],
                ["Marine scientists are studying the effects of ocean acidification on coral reefs in tropical waters."],
            ],
            inputs=text_input,
            label="NO Examples (Not Actionable)"
        )
    
    # Connect the prediction function
    def update_prediction(text):
        label, confidence, status = predict_text(text)
        
        # Create label dict for gradio Label component
        if status == "actionable":
            label_dict = {"YES (Actionable)": confidence / 100, "NO (Not Actionable)": (100 - confidence) / 100}
        elif status == "not_actionable":
            label_dict = {"YES (Actionable)": (100 - confidence) / 100, "NO (Not Actionable)": confidence / 100}
        else:
            label_dict = {"YES (Actionable)": 0.0, "NO (Not Actionable)": 0.0}
        
        explanation = get_explanation(status)
        
        return label_dict, confidence, explanation
    
    submit_btn.click(
        fn=update_prediction,
        inputs=text_input,
        outputs=[prediction_output, confidence_output, explanation_output]
    )
    
    text_input.submit(
        fn=update_prediction,
        inputs=text_input,
        outputs=[prediction_output, confidence_output, explanation_output]
    )
    
    gr.Markdown(
        """
        ---
        ### ℹ️ About
        
        This classifier uses SetFit to identify maritime news articles containing actionable vessel-specific evidence.
        Built for The Outlaw Ocean Project.
        
        **Model**: SetFit (sentence-transformers/all-MiniLM-L6-v2 base)
        """
    )

if __name__ == "__main__":
    app.launch(share=False)