import gradio as gr
import pandas as pd
import numpy as np
import pickle
import json
import tensorflow as tf
from tensorflow.keras.models import load_model
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import os

# Load model artifacts
def load_model_artifacts():
    try:
        # Load the trained model
        model = load_model('final_model.h5')
        
        # Load the scaler
        with open('scaler.pkl', 'rb') as f:
            scaler = pickle.load(f)
        
        # Load metadata
        with open('metadata.json', 'r') as f:
            metadata = json.load(f)
            
        return model, scaler, metadata
    except Exception as e:
        raise Exception(f"Error loading model artifacts: {str(e)}")

# Initialize model components
try:
    model, scaler, metadata = load_model_artifacts()
    feature_names = metadata['feature_names']
    print(f"✅ Model loaded successfully with features: {feature_names}")
except Exception as e:
    print(f"❌ Error loading model: {e}")
    # Fallback values for testing
    model, scaler, metadata = None, None, {}
    feature_names = ['Feature_1', 'Feature_2', 'Feature_3', 'Feature_4']

def predict_student_eligibility(*args):
    """
    Predict student eligibility based on input features
    """
    try:
        if model is None or scaler is None:
            return "Model not loaded", "N/A", "N/A", create_error_plot()
        
        # Create input dictionary from gradio inputs
        input_data = {feature_names[i]: args[i] for i in range(len(feature_names))}
        
        # Convert to DataFrame
        input_df = pd.DataFrame([input_data])
        
        # Scale the input
        input_scaled = scaler.transform(input_df)
        
        # Reshape for CNN
        input_reshaped = input_scaled.reshape(input_scaled.shape[0], input_scaled.shape[1], 1)
        
        # Make prediction
        probability = float(model.predict(input_reshaped)[0][0])
        prediction = "Eligible" if probability > 0.5 else "Not Eligible"
        confidence = abs(probability - 0.5) * 2  # Convert to confidence score
        
        # Create prediction visualization
        fig = create_prediction_viz(probability, prediction, input_data)
        
        return prediction, f"{probability:.4f}", f"{confidence:.4f}", fig
        
    except Exception as e:
        return f"Error: {str(e)}", "N/A", "N/A", create_error_plot()

def create_error_plot():
    """Create a simple error plot"""
    fig = go.Figure()
    fig.add_annotation(
        text="Model not available or error occurred",
        xref="paper", yref="paper",
        x=0.5, y=0.5, xanchor='center', yanchor='middle',
        showarrow=False, font=dict(size=20)
    )
    fig.update_layout(
        xaxis={'visible': False},
        yaxis={'visible': False},
        height=400
    )
    return fig

def create_prediction_viz(probability, prediction, input_data):
    """
    Create visualization for prediction results
    """
    try:
        # Create subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('Prediction Probability', 'Confidence Meter', 'Input Features', 'Probability Distribution'),
            specs=[[{"type": "indicator"}, {"type": "indicator"}],
                   [{"type": "bar"}, {"type": "scatter"}]]
        )
        
        # Prediction probability gauge
        fig.add_trace(
            go.Indicator(
                mode="gauge+number",
                value=probability,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': "Eligibility Probability"},
                gauge={
                    'axis': {'range': [None, 1]},
                    'bar': {'color': "darkblue"},
                    'steps': [
                        {'range': [0, 0.5], 'color': "lightcoral"},
                        {'range': [0.5, 1], 'color': "lightgreen"}
                    ],
                    'threshold': {
                        'line': {'color': "red", 'width': 4},
                        'thickness': 0.75,
                        'value': 0.5
                    }
                }
            ),
            row=1, col=1
        )
        
        # Confidence meter
        confidence = abs(probability - 0.5) * 2
        fig.add_trace(
            go.Indicator(
                mode="gauge+number",
                value=confidence,
                domain={'x': [0, 1], 'y': [0, 1]},
                title={'text': "Prediction Confidence"},
                gauge={
                    'axis': {'range': [None, 1]},
                    'bar': {'color': "orange"},
                    'steps': [
                        {'range': [0, 0.3], 'color': "lightcoral"},
                        {'range': [0.3, 0.7], 'color': "lightyellow"},
                        {'range': [0.7, 1], 'color': "lightgreen"}
                    ]
                }
            ),
            row=1, col=2
        )
        
        # Input features bar chart
        features = list(input_data.keys())
        values = list(input_data.values())
        
        fig.add_trace(
            go.Bar(x=features, y=values, name="Input Values", marker_color="skyblue"),
            row=2, col=1
        )
        
        # Simple probability visualization
        fig.add_trace(
            go.Scatter(
                x=[0, 1], 
                y=[probability, probability], 
                mode='lines+markers',
                name="Probability", 
                line=dict(color="red", width=3),
                marker=dict(size=10)
            ),
            row=2, col=2
        )
        
        fig.update_layout(
            height=800,
            showlegend=False,
            title_text="Student Eligibility Prediction Dashboard",
            title_x=0.5
        )
        
        return fig
    except Exception as e:
        return create_error_plot()

