streamlit_ui.py

"""
Streamlit UI for Multiagent Data Visualization System
Allows users to:
1. Select domain and dataset size
2. Watch real-time agent execution
3. View generated code for all visualizations
4. Select and execute specific visualizations
"""

import streamlit as st
import pandas as pd
import asyncio
import os
from io import StringIO
import sys
import time

# Add project to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from utils.fake_data_generator import generate_fake_data
from utils.sklearn_data_loader import (load_sklearn_dataset, list_available_datasets,
                                       SKLEARN_DATASETS, generate_reference_plot)
from utils.logging_manager import get_logging_manager, reset_logging_manager
from workflows.langgraph_workflow import run_visualization_pipeline
import matplotlib.pyplot as plt

# Page configuration
st.set_page_config(
    page_title="Multiagent Data Visualization",
    page_icon=None,
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .stTabs [data-baseweb="tab-list"] button [data-testid="stMarkdownContainer"] p {
        font-size: 18px;
    }
    .agent-status-active {
        color: #00ff00;
        font-weight: bold;
    }
    .agent-status-complete {
        color: #0099ff;
        font-weight: bold;
    }
    .agent-status-pending {
        color: #cccccc;
        font-weight: normal;
    }
</style>
""", unsafe_allow_html=True)

# Title and description
st.title("Multiagent Data Visualization System")
st.markdown("""
**AI-powered visualization creation with domain-specific intelligence**

Explore how intelligent agents work together to analyze your data and generate optimal visualizations.
""")

# ============================================================================
# SIDEBAR - USER INPUTS
# ============================================================================
st.sidebar.header("Configuration")

# ── Data source selection ────────────────────────────────────────────────
data_source = st.sidebar.radio(
    "Data Source",
    options=["Synthetic Data", "Scikit-learn Dataset"],
    help="Use AI-generated synthetic data or a real scikit-learn dataset"
)

if data_source == "Synthetic Data":
    # Domain selection
    domain = st.sidebar.selectbox(
        "Select Domain",
        options=["healthcare", "finance", "engineering", "general"],
        help="Choose a domain to generate relevant data"
    )
    # Dataset size selection
    dataset_size = st.sidebar.slider(
        "Dataset Size",
        min_value=50,
        max_value=1000,
        value=100,
        step=50,
        help="Number of rows to generate"
    )
    sklearn_dataset_key = None
    enable_validation = False  # No reference plots for synthetic data
else:
    available = list_available_datasets()
    sklearn_options = {d["display_name"]: d["key"] for d in available}
    selected_display = st.sidebar.selectbox(
        "Select Dataset",
        options=list(sklearn_options.keys()),
        help="Real-world scikit-learn dataset for validation"
    )
    sklearn_dataset_key = sklearn_options[selected_display]
    meta = SKLEARN_DATASETS[sklearn_dataset_key]
    domain = meta["domain"]
    dataset_size = None  # determined by the dataset
    st.sidebar.info(
        f"**{meta['display_name']}**\n\n{meta['description']}\n\n"
        f"Reference charts: {', '.join(meta['reference_charts'])}"
    )
    # Validation toggle — only available for scikit-learn datasets
    enable_validation = st.sidebar.checkbox(
        "Run Validation Agent",
        value=True,
        help="Run the 7th ValidationAgent to compare agent plots against reference plots"
    )

# Run button
run_pipeline = st.sidebar.button(
    "Run Pipeline",
    key="run_btn",
    width='stretch'
)

# ============================================================================
# MAIN CONTENT - TABS
# ============================================================================
tab1, tab2, tab3, tab4, tab5, tab6, tab7 = st.tabs([
    "Pipeline Execution",
    "Agent Monitor",
    "Generated Code",
    "Visualizations",
    "Data Preview",
    "Execution Logs",
    "Validation Results"
])

# ============================================================================
# TAB 1: PIPELINE EXECUTION
# ============================================================================
with tab1:
    st.subheader("Agent Execution Pipeline")
    
    if run_pipeline:
        # Reset logging manager for new run
        reset_logging_manager()
        
        # Initialize session state to track pipeline execution
        if 'pipeline_run' not in st.session_state:
            st.session_state.pipeline_run = True
            st.session_state.result = None
            st.session_state.generated_data = None
        
        # Create columns for agent status display
        col1, col2, col3 = st.columns(3)
        
        with col1:
            agent1_container = st.container()
            agent1_placeholder = agent1_container.empty()
        
        with col2:
            agent2_container = st.container()
            agent2_placeholder = agent2_container.empty()
        
        with col3:
            agent3_container = st.container()
            agent3_placeholder = agent3_container.empty()

        # Row 2: Agents 4-5
        col4, col5 = st.columns(2)
        with col4:
            agent4_placeholder = st.empty()
        with col5:
            agent5_placeholder = st.empty()

        # Show agents 4-5 as pending
        for _ph, _name in [
            (agent4_placeholder, "Agent 4: User Collaboration"),
            (agent5_placeholder, "Agent 5: Validation Agent"),
        ]:
            with _ph.container():
                st.markdown(f"""
                <div style='background-color: #1f77b4; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>{_name}</h4>
                    <p style='color: #00ff00;'>Running...</p>
                </div>
                """, unsafe_allow_html=True)
        
        # Step 1: Data Loading / Generation
        _loading_msg = "Loading scikit-learn dataset..." if data_source != "Synthetic Data" else "Generating synthetic data..."
        with agent1_placeholder.container():
            st.markdown(f"""
            <div style='background-color: #1f77b4; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                <h4 style='color: white;'>Agent 1: Data Analysis</h4>
                <p style='color: #00ff00;'>{_loading_msg}</p>
            </div>
            """, unsafe_allow_html=True)
        
        try:
            # ── Generate data based on source ──────────────────────────────
            if data_source == "Synthetic Data":
                generated_data = generate_fake_data(domain=domain, num_rows=dataset_size)
                ref_charts = []
                preferred_chart = ""
                ds_source_label = "synthetic"
                ds_name = domain
            else:
                generated_data, sklearn_meta = load_sklearn_dataset(sklearn_dataset_key)
                ref_charts = sklearn_meta.get("reference_charts", [])
                preferred_chart = sklearn_meta.get("preferred_chart", "")
                ds_source_label = "sklearn"
                ds_name = sklearn_dataset_key

            st.session_state.generated_data = generated_data
            st.session_state.ref_charts = ref_charts
            st.session_state.preferred_chart = preferred_chart
            st.session_state.data_source_label = data_source
            st.session_state.sklearn_dataset_key = sklearn_dataset_key  # None for synthetic
            
            _row_count = len(generated_data)
            _source_label = "Scikit-learn" if data_source != "Synthetic Data" else "Synthetic"
            with agent1_placeholder.container():
                st.markdown(f"""
                <div style='background-color: #2ca02c; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 1: Data Analysis</h4>
                    <p style='color: #00ff00;'>Data Loaded ({_source_label})</p>
                    <p style='color: white; font-size: 14px;'>
                        Domain: <b>{domain.upper()}</b><br>
                        Rows: <b>{_row_count}</b><br>
                        Columns: <b>{len(generated_data.columns)}</b>
                    </p>
                </div>
                """, unsafe_allow_html=True)
            
            # Step 2: Visualization Planning
            with agent2_placeholder.container():
                st.markdown("""
                <div style='background-color: #1f77b4; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 2: Visualization Planner</h4>
                    <p style='color: #00ff00;'>&#x1F916; Querying LLM / planning visualizations...</p>
                </div>
                """, unsafe_allow_html=True)
            
            # Run the workflow
            async def run_workflow():
                data_dict = generated_data.to_dict('records')
                result = await run_visualization_pipeline(
                    data_dict,
                    user_query=f"Create visualizations for {domain} data",
                    reference_charts=ref_charts,
                    preferred_chart=preferred_chart,
                    dataset_source=ds_source_label,
                    dataset_name=ds_name,
                    enable_validation=enable_validation,
                )
                return result
            
            # Run async workflow
            result = asyncio.run(run_workflow())
            st.session_state.result = result
            
            with agent2_placeholder.container():
                recommendations = result.get('chart_recommendations', [])
                rec_types = [rec.get('type', 'Unknown') for rec in recommendations]
                _llm_used = result.get('llm_used', False)
                _llm_model = result.get('llm_model', '')
                _mode_badge = f"&#x1F916; LLM-Powered ({_llm_model})" if _llm_used else "&#x1F4D0; Rule-Based Fallback"
                _badge_color = "#00ff00" if _llm_used else "#ffaa00"
                st.markdown(f"""
                <div style='background-color: #2ca02c; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 2: Visualization Planner</h4>
                    <p style='color: {_badge_color};'>{_mode_badge}</p>
                    <p style='color: white; font-size: 14px;'>
                        Recommended Charts: <b>{', '.join(rec_types)}</b><br>
                        Selected: <b>{result.get('selected_chart_type', 'Unknown').upper()}</b>
                    </p>
                </div>
                """, unsafe_allow_html=True)
            
            # Step 3: Code Generation
            with agent3_placeholder.container():
                st.markdown("""
                <div style='background-color: #1f77b4; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 3: Code Generator</h4>
                    <p style='color: #00ff00;'>Generating code...</p>
                </div>
                """, unsafe_allow_html=True)
            
            all_viz_code = result.get('all_visualizations_code', {})
            
            with agent3_placeholder.container():
                # Count actual libraries generated
                _libs = set()
                for _chart_code in all_viz_code.values():
                    if isinstance(_chart_code, dict):
                        _libs.update(_chart_code.keys())
                _lib_count = len(_libs)
                _lib_names = ', '.join(l.capitalize() for l in sorted(_libs)) if _libs else 'N/A'
                st.markdown(f"""
                <div style='background-color: #2ca02c; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 3: Code Generator</h4>
                    <p style='color: #00ff00;'>Code Generation Complete</p>
                    <p style='color: white; font-size: 14px;'>
                        Total Visualizations: <b>{len(all_viz_code)}</b><br>
                        Libraries Supported: <b>{_lib_count}</b> ({_lib_names})
                    </p>
                </div>
                """, unsafe_allow_html=True)

            # Update Agents 4-5 cards after workflow completes
            _quality = result.get('quality_score', 'N/A')
            _domain_ctx = result.get('domain_context', domain.capitalize())
            _selected = result.get('selected_chart_type', 'N/A').upper()

            with agent4_placeholder.container():
                st.markdown(f"""
                <div style='background-color: #2ca02c; padding: 20px; border-radius: 10px; margin: 10px 0;'>
                    <h4 style='color: white;'>Agent 4: User Collaboration</h4>
                    <p style='color: #00ff00;'>Collaboration Layer Active</p>
                    <p style='color: white; font-size: 14px;'>
                        Selected Chart: <b>{_selected}</b><br>
                        Human Feedback: <b>{'Required' if result.get('human_feedback_required') else 'Not Required'}</b>
                    </p>
                </div>
                """, unsafe_allow_html=True)

            # Update Agent 5 placeholder with final status
            if enable_validation:
                val = result.get("validation_results", {})
                if val.get("validated"):
                    with agent5_placeholder.container():
                        st.markdown(f"""
                        <div style='background-color: #2ca02c; padding: 15px; border-radius: 10px; margin: 10px 0;'>
                            <h4 style='color: white;'>Agent 5: Validation Agent</h4>
                            <p style='color: #00ff00;'>Validation Complete</p>
                            <p style='color: white; font-size: 14px;'>{val.get('verdict', '')}</p>
                        </div>
                        """, unsafe_allow_html=True)
            else:
                with agent5_placeholder.container():
                    st.markdown(f"""
                    <div style='background-color: #2ca02c; padding: 15px; border-radius: 10px; margin: 10px 0;'>
                        <h4 style='color: white;'>Agent 5: Validation Agent</h4>
                        <p style='color: #00ff00;'>Not Enabled</p>
                        <p style='color: white; font-size: 12px;'>Enable in sidebar for validation</p>
                    </div>
                    """, unsafe_allow_html=True)

            # Summary
            st.success("Pipeline Execution Complete!")
            
            # Show any errors
            errors = result.get('error_log', [])
            if errors:
                st.warning("Some errors occurred during execution:")
                for error in errors:
                    st.write(f"- {error}")
        
        except Exception as e:
            st.error(f"Pipeline Failed: {str(e)}")
            import traceback
            st.error(traceback.format_exc())
    
    else:
        st.info("Configure domain and dataset size in the sidebar, then click **Run Pipeline** to start!")

# ============================================================================
# TAB 2: AGENT MONITOR (NEW)
# ============================================================================
with tab2:
    st.subheader("Agent Execution Monitor")
    
    if 'result' in st.session_state and st.session_state.result:
        result = st.session_state.result
        logger = get_logging_manager()
        
        # Display agent timeline
        timeline = logger.get_agent_timeline()
        
        if timeline:
            st.markdown("### Agent Execution Timeline")
            
            # Create a timeline view
            for idx, entry in enumerate(timeline, 1):
                col1, col2, col3, col4 = st.columns([1, 2, 1.5, 1.5])
                
                with col1:
                    st.markdown(f"**{idx}**")
                
                with col2:
                    status_emoji = {
                        "started": "▶️",
                        "processing": "⚙️",
                        "completed": "✓",
                        "error": "✗"
                    }.get(entry.get("status", ""), "•")
                    st.markdown(f"{status_emoji} **{entry.get('agent', 'Unknown')}**")
                
                with col3:
                    status = entry.get("status", "").upper()
                    status_color = {
                        "STARTED": "#FFA500",
                        "PROCESSING": "#0099FF",
                        "COMPLETED": "#00CC00",
                        "ERROR": "#FF0000"
                    }.get(status, "#CCCCCC")
                    st.markdown(f"<span style='color: {status_color};'>{status}</span>", unsafe_allow_html=True)
                
                with col4:
                    duration = entry.get("duration_ms", 0)
                    if duration > 0:
                        st.markdown(f"{duration:.0f}ms")
                    else:
                        st.markdown(entry.get("timestamp", "-"))
                
                # Show input/output if available
                if entry.get("input") or entry.get("output"):
                    with st.expander("Details"):
                        if entry.get("input"):
                            st.markdown(f"**Input:** {entry.get('input')}")
                        if entry.get("output"):
                            st.markdown(f"**Output:** {entry.get('output')}")
                        if entry.get("error"):
                            st.error(f"**Error:** {entry.get('error')}")
                
                st.markdown("---")
            
            # Display agent interactions
            interactions = logger.get_interactions()
            if interactions:
                st.markdown("### Agent-to-Agent Communications")
                
                for interaction in interactions:
                    from_agent = interaction.get("from_agent", "Unknown")
                    to_agent = interaction.get("to_agent", "Unknown")
                    data_summary = interaction.get("data_summary", "")
                    interaction_type = interaction.get("interaction_type", "")
                    status = interaction.get("status", "pending")
                    
                    # Color based on status
                    status_color = {
                        "pending": "#FFA500",
                        "processing": "#0099FF",
                        "complete": "#00CC00",
                        "error": "#FF0000"
                    }.get(status, "#CCCCCC")
                    
                    st.markdown(f"""
                    <div style='background-color: #1a1a1a; padding: 12px; border-radius: 8px; margin: 8px 0; border-left: 4px solid {status_color};'>
                        <b style='color: #00FF00;'>{from_agent}</b> <span style='color: gray;'>→</span> <b style='color: #00FF00;'>{to_agent}</b><br>
                        <span style='color: #CCCCCC; font-size: 12px;'>{interaction_type}</span><br>
                        <span style='color: #AAAAAA;'>{data_summary}</span>
                    </div>
                    """, unsafe_allow_html=True)
        else:
            st.info("No agent timeline data available. Run the pipeline first!")
    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# TAB 3: GENERATED CODE
# ============================================================================
with tab3:
    st.subheader("Generated Visualization Code")
    
    if 'result' in st.session_state and st.session_state.result:
        result = st.session_state.result
        all_viz_code = result.get('all_visualizations_code', {})
        
        if all_viz_code:
            # Chart selection
            chart_options = list(all_viz_code.keys())
            selected_chart = st.selectbox(
                "Select visualization type:",
                options=chart_options
            )
            
            # Library selection
            library_options = ["matplotlib", "plotly", "seaborn"]
            selected_library = st.selectbox(
                "Select library/language:",
                options=library_options
            )
            
            # Get and display code
            if selected_chart in all_viz_code:
                code_dict = all_viz_code[selected_chart]
                
                if isinstance(code_dict, dict) and selected_library in code_dict:
                    code_text = code_dict[selected_library]
                    
                    st.markdown(f"### {selected_library.upper()} Code for {selected_chart.upper()}")
                    st.code(code_text, language="python")
                    
                    # Copy button
                    col1, col2 = st.columns(2)
                    with col1:
                        if st.button(f"Copy to Clipboard"):
                            st.success("Code copied to clipboard!")
                    
                    with col2:
                        if st.button(f"Execute Visualization"):
                            st.info("Executing visualization...")
                            try:
                                # Execute the code
                                exec_globals = {
                                    "df": st.session_state.generated_data,
                                    "pd": pd,
                                    "plt": plt,
                                    "st": st
                                }
                                
                                # Modify code to display in Streamlit
                                modified_code = code_text.replace("plt.show()", "st.pyplot(plt)")
                                
                                exec(modified_code, exec_globals)
                                st.success("Visualization executed successfully!")
                            
                            except Exception as e:
                                st.error(f"Execution failed: {str(e)}")
                else:
                    st.warning(f"Code not available for {selected_library}")
            else:
                st.warning("Code not available for selected chart")
        else:
            st.info("No visualizations generated yet. Run the pipeline first!")
    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# TAB 4: VISUALIZATIONS
# ============================================================================
with tab4:
    st.subheader("Interactive Visualizations")
    
    if 'result' in st.session_state and st.session_state.result and 'generated_data' in st.session_state:
        result = st.session_state.result
        generated_data = st.session_state.generated_data
        all_viz_code = result.get('all_visualizations_code', {})
        
        if all_viz_code:
            # Create tabs for each visualization
            viz_tabs = st.tabs([f"{chart.upper()}" for chart in all_viz_code.keys()])
            
            for idx, (chart_type, viz_tab) in enumerate(zip(all_viz_code.keys(), viz_tabs)):
                with viz_tab:
                    st.markdown(f"### {chart_type.upper()} Visualization")
                    
                    code_dict = all_viz_code[chart_type]
                    
                    # For matplotlib/seaborn
                    if isinstance(code_dict, dict) and "matplotlib" in code_dict:
                        try:
                            exec_globals = {
                                "df": generated_data,
                                "pd": pd,
                                "plt": plt
                            }
                            
                            matplotlib_code = code_dict["matplotlib"]
                            modified_code = matplotlib_code.replace("plt.show()", "")
                            
                            exec(modified_code, exec_globals)
                            st.pyplot(plt.gcf())
                            plt.clf()
                        
                        except Exception as e:
                            st.error(f"Error displaying visualization: {str(e)}")
                    
                    # Show code snippet
                    with st.expander("Show Code"):
                        if "matplotlib" in code_dict:
                            st.code(code_dict["matplotlib"], language="python")
        else:
            st.info("No visualizations generated yet. Run the pipeline first!")
    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# TAB 5: DATA PREVIEW
# ============================================================================
with tab5:
    st.subheader("Data Preview & Statistics")
    
    if 'generated_data' in st.session_state:
        generated_data = st.session_state.generated_data
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.write("### Dataset Info")
            st.write(f"**Rows:** {len(generated_data)}")
            st.write(f"**Columns:** {len(generated_data.columns)}")
            st.write(f"**Domain:** {domain.upper()}")
            st.write(f"**Memory Usage:** {generated_data.memory_usage(deep=True).sum() / 1024:.2f} KB")
        
        with col2:
            st.write("### Column Info")
            col_info = pd.DataFrame({
                'Column': generated_data.columns,
                'Type': [str(dtype) for dtype in generated_data.dtypes],
                'Non-Null': generated_data.count().values
            })
            st.dataframe(col_info, width='stretch')
        
        # Data preview
        st.write("### Data Preview (First 10 Rows)")
        # Convert datetime columns to string for Arrow compatibility
        preview_data = generated_data.head(10).copy()
        for col in preview_data.columns:
            if preview_data[col].dtype == 'object' or 'datetime' in str(preview_data[col].dtype):
                preview_data[col] = preview_data[col].astype(str)
        st.dataframe(preview_data, width='stretch')
        
        # Statistics
        st.write("### Statistical Summary")
        stats = generated_data.describe()
        # Convert stats columns with datetime or object types
        for col in stats.columns:
            if stats[col].dtype == 'object' or 'datetime' in str(stats[col].dtype):
                stats[col] = stats[col].astype(str)
        st.dataframe(stats, width='stretch')
        
        # Download data
        csv = generated_data.to_csv(index=False)
        ds_label = st.session_state.get("data_source_label", "synthetic")
        file_label = f"{domain}_data" if ds_label == "synthetic" else f"sklearn_{domain}"
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name=f"{file_label}_{len(generated_data)}_rows.csv",
            mime="text/csv"
        )
    
    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# TAB 6: EXECUTION LOGS
# ============================================================================
with tab6:
    st.subheader("System Execution Logs")
    
    if 'result' in st.session_state and st.session_state.result:
        logger = get_logging_manager()
        logs = logger.get_logs()
        
        if logs:
            # Display summary
            summary = logger.get_summary()
            col1, col2, col3, col4 = st.columns(4)
            
            with col1:
                st.metric("Total Logs", summary.get("total_logs", 0))
            with col2:
                st.metric("INFO", summary.get("log_levels", {}).get("INFO", 0))
            with col3:
                st.metric("SUCCESS", summary.get("log_levels", {}).get("SUCCESS", 0))
            with col4:
                st.metric("ERROR", summary.get("log_levels", {}).get("ERROR", 0))
            
            st.markdown("---")
            
            # Filter options
            col1, col2 = st.columns(2)
            with col1:
                filter_source = st.selectbox(
                    "Filter by source:",
                    options=["ALL"] + list(set(log["source"] for log in logs)),
                    key="log_filter_source"
                )
            with col2:
                search_text = st.text_input("Search logs:", key="log_search")
            
            # Apply filters
            filtered_logs = logs
            if filter_source != "ALL":
                filtered_logs = [log for log in filtered_logs if log["source"] == filter_source]
            if search_text:
                filtered_logs = [log for log in filtered_logs if search_text.lower() in log["message"].lower()]
            
            # Display logs with color coding
            st.markdown("""<style>
                .log-container { max-height: 600px; overflow-y: auto; }
                .log-entry { padding: 10px; margin: 5px 0; border-radius: 5px; font-family: monospace; }
                .log-info { background-color: #1a3a52; border-left: 4px solid #0099FF; }
                .log-success { background-color: #1a3a1a; border-left: 4px solid #00CC00; }
                .log-warning { background-color: #3a2a1a; border-left: 4px solid #FFA500; }
                .log-error { background-color: #3a1a1a; border-left: 4px solid #FF0000; }
                .log-time { color: #888888; font-size: 12px; }
                .log-source { color: #00FF00; font-weight: bold; }
                .log-message { color: #CCCCCC; }
            </style>""", unsafe_allow_html=True)
            
            # Display filtered logs
            for log in reversed(filtered_logs):  # Show newest first
                level = log.get("level", "INFO")
                css_class = {
                    "INFO": "log-info",
                    "SUCCESS": "log-success",
                    "WARNING": "log-warning",
                    "ERROR": "log-error"
                }.get(level, "log-info")
                
                st.markdown(f"""
                <div class='log-entry {css_class}'>
                    <span class='log-time'>[{log.get('timestamp', '')}]</span> 
                    <span class='log-source'>{log.get('source', 'SYSTEM'):12}</span> 
                    <span class='log-message'>{log.get('message', '')}</span>
                </div>
                """, unsafe_allow_html=True)
            
            # Download logs
            logs_csv = pd.DataFrame(filtered_logs).to_csv(index=False)
            st.download_button(
                label="Download Logs (CSV)",
                data=logs_csv,
                file_name="execution_logs.csv",
                mime="text/csv"
            )
        else:
            st.info("No logs available. Run the pipeline first!")
    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# TAB 7: VALIDATION RESULTS
# ============================================================================
with tab7:
    st.subheader("Validation Agent Results")

    if 'result' in st.session_state and st.session_state.result:
        result = st.session_state.result
        val = result.get("validation_results", {})

        if not val:
            st.info("Validation was not enabled for this run. Enable **Run Validation Agent** in the sidebar and re-run.")
        elif not val.get("validated", False):
            st.error(f"Validation failed: {val.get('error', 'Unknown error')}")
        else:
            # Data source context
            ds_label = st.session_state.get("data_source_label", "Synthetic Data")
            ref_charts = st.session_state.get("ref_charts", [])
            preferred = st.session_state.get("preferred_chart", "")
            is_sklearn = (ds_label == "Scikit-learn Dataset")

            overall = val.get("overall_validation_score", 0)

            if is_sklearn:
                st.markdown(f"**Dataset:** `{result.get('dataset_name', '')}` &nbsp;|&nbsp; "
                            f"**Reference charts:** `{', '.join(ref_charts)}` &nbsp;|&nbsp; "
                            f"**Preferred:** `{preferred}`")

            st.markdown("---")

            # -- Side-by-side visual comparison (sklearn only) ---------
            ds_key = st.session_state.get("sklearn_dataset_key")
            if is_sklearn and ds_key and ref_charts:
                st.markdown("---")
                st.markdown("#### Reference Plots vs Agent-Generated Plots")
                st.caption(
                    "Each chart type is compared side-by-side: "
                    "reference (left) vs agent output (right). "
                    "Only chart types that have a reference plot are shown."
                )

                import glob
                all_viz_code        = result.get("all_visualizations_code", {})
                chart_visual_scores = val.get("chart_visual_scores", {})
                selected_chart      = result.get("selected_chart_type", "")

                # Show official PNG images (e.g. iris) in a collapsible section
                ref_images_dir = os.path.join(
                    os.path.dirname(os.path.abspath(__file__)),
                    "reference_plots", ds_key
                )
                image_files = sorted(
                    f for ext in ("*.png", "*.jpg", "*.jpeg", "*.gif", "*.webp")
                    for f in glob.glob(os.path.join(ref_images_dir, ext))
                ) if os.path.isdir(ref_images_dir) else []
                if image_files:
                    with st.expander("Official documentation images"):
                        for img_path in image_files:
                            img_name = (os.path.splitext(os.path.basename(img_path))[0]
                                        .replace("_", " ").title())
                            st.caption(img_name)
                            st.image(img_path, use_container_width=True)

                # Load code-generated reference plots (tagged with chart type)
                try:
                    ref_figures = generate_reference_plot(
                        st.session_state.generated_data, ds_key
                    )
                except Exception as ref_err:
                    st.error(f"Could not load reference plots: {ref_err}")
                    ref_figures = []

                # Group reference figures by chart type
                ref_by_type: dict = {}   # chart_type -> [(title, fig)]
                for item in ref_figures:
                    r_title, r_fig = item[0], item[1]
                    r_type = item[2] if len(item) > 2 else "reference"
                    ref_by_type.setdefault(r_type, []).append((r_title, r_fig))

                # Helper: normalise chart type name for loose matching
                def _norm(s):
                    return s.lower().replace("_", "").replace(" ", "").replace("-", "")

                # For each reference chart type make a paired row
                for ref_type, ref_figs_for_type in ref_by_type.items():
                    ref_type_display = ref_type.capitalize()

                    # Find the agent's code block whose type matches ref_type
                    agent_code       = None
                    agent_chart_key  = None
                    for ct, lib_dict in all_viz_code.items():
                        if not isinstance(lib_dict, dict):
                            continue
                        if _norm(ct) == _norm(ref_type) or _norm(ref_type) in _norm(ct) or _norm(ct) in _norm(ref_type):
                            c = lib_dict.get("matplotlib")
                            if c:
                                agent_code      = c
                                agent_chart_key = ct
                                break

                    st.markdown(f"##### {ref_type_display}")
                    left, right = st.columns(2)

                    # ── LEFT: reference plot(s) for this type ────────────
                    with left:
                        st.markdown(f"**Reference — {ref_type_display}**")
                        for r_title, r_fig in ref_figs_for_type:
                            st.caption(r_title)
                            st.pyplot(r_fig)
                            plt.close(r_fig)

                    # ── RIGHT: agent plot for matching type ──────────────
                    with right:
                        st.markdown(f"**Agent — {ref_type_display}**")
                        if agent_code:
                            is_selected = (
                                agent_chart_key and
                                agent_chart_key.lower() == selected_chart.lower()
                            )
                            if is_selected:
                                st.caption("Selected chart")
                            try:
                                exec_globals = {
                                    "df": st.session_state.generated_data,
                                    "pd": pd,
                                    "plt": plt,
                                }
                                exec(agent_code.replace("plt.show()", ""), exec_globals)
                                agent_fig = plt.gcf()
                                st.pyplot(agent_fig)
                                plt.clf()


                            except Exception as agent_err:
                                st.warning(f"Could not render agent {ref_type} plot: {agent_err}")
                                with st.expander("Show code"):
                                    st.code(agent_code, language="python")
                        else:
                            st.info(f"Agent did not generate a {ref_type_display} chart.")

                    st.markdown("---")

    else:
        st.info("Run the pipeline in the **Pipeline Execution** tab first!")

# ============================================================================
# FOOTER
# ============================================================================
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: gray;'>
    <p>Multiagent Data Visualization System | Addressing 6 Research Gaps in Data Visualization</p>
    <p style='font-size: 12px;'>
        Gap 1: Domain Adaptability | Gap 2: Scalability | Gap 3: Human-AI Collaboration | 
        Gap 4: Integration | Gap 5: Technical Robustness | Gap 6: Evaluation Framework
    </p>
</div>
""", unsafe_allow_html=True)