"""
Dashboard Builder

Interactive dashboard builder for ML experiments and model monitoring
using Streamlit and Plotly for real-time visualization.
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import time
from typing import Dict, List, Any, Optional, Callable
import json
from datetime import datetime, timedelta
import sqlite3
import warnings


class DashboardBuilder:
    """
    Interactive dashboard builder for ML experiments and model monitoring
    with real-time updates and comprehensive visualizations.
    """
    
    def __init__(self, title: str = "ML Dashboard", layout: str = "wide"):
        """
        Initialize dashboard builder
        
        Args:
            title: Dashboard title
            layout: Streamlit layout ('wide' or 'centered')
        """
        self.title = title
        self.layout = layout
        self.data_sources = {}
        self.callbacks = {}
        self.filters = {}
        
    def setup_page(self):
        """Setup Streamlit page configuration"""
        st.set_page_config(
            page_title=self.title,
            page_icon="🤖",
            layout=self.layout,
            initial_sidebar_state="expanded"
        )
        
        # Custom CSS
        st.markdown("""
        <style>
        .main {
            padding-top: 2rem;
        }
        .stAlert {
            margin-top: 1rem;
        }
        </style>
        """, unsafe_allow_html=True)
    
    def add_data_source(self, name: str, data_source: Any, description: str = ""):
        """
        Add a data source to the dashboard
        
        Args:
            name: Name of the data source
            data_source: Data (DataFrame, dict, etc.)
            description: Description of the data source
        """
        self.data_sources[name] = {
            'data': data_source,
            'description': description,
            'last_updated': datetime.now()
        }
    
    def add_filter(self, name: str, filter_type: str, options: List[Any] = None,
                  default: Any = None, description: str = ""):
        """
        Add a filter to the dashboard
        
        Args:
            name: Filter name
            filter_type: Type of filter ('selectbox', 'multiselect', 'slider', 'date_range')
            options: Options for select/multiselect filters
            default: Default value
            description: Filter description
        """
        self.filters[name] = {
            'type': filter_type,
            'options': options,
            'default': default,
            'description': description
        }
    
    def render_filters(self):
        """Render all filters in the sidebar"""
        st.sidebar.markdown("### Filters")
        
        filter_values = {}
        
        for name, filter_config in self.filters.items():
            if filter_config['description']:
                st.sidebar.markdown(f"**{filter_config['description']}**")
            
            if filter_config['type'] == 'selectbox':
                value = st.sidebar.selectbox(
                    name,
                    options=filter_config['options'],
                    index=filter_config['options'].index(filter_config['default']) if filter_config['default'] in filter_config['options'] else 0
                )
                filter_values[name] = value
                
            elif filter_config['type'] == 'multiselect':
                value = st.sidebar.multiselect(
                    name,
                    options=filter_config['options'],
                    default=filter_config['default'] or []
                )
                filter_values[name] = value
                
            elif filter_config['type'] == 'slider':
                value = st.sidebar.slider(
                    name,
                    min_value=filter_config['options'][0],
                    max_value=filter_config['options'][1],
                    value=filter_config['default'] or filter_config['options'][0]
                )
                filter_values[name] = value
                
            elif filter_config['type'] == 'date_range':
                value = st.sidebar.date_input(
                    name,
                    value=filter_config['default'] or (datetime.now() - timedelta(days=30), datetime.now())
                )
                filter_values[name] = value
            
            st.sidebar.markdown("---")
        
        return filter_values
    
    def create_experiment_dashboard(self, experiments_df: pd.DataFrame) -> None:
        """
        Create comprehensive experiment tracking dashboard
        
        Args:
            experiments_df: DataFrame containing experiment data
        """
        st.title("🧪 Experiment Tracking Dashboard")
        
        # Add filters
        self.add_filter("status", "multiselect", 
                       options=experiments_df['status'].unique().tolist(),
                       default=['completed'],
                       description="Experiment Status")
        
        self.add_filter("date_range", "date_range",
                       description="Date Range")
        
        # Render filters
        filters = self.render_filters()
        
        # Filter data
        filtered_df = experiments_df.copy()
        
        if 'status' in filters and filters['status']:
            filtered_df = filtered_df[filtered_df['status'].isin(filters['status'])]
        
        if 'date_range' in filters and filters['date_range']:
            start_date, end_date = filters['date_range']
            filtered_df['created_at'] = pd.to_datetime(filtered_df['created_at'])
            filtered_df = filtered_df[
                (filtered_df['created_at'].dt.date >= start_date) &
                (filtered_df['created_at'].dt.date <= end_date)
            ]
        
        # Key metrics
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("Total Experiments", len(filtered_df))
        
        with col2:
            completed = len(filtered_df[filtered_df['status'] == 'completed'])
            st.metric("Completed", completed)
        
        with col3:
            running = len(filtered_df[filtered_df['status'] == 'running'])
            st.metric("Running", running)
        
        with col4:
            failed = len(filtered_df[filtered_df['status'] == 'failed'])
            st.metric("Failed", failed)
        
        # Charts
        col1, col2 = st.columns(2)
        
        with col1:
            # Status distribution
            status_counts = filtered_df['status'].value_counts()
            fig = px.pie(values=status_counts.values, names=status_counts.index,
                        title="Experiment Status Distribution")
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            # Experiments over time
            if 'created_at' in filtered_df.columns:
                filtered_df['date'] = pd.to_datetime(filtered_df['created_at']).dt.date
                daily_counts = filtered_df.groupby('date').size().reset_index(name='count')
                fig = px.line(daily_counts, x='date', y='count',
                             title="Experiments Over Time")
                st.plotly_chart(fig, use_container_width=True)
        
        # Experiment details table
        st.subheader("Experiment Details")
        
        # Select columns to display
        display_columns = ['experiment_id', 'name', 'status', 'created_at', 'completed_at']
        available_columns = [col for col in display_columns if col in filtered_df.columns]
        
        if available_columns:
            st.dataframe(filtered_df[available_columns], use_container_width=True)
    
    def create_model_performance_dashboard(self, models_df: pd.DataFrame) -> None:
        """
        Create model performance comparison dashboard
        
        Args:
            models_df: DataFrame containing model performance data
        """
        st.title("📊 Model Performance Dashboard")
        
        # Add filters
        if 'model_type' in models_df.columns:
            self.add_filter("model_type", "multiselect",
                           options=models_df['model_type'].unique().tolist(),
                           description="Model Type")
        
        if 'problem_type' in models_df.columns:
            self.add_filter("problem_type", "selectbox",
                           options=models_df['problem_type'].unique().tolist(),
                           description="Problem Type")
        
        # Render filters
        filters = self.render_filters()
        
        # Filter data
        filtered_df = models_df.copy()
        
        if 'model_type' in filters and filters['model_type']:
            filtered_df = filtered_df[filtered_df['model_type'].isin(filters['model_type'])]
        
        if 'problem_type' in filters and filters['problem_type']:
            filtered_df = filtered_df[filtered_df['problem_type'] == filters['problem_type']]
        
        # Performance metrics
        metric_columns = [col for col in filtered_df.columns 
                         if col.endswith('_score') or col.endswith('_accuracy') 
                         or col.endswith('_loss') or col in ['accuracy', 'precision', 'recall', 'f1']]
        
        if metric_columns:
            # Model comparison chart
            st.subheader("Model Performance Comparison")
            
            # Select metric to compare
            selected_metric = st.selectbox("Select Metric", metric_columns)
            
            fig = px.bar(filtered_df, x='model_name', y=selected_metric,
                        color='model_type', barmode='group',
                        title=f"Model Comparison - {selected_metric}")
            st.plotly_chart(fig, use_container_width=True)
            
            # Performance table
            st.subheader("Performance Metrics Table")
            display_columns = ['model_name', 'model_type'] + metric_columns
            available_columns = [col for col in display_columns if col in filtered_df.columns]
            st.dataframe(filtered_df[available_columns], use_container_width=True)
    
    def create_training_monitoring_dashboard(self, training_data: Dict[str, Any]) -> None:
        """
        Create real-time training monitoring dashboard
        
        Args:
            training_data: Dictionary containing training history and metrics
        """
        st.title("🎯 Training Monitoring Dashboard")
        
        # Auto-refresh
        auto_refresh = st.checkbox("Auto Refresh (5s)", value=True)
        
        if auto_refresh:
            time.sleep(5)
            st.experimental_rerun()
        
        # Training progress
        if 'history' in training_data:
            history = training_data['history']
            
            # Current epoch and metrics
            if history:
                current_epoch = len(history)
                latest_metrics = history[-1]
                
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    st.metric("Current Epoch", current_epoch)
                
                if 'loss' in latest_metrics:
                    with col2:
                        st.metric("Current Loss", f"{latest_metrics['loss']:.4f}")
                
                if 'accuracy' in latest_metrics:
                    with col3:
                        st.metric("Current Accuracy", f"{latest_metrics['accuracy']:.4f}")
                
                # Training curves
                st.subheader("Training Curves")
                
                fig = make_subplots(
                    rows=2, cols=2,
                    subplot_titles=("Loss", "Accuracy", "Learning Rate", "Metrics"),
                    specs=[[{"secondary_y": False}, {"secondary_y": False}],
                           [{"secondary_y": False}, {"secondary_y": False}]]
                )
                
                epochs = list(range(1, len(history) + 1))
                
                # Loss
                if 'loss' in history[0]:
                    train_loss = [h['loss'] for h in history]
                    fig.add_trace(go.Scatter(x=epochs, y=train_loss, name='Training Loss'),
                                 row=1, col=1)
                
                if 'val_loss' in history[0]:
                    val_loss = [h['val_loss'] for h in history if 'val_loss' in h]
                    if val_loss:
                        val_epochs = epochs[:len(val_loss)]
                        fig.add_trace(go.Scatter(x=val_epochs, y=val_loss, name='Validation Loss'),
                                     row=1, col=1)
                
                # Accuracy
                if 'accuracy' in history[0]:
                    train_acc = [h['accuracy'] for h in history]
                    fig.add_trace(go.Scatter(x=epochs, y=train_acc, name='Training Accuracy'),
                                 row=1, col=2)
                
                if 'val_accuracy' in history[0]:
                    val_acc = [h['val_accuracy'] for h in history if 'val_accuracy' in h]
                    if val_acc:
                        val_epochs = epochs[:len(val_acc)]
                        fig.add_trace(go.Scatter(x=val_epochs, y=val_acc, name='Validation Accuracy'),
                                     row=1, col=2)
                
                # Learning rate
                if 'lr' in history[0]:
                    lr_values = [h['lr'] for h in history if 'lr' in h]
                    if lr_values:
                        lr_epochs = epochs[:len(lr_values)]
                        fig.add_trace(go.Scatter(x=lr_epochs, y=lr_values, name='Learning Rate'),
                                     row=2, col=1)
                
                # Additional metrics
                other_metrics = [k for k in history[0].keys() 
                               if k not in ['loss', 'val_loss', 'accuracy', 'val_accuracy', 'lr']]
                
                if other_metrics:
                    for metric in other_metrics[:2]:  # Show first 2 additional metrics
                        metric_values = [h.get(metric) for h in history if metric in h]
                        if metric_values:
                            fig.add_trace(go.Scatter(x=epochs, y=metric_values, name=metric),
                                         row=2, col=2)
                
                fig.update_layout(height=600, showlegend=True)
                st.plotly_chart(fig, use_container_width=True)
        
        # System metrics
        if 'system_metrics' in training_data:
            st.subheader("System Metrics")
            
            system_metrics = training_data['system_metrics']
            
            col1, col2 = st.columns(2)
            
            with col1:
                if 'cpu_usage' in system_metrics:
                    st.metric("CPU Usage", f"{system_metrics['cpu_usage']:.1f}%")
                if 'memory_usage' in system_metrics:
                    st.metric("Memory Usage", f"{system_metrics['memory_usage']:.1f}%")
            
            with col2:
                if 'gpu_usage' in system_metrics:
                    st.metric("GPU Usage", f"{system_metrics['gpu_usage']:.1f}%")
                if 'gpu_memory' in system_metrics:
                    st.metric("GPU Memory", f"{system_metrics['gpu_memory']:.1f}%")
    
    def create_data_exploration_dashboard(self, df: pd.DataFrame) -> None:
        """
        Create interactive data exploration dashboard
        
        Args:
            df: Input DataFrame
        """
        st.title("🔍 Data Exploration Dashboard")
        
        # Dataset overview
        st.subheader("Dataset Overview")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            st.metric("Rows", df.shape[0])
        
        with col2:
            st.metric("Columns", df.shape[1])
        
        with col3:
            missing_pct = (df.isnull().sum().sum() / (df.shape[0] * df.shape[1])) * 100
            st.metric("Missing Values", f"{missing_pct:.1f}%")
        
        # Data preview
        st.subheader("Data Preview")
        
        # Show sample data
        st.write("**Sample Data:**")
        st.dataframe(df.head(), use_container_width=True)
        
        # Column information
        st.subheader("Column Information")
        
        col_info = []
        for col in df.columns:
            col_info.append({
                'Column': col,
                'Type': str(df[col].dtype),
                'Non-Null Count': df[col].count(),
                'Unique Values': df[col].nunique(),
                'Missing Count': df[col].isnull().sum()
            })
        
        col_info_df = pd.DataFrame(col_info)
        st.dataframe(col_info_df, use_container_width=True)
        
        # Visualizations
        numerical_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        categorical_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
        
        if numerical_cols:
            st.subheader("Numerical Features")
            
            # Select feature to visualize
            selected_num = st.selectbox("Select Numerical Feature", numerical_cols)
            
            col1, col2 = st.columns(2)
            
            with col1:
                # Distribution
                fig = px.histogram(df, x=selected_num, title=f"Distribution of {selected_num}")
                st.plotly_chart(fig, use_container_width=True)
            
            with col2:
                # Box plot
                fig = px.box(df, y=selected_num, title=f"Box Plot of {selected_num}")
                st.plotly_chart(fig, use_container_width=True)
        
        if categorical_cols:
            st.subheader("Categorical Features")
            
            # Select feature to visualize
            selected_cat = st.selectbox("Select Categorical Feature", categorical_cols)
            
            col1, col2 = st.columns(2)
            
            with col1:
                # Value counts
                value_counts = df[selected_cat].value_counts()
                fig = px.bar(x=value_counts.index, y=value_counts.values,
                            title=f"Value Counts of {selected_cat}")
                st.plotly_chart(fig, use_container_width=True)
            
            with col2:
                # Pie chart
                fig = px.pie(values=value_counts.values, names=value_counts.index,
                            title=f"Distribution of {selected_cat}")
                st.plotly_chart(fig, use_container_width=True)
        
        # Correlation matrix
        if len(numerical_cols) > 1:
            st.subheader("Correlation Matrix")
            
            corr_matrix = df[numerical_cols].corr()
            fig = px.imshow(corr_matrix, text_auto=True, aspect="auto",
                           title="Feature Correlation Matrix")
            st.plotly_chart(fig, use_container_width=True)
    
    def create_model_interpretability_dashboard(self, model, X: pd.DataFrame, 
                                              feature_names: List[str]) -> None:
        """
        Create model interpretability dashboard
        
        Args:
            model: Trained model
            X: Feature matrix
            feature_names: List of feature names
        """
        st.title("🔮 Model Interpretability Dashboard")
        
        # Feature importance
        if hasattr(model, 'feature_importances_'):
            st.subheader("Feature Importance")
            
            importance_scores = model.feature_importances_
            importance_df = pd.DataFrame({
                'Feature': feature_names,
                'Importance': importance_scores
            }).sort_values('Importance', ascending=False)
            
            # Bar chart
            fig = px.bar(importance_df.head(20), x='Importance', y='Feature',
                        orientation='h', title="Top 20 Feature Importance")
            st.plotly_chart(fig, use_container_width=True)
            
            # Table
            st.dataframe(importance_df.head(20), use_container_width=True)
        
        # SHAP values (if available)
        try:
            import shap
            
            st.subheader("SHAP Analysis")
            
            # Create SHAP explainer
            explainer = shap.Explainer(model, X)
            shap_values = explainer(X)
            
            # Summary plot
            fig, ax = plt.subplots(figsize=(10, 8))
            shap.summary_plot(shap_values, X, feature_names=feature_names, show=False)
            st.pyplot(fig)
            plt.close()
            
            # Feature importance based on SHAP
            shap_importance = np.abs(shap_values.values).mean(0)
            shap_df = pd.DataFrame({
                'Feature': feature_names,
                'SHAP Importance': shap_importance
            }).sort_values('SHAP Importance', ascending=False)
            
            fig = px.bar(shap_df.head(20), x='SHAP Importance', y='Feature',
                        orientation='h', title="Top 20 SHAP Feature Importance")
            st.plotly_chart(fig, use_container_width=True)
            
        except ImportError:
            st.warning("SHAP library not installed. Install with: pip install shap")
        except Exception as e:
            st.error(f"Error computing SHAP values: {str(e)}")
        
        # Prediction distribution
        st.subheader("Prediction Distribution")
        
        predictions = model.predict(X)
        
        col1, col2 = st.columns(2)
        
        with col1:
            fig = px.histogram(predictions, title="Prediction Distribution")
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            if hasattr(model, 'predict_proba'):
                probabilities = model.predict_proba(X)
                fig = px.histogram(probabilities.max(axis=1), 
                                 title="Maximum Prediction Probability")
                st.plotly_chart(fig, use_container_width=True)
    
    def create_real_time_monitoring_dashboard(self, metrics_db: str) -> None:
        """
        Create real-time monitoring dashboard
        
        Args:
            metrics_db: Path to SQLite database containing metrics
        """
        st.title("📈 Real-time Monitoring Dashboard")
        
        # Auto-refresh
        auto_refresh = st.checkbox("Auto Refresh (10s)", value=True)
        
        if auto_refresh:
            time.sleep(10)
            st.experimental_rerun()
        
        try:
            conn = sqlite3.connect(metrics_db)
            
            # Load recent metrics
            recent_metrics = pd.read_sql_query("""
                SELECT * FROM metrics 
                WHERE timestamp > datetime('now', '-1 hour')
                ORDER BY timestamp DESC
            """, conn)
            
            if not recent_metrics.empty:
                # Current metrics
                st.subheader("Current Metrics")
                
                latest = recent_metrics.iloc[0]
                
                col1, col2, col3, col4 = st.columns(4)
                
                with col1:
                    st.metric("Requests/min", latest.get('requests_per_minute', 0))
                
                with col2:
                    st.metric("Avg Response Time", f"{latest.get('avg_response_time', 0):.2f}s")
                
                with col3:
                    st.metric("Error Rate", f"{latest.get('error_rate', 0):.2%}")
                
                with col4:
                    st.metric("CPU Usage", f"{latest.get('cpu_usage', 0):.1f}%")
                
                # Time series charts
                st.subheader("Metrics Over Time")
                
                fig = make_subplots(
                    rows=2, cols=2,
                    subplot_titles=("Requests/min", "Response Time", "Error Rate", "CPU Usage"),
                    specs=[[{"secondary_y": False}, {"secondary_y": False}],
                           [{"secondary_y": False}, {"secondary_y": False}]]
                )
                
                recent_metrics['timestamp'] = pd.to_datetime(recent_metrics['timestamp'])
                
                # Requests per minute
                fig.add_trace(go.Scatter(x=recent_metrics['timestamp'], 
                                       y=recent_metrics['requests_per_minute'],
                                       name='Requests/min'), row=1, col=1)
                
                # Response time
                fig.add_trace(go.Scatter(x=recent_metrics['timestamp'], 
                                       y=recent_metrics['avg_response_time'],
                                       name='Response Time'), row=1, col=2)
                
                # Error rate
                fig.add_trace(go.Scatter(x=recent_metrics['timestamp'], 
                                       y=recent_metrics['error_rate'] * 100,
                                       name='Error Rate (%)'), row=2, col=1)
                
                # CPU usage
                fig.add_trace(go.Scatter(x=recent_metrics['timestamp'], 
                                       y=recent_metrics['cpu_usage'],
                                       name='CPU Usage (%)'), row=2, col=2)
                
                fig.update_layout(height=600, showlegend=False)
                st.plotly_chart(fig, use_container_width=True)
            
            else:
                st.warning("No recent metrics data available")
            
            conn.close()
            
        except Exception as e:
            st.error(f"Error loading metrics: {str(e)}")
    
    def run_dashboard(self, dashboard_type: str = "experiment", **kwargs):
        """
        Run the dashboard
        
        Args:
            dashboard_type: Type of dashboard to run
            **kwargs: Additional arguments for specific dashboard types
        """
        self.setup_page()
        
        if dashboard_type == "experiment":
            if 'experiments_df' not in kwargs:
                st.error("experiments_df required for experiment dashboard")
                return
            self.create_experiment_dashboard(kwargs['experiments_df'])
        
        elif dashboard_type == "performance":
            if 'models_df' not in kwargs:
                st.error("models_df required for performance dashboard")
                return
            self.create_model_performance_dashboard(kwargs['models_df'])
        
        elif dashboard_type == "training":
            if 'training_data' not in kwargs:
                st.error("training_data required for training dashboard")
                return
            self.create_training_monitoring_dashboard(kwargs['training_data'])
        
        elif dashboard_type == "exploration":
            if 'df' not in kwargs:
                st.error("df required for exploration dashboard")
                return
            self.create_data_exploration_dashboard(kwargs['df'])
        
        elif dashboard_type == "interpretability":
            if any(k not in kwargs for k in ['model', 'X', 'feature_names']):
                st.error("model, X, and feature_names required for interpretability dashboard")
                return
            self.create_model_interpretability_dashboard(
                kwargs['model'], kwargs['X'], kwargs['feature_names'])
        
        elif dashboard_type == "monitoring":
            if 'metrics_db' not in kwargs:
                st.error("metrics_db required for monitoring dashboard")
                return
            self.create_real_time_monitoring_dashboard(kwargs['metrics_db'])
        
        else:
            st.error(f"Unknown dashboard type: {dashboard_type}")
    
    def export_dashboard_config(self, filepath: str):
        """Export dashboard configuration to file"""
        config = {
            'title': self.title,
            'layout': self.layout,
            'data_sources': {name: {'description': info['description']} 
                            for name, info in self.data_sources.items()},
            'filters': self.filters
        }
        
        with open(filepath, 'w') as f:
            json.dump(config, f, indent=2)
    
    def import_dashboard_config(self, filepath: str):
        """Import dashboard configuration from file"""
        with open(filepath, 'r') as f:
            config = json.load(f)
        
        self.title = config.get('title', 'ML Dashboard')
        self.layout = config.get('layout', 'wide')
        self.filters = config.get('filters', {})