def create_model_info():
    """
    Create model information display
    """
    if metadata:
        info_html = f"""
        <div style="padding: 20px; background-color: #f0f2f6; border-radius: 10px; margin: 10px 0;">
            <h3>🤖 Model Information</h3>
            <ul>
                <li><strong>Model Type:</strong> {metadata.get('model_type', 'CNN')}</li>
                <li><strong>Test Accuracy:</strong> {metadata.get('performance_metrics', {}).get('test_accuracy', 'N/A')}</li>
                <li><strong>AUC Score:</strong> {metadata.get('performance_metrics', {}).get('auc_score', 'N/A')}</li>
                <li><strong>Creation Date:</strong> {metadata.get('creation_date', 'N/A')}</li>
                <li><strong>Features:</strong> {len(feature_names)} input features</li>
            </ul>
        </div>
        """
    else:
        info_html = """
        <div style="padding: 20px; background-color: #ffebee; border-radius: 10px; margin: 10px 0;">
            <h3>⚠️ Model Information</h3>
            <p>Model artifacts not loaded. Please ensure all required files are uploaded.</p>
        </div>
        """
    return info_html

def batch_predict(file):
    """
    Batch prediction from uploaded CSV file
    """
    try:
        if model is None or scaler is None:
            return "Model not loaded. Please check if all model files are uploaded.", None
        
        if file is None:
            return "Please upload a CSV file.", None
        
        # Read the uploaded file
        df = pd.read_csv(file.name)
        
        # Check if all required features are present
        missing_features = set(feature_names) - set(df.columns)
        if missing_features:
            return f"Missing features: {missing_features}", None
        
        # Select only the required features
        df_features = df[feature_names]
        
        # Scale the features
        df_scaled = scaler.transform(df_features)
        
        # Reshape for CNN
        df_reshaped = df_scaled.reshape(df_scaled.shape[0], df_scaled.shape[1], 1)
        
        # Make predictions
        probabilities = model.predict(df_reshaped).flatten()
        predictions = ["Eligible" if p > 0.5 else "Not Eligible" for p in probabilities]
        
        # Create results dataframe
        results_df = df_features.copy()
        results_df['Probability'] = probabilities
        results_df['Prediction'] = predictions
        results_df['Confidence'] = np.abs(probabilities - 0.5) * 2
        
        # Save results
        output_file = "batch_predictions.csv"
        results_df.to_csv(output_file, index=False)
        
        # Create summary statistics
        eligible_count = sum(1 for p in predictions if p == 'Eligible')
        not_eligible_count = len(predictions) - eligible_count
        
        summary = f"""Batch Prediction Summary:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
📊 Total predictions: {len(results_df)}
✅ Eligible: {eligible_count} ({eligible_count/len(predictions)*100:.1f}%)
❌ Not Eligible: {not_eligible_count} ({not_eligible_count/len(predictions)*100:.1f}%)
📈 Average Probability: {np.mean(probabilities):.4f}
🎯 Average Confidence: {np.mean(np.abs(probabilities - 0.5) * 2):.4f}
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Results saved to: {output_file}
        """
        
        return summary, output_file
        
    except Exception as e:
        return f"Error processing file: {str(e)}", None

