import gradio as gr
import joblib
import numpy as np
import pandas as pd
from huggingface_hub import hf_hub_download

# Load model and preprocessing pipeline
def load_model_and_preprocessor():
    print("Loading model and preprocessor...")
    
    # Download files from HuggingFace
    model_path = hf_hub_download(
        repo_id="EricCRX/grape-firmness-automl",
        filename="model.joblib"
    )
    
    preprocess_path = hf_hub_download(
        repo_id="EricCRX/grape-firmness-automl",
        filename="preprocess.joblib"
    )
    
    # Load the model and preprocessor
    model = joblib.load(model_path)
    preprocessor = joblib.load(preprocess_path)
    
    print("Model and preprocessor loaded successfully!")
    return model, preprocessor

# Load at startup
model, preprocessor = load_model_and_preprocessor()

# Firmness level mapping
FIRMNESS_LEVELS = {
    1: "Soft",
    2: "Medium",
    3: "Firm"
}

def predict_firmness(width_cm, length_cm, blemish_count, color):
    """
    Predict grape firmness based on physical characteristics
    
    Args:
        width_cm: Grape width in centimeters
        length_cm: Grape length in centimeters
        blemish_count: Number of visible blemishes
        color: Grape color (red or green)
    
    Returns:
        Prediction text with firmness level and confidence
    """
    try:
        # Create input dataframe
        input_data = pd.DataFrame([{
            'width_cm': float(width_cm),
            'length_cm': float(length_cm),
            'blemish_count': int(blemish_count),
            'color': str(color)
        }])
        
        # Apply preprocessing
        X_preprocessed = preprocessor.transform(input_data)
        
        # Make prediction
        prediction = model.predict(X_preprocessed)[0]
        
        # Round to nearest firmness level (1, 2, or 3)
        firmness_level = int(round(prediction))
        firmness_level = max(1, min(3, firmness_level))  # Clamp between 1 and 3
        
        firmness_label = FIRMNESS_LEVELS[firmness_level]
        
        # Create result message
        result = f"""
## Prediction Result

### Predicted Firmness: **{firmness_label}** (Level {firmness_level})

**Raw Prediction Value:** {prediction:.3f}

### Interpretation:
- **Level 1 (Soft)**: Grape is very soft, may be overripe
- **Level 2 (Medium)**: Grape has moderate firmness, good for eating
- **Level 3 (Firm)**: Grape is very firm, fresh and crisp

### Input Characteristics:
- Width: {width_cm} cm
- Length: {length_cm} cm
- Blemishes: {blemish_count}
- Color: {color.capitalize()}

### Quality Assessment:
"""
        
        # Add quality assessment based on characteristics
        if blemish_count == 0:
            result += "- ✅ No blemishes detected - excellent visual quality\n"
        elif blemish_count <= 2:
            result += f"- ⚠️ {blemish_count} blemish(es) detected - acceptable quality\n"
        else:
            result += f"- ❌ {blemish_count} blemishes detected - lower quality\n"
        
        # Size assessment
        avg_size = (width_cm + length_cm) / 2
        if avg_size > 2.4:
            result += "- 📏 Large grape size\n"
        elif avg_size > 2.1:
            result += "- 📏 Medium grape size\n"
        else:
            result += "- 📏 Small grape size\n"
        
        return result
        
    except Exception as e:
        return f"**Error making prediction:**\n\n{str(e)}\n\nPlease check your inputs and try again."

