import gradio as gr
import json
from fastapi import FastAPI
import uvicorn
import pickle
import numpy as np
import pandas as pd
from typing import Dict, List, Optional
import os

app = FastAPI()

class DiabetesPredictor:
    def __init__(self, model_path: str = "diabetes_model.pkl", 
                 scaler_path: str = "scaler.pkl"):
        """
        Initialize the diabetes predictor with model and scaler.
        
        Args:
            model_path: Path to the trained model .pkl file
            scaler_path: Path to the scaler .pkl file
        """
        self.model = None
        self.scaler = None
        self.feature_names = None
        
        # Try to load the model
        try:
            if os.path.exists(model_path):
                with open(model_path, 'rb') as f:
                    self.model = pickle.load(f)
                print(f"✓ Model loaded successfully from {model_path}")
            else:
                print(f"⚠ Warning: Model file not found at {model_path}")
        except Exception as e:
            print(f"✗ Error loading model: {e}")
        
        # Try to load the scaler
        try:
            if os.path.exists(scaler_path):
                with open(scaler_path, 'rb') as f:
                    self.scaler = pickle.load(f)
                print(f"✓ Scaler loaded successfully from {scaler_path}")
            else:
                print(f"⚠ Warning: Scaler file not found at {scaler_path}")
        except Exception as e:
            print(f"✗ Error loading scaler: {e}")
    
    def prepare_features(self, data: Dict) -> np.ndarray:
        """
        Prepare input features for prediction.
        Expected features in order:
        - Pregnancies
        - Glucose
        - BloodPressure
        - SkinThickness
        - Insulin
        - BMI
        - DiabetesPedigreeFunction
        - Age
        """
        try:
            # Extract features from input data
            # If your model expects different features, modify this mapping
            features = [
                data.get("pregnancies", 0),          # Usually needed for diabetes prediction
                data.get("glucose", 100.0),
                data.get("blood_pressure", 120.0),
                data.get("skin_thickness", 20.0),    # Common diabetes dataset feature
                data.get("insulin", 15.0),
                data.get("bmi", 25.0),
                data.get("diabetes_pedigree", 0.5),  # Common diabetes dataset feature
                data.get("age", 30)
            ]
            
            return np.array(features).reshape(1, -1)
        except Exception as e:
            print(f"Error preparing features: {e}")
            return None
    
    def predict(self, data: Dict) -> Dict:
        """
        Make prediction using the loaded model.
        """
        if self.model is None:
            return {
                "success": False,
                "error": "Model not loaded. Using fallback prediction.",
                "fallback_used": True,
                "risk_score": self.fallback_prediction(data)
            }
        
        try:
            # Prepare features
            features = self.prepare_features(data)
            if features is None:
                raise ValueError("Could not prepare features")
            
            # Scale features if scaler is available
            if self.scaler is not None:
                features = self.scaler.transform(features)
            
            # Make prediction
            prediction = self.model.predict(features)[0]
            prediction_proba = self.model.predict_proba(features)[0]
            
            # Get probability for positive class (diabetes)
            # Assuming class 1 is diabetes
            risk_score = float(prediction_proba[1] * 100) if len(prediction_proba) > 1 else float(prediction * 100)
            
            is_high_risk = prediction == 1 or risk_score >= 50
            
            return {
                "success": True,
                "model_used": True,
                "prediction": int(prediction),
                "risk_score": risk_score,
                "is_high_risk": bool(is_high_risk),
                "risk_level": "High Risk" if is_high_risk else "Low Risk",
                "confidence": float(max(prediction_proba) * 100) if len(prediction_proba) > 1 else None,
                "message": self.get_recommendation(is_high_risk, risk_score)
            }
            
        except Exception as e:
            print(f"Prediction error: {e}")
            return {
                "success": False,
                "error": str(e),
                "fallback_used": True,
                "risk_score": self.fallback_prediction(data)
            }
    
    def fallback_prediction(self, data: Dict) -> float:
        """
        Fallback prediction logic when model fails to load.
        This is your original logic.
        """
        try:
            age = int(data.get("age", 30))
            bmi = float(data.get("bmi", 25.0))
            glucose = float(data.get("glucose", 100.0))
            
            score = 0
            if glucose > 140:
                score += 40
            if bmi > 30:
                score += 20
            if age > 45:
                score += 10
            
            # Add symptoms
            if data.get("increased_thirst"):
                score += 10
            if data.get("increased_hunger"):
                score += 5
            if data.get("fatigue"):
                score += 5
            if data.get("blurred_vision"):
                score += 10
            if data.get("weight_loss"):
                score += 15
            
            return min(score, 100)
        except:
            return 0.0
    
    def get_recommendation(self, is_high_risk: bool, risk_score: float) -> str:
        """Generate recommendation based on risk level."""
        if is_high_risk:
            if risk_score > 80:
                return "URGENT: Very high diabetes risk detected. Please consult a healthcare professional immediately."
            elif risk_score > 60:
                return "High diabetes risk detected. Schedule an appointment with your doctor soon."
            else:
                return "Moderate diabetes risk. Consider lifestyle changes and regular monitoring."
        else:
            if risk_score < 20:
                return "Low diabetes risk. Keep maintaining your healthy lifestyle!"
            else:
                return "Some risk factors present. Consider preventive measures and regular check-ups."