# Create Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Student Eligibility Prediction",
    css="""
    .gradio-container {
        max-width: 1200px !important;
    }
    .main-header {
        text-align: center;
        padding: 20px;
        background: linear-gradient(45deg, #667eea 0%, #764ba2 100%);
        color: white;
        border-radius: 10px;
        margin-bottom: 20px;
    }
    .feature-input {
        margin: 5px 0;
    }
    """
) as demo:
    
    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>🎓 Student Eligibility Prediction System</h1>
        <p>AI-powered CNN model for predicting student eligibility with advanced analytics</p>
    </div>
    """)
    
    with gr.Tabs():
        # Single Prediction Tab
        with gr.TabItem("🔮 Single Prediction"):
            gr.Markdown("### Enter student information to predict eligibility")
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("#### Input Features")
                    # Create input components dynamically based on features
                    inputs = []
                    for i, feature in enumerate(feature_names):
                        inputs.append(
                            gr.Number(
                                label=f"📊 {feature}",
                                value=75 + i*5,  # Different default values
                                minimum=0,
                                maximum=100,
                                step=0.1,
                                elem_classes=["feature-input"]
                            )
                        )
                    
                    predict_btn = gr.Button(
                        "🔮 Predict Eligibility", 
                        variant="primary", 
                        size="lg",
                        elem_id="predict-btn"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("#### Prediction Results")
                    with gr.Row():
                        prediction_output = gr.Textbox(label="🎯 Prediction", scale=1)
                        probability_output = gr.Textbox(label="📊 Probability", scale=1)
                        confidence_output = gr.Textbox(label="🎯 Confidence", scale=1)
                    
                    prediction_plot = gr.Plot(label="📈 Prediction Visualization")
            
            # Model information
            gr.HTML(create_model_info())
        
        # Batch Prediction Tab
        with gr.TabItem("📊 Batch Prediction"):
            gr.Markdown("### Upload a CSV file for batch predictions")
            gr.Markdown(f"**Required columns:** `{', '.join(feature_names)}`")
            
            # Sample CSV format
            gr.Markdown("""
            **Example CSV format:**
            ```csv
            Feature_1,Feature_2,Feature_3,Feature_4
            85,90,75,88
            92,78,85,91
            ```
            """)
            
            with gr.Row():
                with gr.Column():
                    file_input = gr.File(
                        label="📁 Upload CSV File",
                        file_types=[".csv"],
                        type="file"
                    )
                    batch_predict_btn = gr.Button(
                        "📊 Process Batch", 
                        variant="primary",
                        size="lg"
                    )
                
                with gr.Column():
                    batch_output = gr.Textbox(
                        label="📋 Batch Results Summary", 
                        lines=15,
                        max_lines=20
                    )
                    download_file = gr.File(label="⬇️ Download Results")
        
        # Model Analytics Tab
        with gr.TabItem("📈 Model Analytics"):
            gr.Markdown("### Model Performance Metrics")
            
            if metadata and 'performance_metrics' in metadata:
                # Performance metrics
                metrics_data = metadata['performance_metrics']
                metrics_df = pd.DataFrame([{
                    'Metric': k.replace('_', ' ').title(),
                    'Value': f"{v:.4f}" if isinstance(v, float) else str(v)
                } for k, v in metrics_data.items()])
                
                gr.Dataframe(
                    metrics_df, 
                    label="🎯 Performance Metrics",
                    headers=['Metric', 'Value']
                )
            else:
                gr.Markdown("⚠️ **Performance metrics not available**")
            
            # Feature information
            gr.Markdown("### 📊 Model Features")
            feature_info = pd.DataFrame({
                'Feature Name': feature_names,
                'Index': range(len(feature_names)),
                'Type': ['Numerical'] * len(feature_names)
            })
            gr.Dataframe(feature_info, label="Feature Information")
            
            # Model architecture info
            if metadata:
                gr.Markdown("### 🏗️ Model Architecture")
                arch_info = f"""
                - **Model Type**: {metadata.get('model_type', 'CNN')}
                - **Input Shape**: {metadata.get('input_shape', 'N/A')}
                - **Total Features**: {len(feature_names)}
                - **Output Classes**: {len(metadata.get('target_classes', {}))}
                """
                gr.Markdown(arch_info)
    
    # Event handlers
    predict_btn.click(
        fn=predict_student_eligibility,
        inputs=inputs,
        outputs=[prediction_output, probability_output, confidence_output, prediction_plot]
    )
    
    batch_predict_btn.click(
        fn=batch_predict,
        inputs=[file_input],
        outputs=[batch_output, download_file]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860
    )