# Create Gradio interface
with gr.Blocks(title="Grape Firmness Predictor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🍇 Grape Firmness Predictor
    
    Predict grape firmness using a **RandomForest Regressor** trained on physical characteristics.
    
    This model predicts firmness on a 3-level scale:
    - **1 = Soft** (overripe, mushy)
    - **2 = Medium** (optimal ripeness)
    - **3 = Firm** (fresh, crisp)
    
    ### Model Performance
    - **RMSE:** 0.177
    - **R² Score:** 0.943
    - **Framework:** scikit-learn RandomForest
    """)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🔍 Grape Characteristics")
            
            width_input = gr.Slider(
                minimum=1.5,
                maximum=3.0,
                value=2.2,
                step=0.1,
                label="Width (cm)",
                info="Grape width in centimeters"
            )
            
            length_input = gr.Slider(
                minimum=1.5,
                maximum=3.0,
                value=2.5,
                step=0.1,
                label="Length (cm)",
                info="Grape length in centimeters"
            )
            
            blemish_input = gr.Slider(
                minimum=0,
                maximum=10,
                value=2,
                step=1,
                label="Blemish Count",
                info="Number of visible blemishes on grape surface"
            )
            
            color_input = gr.Radio(
                choices=["red", "green"],
                value="red",
                label="Grape Color",
                info="Select the grape variety color"
            )
            
            predict_btn = gr.Button("🔮 Predict Firmness", variant="primary", size="lg")
        
        with gr.Column():
            gr.Markdown("### 📊 Prediction Results")
            output = gr.Markdown()
    
    # Example inputs
    gr.Markdown("### 🎯 Try These Examples")
    gr.Markdown("*Click an example to load it*")
    
    examples = [
        [2.3, 2.5, 4, "red"],      # Typical soft red grape
        [2.2, 2.4, 0, "green"],    # Firm green grape with no blemishes
        [2.5, 2.8, 1, "red"],      # Large firm red grape
        [2.0, 2.3, 5, "red"],      # Small soft red grape with blemishes
        [2.4, 2.7, 2, "green"],    # Medium green grape
    ]
    
    gr.Examples(
        examples=examples,
        inputs=[width_input, length_input, blemish_input, color_input],
        outputs=output,
        fn=predict_firmness,
        cache_examples=False
    )
    
    # Information sections
    with gr.Accordion("📖 About the Model", open=False):
        gr.Markdown("""
        ### Model Architecture
        - **Algorithm**: RandomForest Regressor
        - **Training**: AutoML grid search across multiple models
        - **Preprocessing**: Custom feature engineering and scaling pipeline
        - **Features**: width_cm, length_cm, blemish_count, color
        
        ### Dataset Information
        - **Source**: [rlogh/grape-firmness-dataset](https://huggingface.co/datasets/rlogh/grape-firmness-dataset)
        - **Size**: 300 samples (augmented from 30 original samples)
        - **Augmentation**: SMOTE-NC and numeric jitter for class balance
        - **Splits**: 70% train, 15% validation, 15% test
        
        ### Feature Descriptions
        - **width_cm**: Physical width measurement in centimeters
        - **length_cm**: Physical length measurement in centimeters
        - **blemish_count**: Number of visible surface imperfections
        - **color**: Grape variety color (red or green)
        
        ### Target Variable
        - **firmness**: Integer scale from 1-3
          - 1 = Soft (may be overripe)
          - 2 = Medium (optimal eating ripeness)
          - 3 = Firm (fresh, good for storage)
        """)
    
    with gr.Accordion("⚙️ Model Training Details", open=False):
        gr.Markdown("""
        ### AutoML Search Space
        The model was selected through automated hyperparameter search:
        
        **RandomForest Regressor** (Selected):
        - `n_estimators`: Number of trees in forest
        - `max_depth`: Maximum tree depth
        
        **Support Vector Regressor** (Alternative tested):
        - `C`: Regularization parameter
        - `kernel`: Kernel type (linear, rbf, poly)
        
        ### Evaluation Metric
        - **Primary**: R² score (coefficient of determination)
        - **Secondary**: RMSE (Root Mean Squared Error)
        
        ### Validation Strategy
        - Stratified split for regression
        - Cross-validation during grid search
        - Held-out test set for final evaluation
        """)
    
    with gr.Accordion("⚠️ Limitations & Disclaimers", open=False):
        gr.Markdown("""
        ### Known Limitations
        1. **Small Dataset**: Only 30 original samples, augmented to 300
           - May not generalize to unseen grape varieties
           - Limited geographic/environmental diversity
        
        2. **Feature Engineering**: Minimal domain knowledge incorporated
           - No temperature or humidity data
           - No time-since-harvest information
           - No chemical composition data
        
        3. **Measurement Precision**: Assumes consistent measurement methods
           - Results depend on accurate physical measurements
           - Blemish counting may be subjective
        
        4. **Not Production-Ready**: Academic demonstration only
           - Should not be used for commercial grading
           - Requires validation on real-world data
        
        ### When Predictions May Be Unreliable
        - Grapes outside the training size range (1.5-3.0 cm)
        - Unusual grape varieties not in training data
        - Grapes with surface treatments or coatings
        - Very high blemish counts (>6)
        
        ### Recommended Use
        - Educational demonstrations
        - Proof-of-concept for automated quality assessment
        - Baseline model for further development
        """)
    
    with gr.Accordion("💡 How to Interpret Results", open=False):
        gr.Markdown("""
        ### Understanding Firmness Levels
        
        **Level 1 - Soft**:
        - Grape yields easily to pressure
        - May have wrinkled skin
        - Suitable for immediate consumption or making juice
        - Lower shelf life
        
        **Level 2 - Medium**:
        - Balanced firmness
        - Optimal for eating fresh
        - Good texture and mouthfeel
        - Moderate shelf life
        
        **Level 3 - Firm**:
        - Crisp and crunchy texture
        - Fresh from vine or cold storage
        - Excellent for long-term storage
        - May be slightly tart
        
        ### Factors Affecting Firmness
        1. **Ripeness**: As grapes ripen, they soften
        2. **Water Content**: Higher water = softer texture
        3. **Storage**: Temperature and humidity affect firmness
        4. **Variety**: Some varieties naturally firmer than others
        5. **Blemishes**: Damage can create soft spots
        """)
    
    # Connect button
    predict_btn.click(
        fn=predict_firmness,
        inputs=[width_input, length_input, blemish_input, color_input],
        outputs=output
    )

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