# app.py
import streamlit as st
import pandas as pd
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
import plotly.express as px
import plotly.graph_objects as go
import seaborn as sns
import matplotlib.pyplot as plt

class IrisPredictor:
    def __init__(self):
        self.iris = load_iris()
        self.df = pd.DataFrame(data=np.c_[self.iris['data'], self.iris['target']], 
                             columns=self.iris['feature_names'] + ['target'])
        self.models = {
            'Logistic Regression': LogisticRegression(),
            'Decision Tree': DecisionTreeClassifier(),
            'Random Forest': RandomForestClassifier()
        }
        self.X = self.df.drop('target', axis=1)
        self.y = self.df['target']
        self.scaler = StandardScaler()
        
    def preprocess_data(self):
        # Split the data
        X_train, X_test, y_train, y_test = train_test_split(
            self.X, self.y, test_size=0.2, random_state=42
        )
        
        # Scale the features
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        return X_train_scaled, X_test_scaled, y_train, y_test
    
    def train_model(self, model_name):
        X_train_scaled, X_test_scaled, y_train, y_test = self.preprocess_data()
        
        # Train model
        model = self.models[model_name]
        model.fit(X_train_scaled, y_train)
        
        # Make predictions
        y_pred = model.predict(X_test_scaled)
        
        # Calculate metrics
        accuracy = accuracy_score(y_test, y_pred)
        conf_matrix = confusion_matrix(y_test, y_pred)
        class_report = classification_report(y_test, y_pred)
        
        return {
            'model': model,
            'accuracy': accuracy,
            'confusion_matrix': conf_matrix,
            'classification_report': class_report,
            'X_test': X_test_scaled,
            'y_test': y_test,
            'y_pred': y_pred
        }
    
    def plot_confusion_matrix(self, conf_matrix):
        fig = px.imshow(conf_matrix,
                       labels=dict(x="Predicted", y="Actual"),
                       x=['Setosa', 'Versicolor', 'Virginica'],
                       y=['Setosa', 'Versicolor', 'Virginica'],
                       title="Confusion Matrix")
        return fig
    
    def plot_feature_importance(self, model_name, model):
        if model_name == 'Logistic Regression':
            importance = abs(model.coef_[0])
        else:
            importance = model.feature_importances_
            
        fig = px.bar(x=self.X.columns, y=importance,
                    title=f"Feature Importance - {model_name}",
                    labels={'x': 'Features', 'y': 'Importance'})
        return fig
    
    def predict_single_sample(self, model, features):
        # Scale features
        scaled_features = self.scaler.transform([features])
        # Make prediction
        prediction = model.predict(scaled_features)
        probabilities = model.predict_proba(scaled_features)
        return prediction[0], probabilities[0]

def main():
    st.title("🌸 Iris Flower Prediction App")
    st.write("""
    This app predicts the Iris flower type based on its features.
    Choose a model and see how it performs!
    """)
    
    # Initialize predictor
    predictor = IrisPredictor()
    
    # Model selection
    st.sidebar.header("Model Selection")
    model_name = st.sidebar.selectbox(
        "Choose a model",
        list(predictor.models.keys())
    )
    
    # Train model and show results
    if st.sidebar.button("Train Model"):
        with st.spinner("Training model..."):
            results = predictor.train_model(model_name)
            
            # Display metrics
            st.header("Model Performance")
            st.metric("Accuracy", f"{results['accuracy']:.2%}")
            
            # Display confusion matrix
            st.subheader("Confusion Matrix")
            conf_matrix_fig = predictor.plot_confusion_matrix(results['confusion_matrix'])
            st.plotly_chart(conf_matrix_fig)
            
            # Display feature importance
            st.subheader("Feature Importance")
            importance_fig = predictor.plot_feature_importance(model_name, results['model'])
            st.plotly_chart(importance_fig)
            
            # Display classification report
            st.subheader("Classification Report")
            st.text(results['classification_report'])
            
            # Store the trained model in session state
            st.session_state['trained_model'] = results['model']
    
    # Make predictions
    st.header("Make Predictions")
    col1, col2 = st.columns(2)
    
    with col1:
        sepal_length = st.slider("Sepal Length", 4.0, 8.0, 5.4)
        sepal_width = st.slider("Sepal Width", 2.0, 4.5, 3.4)
    
    with col2:
        petal_length = st.slider("Petal Length", 1.0, 7.0, 4.7)
        petal_width = st.slider("Petal Width", 0.1, 2.5, 1.4)
    
    if st.button("Predict"):
        if 'trained_model' in st.session_state:
            features = [sepal_length, sepal_width, petal_length, petal_width]
            prediction, probabilities = predictor.predict_single_sample(
                st.session_state['trained_model'], features
            )
            
            # Display prediction
            iris_types = ['Setosa', 'Versicolor', 'Virginica']
            st.success(f"Predicted Iris Type: {iris_types[int(prediction)]}")
            
            # Display probability distribution
            st.subheader("Prediction Probabilities")
            prob_fig = px.bar(x=iris_types, y=probabilities,
                            title="Prediction Probability Distribution",
                            labels={'x': 'Iris Type', 'y': 'Probability'})
            st.plotly_chart(prob_fig)
        else:
            st.warning("Please train a model first!")

if __name__ == "__main__":
    main()