# Initialize predictor
predictor = DiabetesPredictor(
    model_path="diabetes_model.pkl",
    scaler_path="scaler.pkl"
)

def calculate_diabetes_risk_api(data: dict) -> dict:
    """API endpoint for diabetes risk prediction using ML model."""
    try:
        # Use the predictor
        result = predictor.predict(data)
        
        # If model prediction failed but we have fallback, format it
        if not result.get("success", False) and "fallback_used" in result:
            risk_score = result.get("risk_score", 0)
            is_high_risk = risk_score >= 50
            
            return {
                "success": True,
                "model_used": False,
                "fallback_used": True,
                "risk_score": risk_score,
                "is_high_risk": is_high_risk,
                "risk_level": "High Risk" if is_high_risk else "Low Risk",
                "message": predictor.get_recommendation(is_high_risk, risk_score)
            }
        
        return result
        
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }

# Create a comprehensive Gradio interface
with gr.Blocks(
    title="GlucoCheck AI - Diabetes Prediction",
    css="""
    .gradio-container {
        max-width: 900px;
        margin: auto;
    }
    .header {
        text-align: center;
        margin-bottom: 30px;
    }
    .header h1 {
        color: #2E384D;
        font-size: 36px;
        margin-bottom: 10px;
    }
    .header p {
        color: #6B7280;
        font-size: 16px;
    }
    .metric-card {
        background: linear-gradient(135deg, #f8fafc, #f1f5f9);
        padding: 15px;
        border-radius: 10px;
        border: 1px solid #e2e8f0;
        margin-bottom: 10px;
    }
    .vital-metric {
        background: linear-gradient(135deg, #fef2f2, #fef7ed);
        padding: 20px;
        border-radius: 12px;
        border: 2px solid #fecaca;
        margin-bottom: 15px;
    }
    .result-high-risk {
        background: linear-gradient(135deg, #fef2f2, #fee2e2);
        border-left: 5px solid #EF4444;
        padding: 20px;
        border-radius: 10px;
        margin: 15px 0;
    }
    .result-low-risk {
        background: linear-gradient(135deg, #f0fdf4, #dcfce7);
        border-left: 5px solid #10B981;
        padding: 20px;
        border-radius: 10px;
        margin: 15px 0;
    }
    .analyze-btn {
        background: linear-gradient(135deg, #4361ee, #3a56d4);
        color: white;
        padding: 15px 30px;
        border-radius: 12px;
        font-weight: 600;
        font-size: 16px;
        border: none;
        margin-top: 20px;
        width: 100%;
    }
    .analyze-btn:hover {
        background: linear-gradient(135deg, #3a56d4, #304bc0);
    }
    .disclaimer {
        margin-top: 30px;
        padding-top: 20px;
        border-top: 1px solid #e2e8f0;
        color: #6B7280;
        font-size: 12px;
        text-align: center;
    }
    .model-status {
        padding: 10px;
        border-radius: 8px;
        margin: 10px 0;
        text-align: center;
    }
    .model-success {
        background: #10B98120;
        color: #10B981;
        border: 1px solid #10B981;
    }
    .model-warning {
        background: #F59E0B20;
        color: #F59E0B;
        border: 1px solid #F59E0B;
    }
    """
) as demo:
    
    # Header
    gr.HTML("""
    <div class="header">
        <h1> GlucoCheck AI - Diabetes Prediction</h1>
        <p>Advanced ML-based diabetes risk assessment using trained models</p>
    </div>
    """)
    
    # Model status display
    model_status = gr.HTML("""
    <div class="model-status model-success">
        <strong>✓ ML Model Status:</strong> Ready for predictions
    </div>
    """) if predictor.model is not None else gr.HTML("""
    <div class="model-status model-warning">
        <strong>⚠ ML Model Status:</strong> Using fallback prediction logic
    </div>
    """)
    
    gr.Markdown("### Enter Patient Information")
    
    # Input fields in two columns
    with gr.Row():
        with gr.Column():
            age = gr.Number(
                label="Age (Years)",
                value=30,
                minimum=0,
                maximum=120,
                step=1,
                elem_classes="metric-card"
            )
            
            bmi = gr.Number(
                label="BMI (kg/m²)",
                value=25.0,
                minimum=10,
                maximum=60,
                step=0.1,
                elem_classes="metric-card"
            )
            
            pregnancies = gr.Number(
                label="Number of Pregnancies",
                value=0,
                minimum=0,
                maximum=20,
                step=1,
                elem_classes="metric-card"
            )
        
        with gr.Column():
            glucose = gr.Number(
                label="Glucose Level (mg/dL)",
                value=100.0,
                minimum=50,
                maximum=300,
                step=1.0,
                elem_classes="vital-metric"
            )
            
            blood_pressure = gr.Number(
                label="Blood Pressure (mm Hg)",
                value=120.0,
                minimum=60,
                maximum=200,
                step=1.0,
                elem_classes="metric-card"
            )
            
            insulin = gr.Number(
                label="Insulin Level (mu U/ml)",
                value=15.0,
                minimum=0,
                maximum=100,
                step=0.1,
                elem_classes="metric-card"
            )
    
    # Additional features that might be in your model
    with gr.Row():
        with gr.Column():
            skin_thickness = gr.Number(
                label="Skin Thickness (mm)",
                value=20.0,
                minimum=0,
                maximum=100,
                step=0.1,
                elem_classes="metric-card"
            )
        
        with gr.Column():
            diabetes_pedigree = gr.Number(
                label="Diabetes Pedigree Function",
                value=0.5,
                minimum=0,
                maximum=2.5,
                step=0.01,
                elem_classes="metric-card"
            )
    
    # Symptoms section
    gr.Markdown("### Symptoms")
    with gr.Row():
        increased_thirst = gr.Checkbox(label="Increased Thirst")
        increased_hunger = gr.Checkbox(label="Increased Hunger")
        fatigue = gr.Checkbox(label="Fatigue")
    
    with gr.Row():
        blurred_vision = gr.Checkbox(label="Blurred Vision")
        weight_loss = gr.Checkbox(label="Weight Loss")
    
    # Predict button
    predict_btn = gr.Button(" Analyze Diabetes Risk", variant="primary", elem_classes="analyze-btn")
    
    # Output sections
    gr.Markdown("### Prediction Results")
    
    with gr.Row():
        with gr.Column():
            risk_score_output = gr.Number(label="Risk Score (%)", interactive=False)
            risk_level_output = gr.Textbox(label="Risk Level", interactive=False)
            model_used_output = gr.Textbox(label="Prediction Method", interactive=False)
        
        with gr.Column():
            result_output = gr.HTML(label="Detailed Analysis")
    
    # Recommendations output
    recommendations_output = gr.Textbox(
        label="Recommendations",
        interactive=False,
        lines=4
    )
    
    # Raw JSON output for debugging/API
    json_output = gr.JSON(label="Raw API Response")
    
    # Prediction function
    def predict_risk(
        age_val, bmi_val, glucose_val, bp_val, insulin_val,
        pregnancies_val, skin_val, pedigree_val,
        thirst, hunger, fatigue_val, vision, weight
    ):
        # Prepare data dictionary
        data = {
            "age": age_val,
            "bmi": bmi_val,
            "glucose": glucose_val,
            "blood_pressure": bp_val,
            "insulin": insulin_val,
            "pregnancies": pregnancies_val,
            "skin_thickness": skin_val,
            "diabetes_pedigree": pedigree_val,
            "increased_thirst": thirst,
            "increased_hunger": hunger,
            "fatigue": fatigue_val,
            "blurred_vision": vision,
            "weight_loss": weight
        }
        
        # Get prediction
        result = calculate_diabetes_risk_api(data)
        
        # Prepare outputs
        if result.get("success", False):
            risk_score = result.get("risk_score", 0)
            risk_level = result.get("risk_level", "Unknown")
            model_used = "ML Model" if result.get("model_used", False) else "Fallback Logic"
            message = result.get("message", "")
            
            # Create HTML result display
            if result.get("is_high_risk", False):
                result_html = f"""
                <div class="result-high-risk">
                    <h3 style="color: #EF4444; margin-top: 0;"> HIGH RISK DETECTED</h3>
                    <p><strong>Risk Score:</strong> {risk_score:.1f}%</p>
                    <p><strong>Confidence:</strong> {result.get('confidence', 'N/A')}%</p>
                    <p><strong>Prediction:</strong> Diabetes likely present</p>
                </div>
                """
            else:
                result_html = f"""
                <div class="result-low-risk">
                    <h3 style="color: #10B981; margin-top: 0;"> LOW RISK</h3>
                    <p><strong>Risk Score:</strong> {risk_score:.1f}%</p>
                    <p><strong>Confidence:</strong> {result.get('confidence', 'N/A')}%</p>
                    <p><strong>Prediction:</strong> Diabetes unlikely</p>
                </div>
                """
            
            return {
                risk_score_output: risk_score,
                risk_level_output: risk_level,
                model_used_output: model_used,
                result_output: result_html,
                recommendations_output: message,
                json_output: result
            }
        else:
            error_html = f"""
            <div style="
                background: #FEF2F2;
                border-left: 5px solid #EF4444;
                padding: 20px;
                border-radius: 10px;
                margin: 15px 0;
            ">
                <h3 style="color: #EF4444; margin-top: 0;"> Error</h3>
                <p>{result.get('error', 'Unknown error occurred')}</p>
            </div>
            """
            
            return {
                risk_score_output: 0,
                risk_level_output: "Error",
                model_used_output: "Error",
                result_output: error_html,
                recommendations_output: "Please check your inputs and try again.",
                json_output: result
            }
    
    # Connect predict button
    predict_btn.click(
        predict_risk,
        inputs=[
            age, bmi, glucose, blood_pressure, insulin,
            pregnancies, skin_thickness, diabetes_pedigree,
            increased_thirst, increased_hunger, fatigue, blurred_vision, weight_loss
        ],
        outputs=[
            risk_score_output, risk_level_output, model_used_output,
            result_output, recommendations_output, json_output
        ]
    )
    
    # API documentation
    gr.Markdown("### API Usage")
    gr.Markdown("""
    You can also use this as an API endpoint:
    
    ```bash
    curl -X POST https://your-space.hf.space/api/predict \\
      -H "Content-Type: application/json" \\
      -d '{
        "age": 45,
        "bmi": 28.5,
        "glucose": 150,
        "blood_pressure": 130,
        "insulin": 20,
        "pregnancies": 0,
        "skin_thickness": 25,
        "diabetes_pedigree": 0.6,
        "increased_thirst": true,
        "increased_hunger": false,
        "fatigue": true,
        "blurred_vision": false,
        "weight_loss": true
      }'
    ```
    """)
    
    # Footer
    gr.HTML("""
    <div class="disclaimer">
        <p><strong> Medical Disclaimer:</strong> This tool is for informational purposes only and is not a substitute for professional medical advice, diagnosis, or treatment.</p>
        <p>Always seek the advice of your physician or other qualified health provider with any questions you may have regarding a medical condition.</p>
        <p>Model file: {'Loaded' if predictor.model is not None else 'Not found'}</p>
    </div>
    """)

# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

# For Hugging Face Spaces
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)