import gradio
import pandas as pd
import pathlib
import os  # For filesystem operations
import shutil  # For directory cleanup
import zipfile  # For extracting model archives
import pathlib
import huggingface_hub
# import huggingface
from huggingface_hub import login, create_repo, HfApi
import autogluon
from autogluon import tabular

# Model config
MODEL_REPO_ID = "nadakandrew/places-automl"
ZIP_FILENAME = "autogluon_predictor_dir.zip"
CACHE_DIR = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"

# Download & load the native predictor
def _prepare_predictor_dir() -> str:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    local_zip = huggingface_hub.hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
        local_dir=str(CACHE_DIR),
        local_dir_use_symlinks=False,
    )
    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(local_zip, "r") as zf:
        zf.extractall(str(EXTRACT_DIR))
    contents = list(EXTRACT_DIR.iterdir())
    predictor_root = contents[0] if (len(contents) == 1 and contents[0].is_dir()) else EXTRACT_DIR
    return str(predictor_root)

PREDICTOR_DIR = _prepare_predictor_dir()
PREDICTOR = autogluon.tabular.TabularPredictor.load(PREDICTOR_DIR, require_py_version_match=False)

clf = PREDICTOR # autogluon.tabular predictor
print(clf.feature_metadata)

examples = [
    ["The Hidden Mall", "Inside The Hidden Mall, you will find a mix of modern and traditional design. The mall has a reputation for being a popular gathering spot for students and locals."],
    ["The Red Cinema", "Visitors enjoy spending time at The Red Cinema because it offers both comfort and energy. The cinema has a reputation for being a popular gathering spot for students and locals."],
    ["Moonlight Museum", "Visitors enjoy spending time at Moonlight Museum because it offers both comfort and energy. Inside Moonlight Museum, you will find a mix of modern and traditional design."],
    ["Central Park", "A large public park in the heart of the city with walking trails, playgrounds, and open green spaces where families gather for picnics and recreation."],
    ["Downtown Library", "A quiet academic building with extensive book collections, study rooms, and computer access where students and researchers come to work and learn."]
]

def predict_location(place_name, description, confidence_threshold=0.5, return_probabilities=False):
    """Function to predict location type based on inputs with exposed inference parameters"""
    if not description.strip():
        return "Please provide a description", "**Error:** Description cannot be empty"
    
    try:
        input_df = pd.DataFrame({"Place_Name": [place_name], "Description": [description]})
        
        # Make prediction using the classifier
        prediction = clf.predict(input_df)[0]
        
        # Get prediction probabilities if available
        detailed_output = f"**Predicted Location Type:** {prediction}"
        
        if return_probabilities:
            try:
                # Try to get prediction probabilities
                probabilities = clf.predict_proba(input_df)
                if probabilities is not None:
                    if hasattr(probabilities, 'iloc'):
                        prob_dict = probabilities.iloc[0].to_dict()
                    else:
                        prob_dict = probabilities
                    
                    # Sort by probability and show top predictions
                    sorted_probs = sorted(prob_dict.items(), key=lambda x: x[1], reverse=True)
                    
                    detailed_output += "\n\n**Prediction Probabilities:**\n"
                    for class_name, prob in sorted_probs[:5]:  # Show top 5
                        if prob >= confidence_threshold:
                            detailed_output += f"- {class_name}: {prob:.3f}\n"
                
            except Exception as prob_error:
                detailed_output += f"\n\nNote: Could not retrieve probabilities ({str(prob_error)})"
        
        return prediction, detailed_output
        
    except Exception as e:
        return f"Error making prediction: {str(e)}", f"**Error:** {str(e)}"

# Create the Gradio interface
with gradio.Blocks(title="Location Type Predictor", theme=gradio.themes.Soft()) as demo:
    # Header
    gradio.Markdown("""
    # 🌍 What type of place are you talking about?
    
    ## This AI model determines the location context of your description.
    
    Enter a place name and its description to get an AI-powered prediction of what type of location it is.
    """)
    
    with gradio.Row():
        with gradio.Column():
            place_name = gradio.Textbox(
                label="Place Name", 
                placeholder="Enter the name of the place (e.g., 'Central Library')",
                lines=1
            )
            description = gradio.Textbox(
                label="Description", 
                placeholder="Describe the place in detail...",
                lines=4
            )
            
            # Inference parameters
            gradio.Markdown("### 🔧 Inference Parameters")
            with gradio.Row():
                confidence_threshold = gradio.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.1,
                    step=0.05,
                    label="Confidence Threshold",
                    info="Minimum confidence to display probability scores"
                )
                return_probabilities = gradio.Checkbox(
                    label="Show Prediction Probabilities",
                    value=True,
                    info="Display confidence scores for top predictions"
                )
            
            predict_btn = gradio.Button("🎯 Predict Location Type", variant="primary", size="lg")
        
        with gradio.Column():
            # Outputs
            prediction_output = gradio.Textbox(
                label="Quick Result", 
                interactive=False,
                lines=1
            )
            detailed_output = gradio.Markdown(
                label="Detailed Analysis",
                value="Results will appear here..."
            )
    
    # Examples section
    gradio.Markdown("### 📝 Try These Examples")
    gradio.Examples(
        examples=examples,
        inputs=[place_name, description],
        label="Click on any example to try it",
        examples_per_page=5,
        cache_examples=False
    )
    
    # Connect the prediction function
    predict_btn.click(
        fn=predict_location,
        inputs=[place_name, description, confidence_threshold, return_probabilities],
        outputs=[prediction_output, detailed_output]
    )
    
demo.launch()