analysis_viewer.py

"""
Analysis Viewer Component for Pharmaceutical Analytics Application

This module provides a dedicated view for watching the agents work in real-time,
displaying generated code, SQL queries, and visualization steps as they happen.
"""

import streamlit as st
import pandas as pd
import time
import plotly.graph_objects as go
from datetime import datetime
import json

def render_progress_bar(current_step, steps=None):
    """Render a progress bar for the current workflow state"""
    if steps is None:
        steps = ["planning", "data_collection", "analysis", "validation", "insights", "complete"]
    
    step_idx = steps.index(current_step) if current_step in steps else 0
    progress = (step_idx + 1) / len(steps)
    
    st.progress(progress)
    
    # Step indicators with more detailed descriptions
    cols = st.columns(5)
    
    step_descriptions = {
        "planning": "Decomposes problem & plans analysis approach",
        "data_collection": "Translates to SQL & retrieves data",
        "analysis": "Performs statistical analysis & modeling",
        "validation": "Validates results & checks quality",
        "insights": "Creates visualizations & recommendations"
    }
    
    with cols[0]:
        check = "✅" if step_idx >= 0 else "🔄"
        check = "🔄" if step_idx == 0 else check
        st.markdown(f"{check} **Planning**")
        if current_step == "planning":
            st.caption(step_descriptions["planning"])
    
    with cols[1]:
        check = "✅" if step_idx >= 1 else "⏳"
        check = "🔄" if step_idx == 1 else check
        st.markdown(f"{check} **Data Collection**")
        if current_step == "data_collection":
            st.caption(step_descriptions["data_collection"])
    
    with cols[2]:
        check = "✅" if step_idx >= 2 else "⏳"
        check = "🔄" if step_idx == 2 else check
        st.markdown(f"{check} **Analysis**")
        if current_step == "analysis":
            st.caption(step_descriptions["analysis"])
    
    with cols[3]:
        check = "✅" if step_idx >= 3 else "⏳"
        check = "🔄" if step_idx == 3 else check
        st.markdown(f"{check} **Validation**")
        if current_step == "validation":
            st.caption(step_descriptions["validation"])
    
    with cols[4]:
        check = "✅" if step_idx >= 4 else "⏳"
        check = "🔄" if step_idx == 4 else check
        st.markdown(f"{check} **Insights**")
        if current_step == "insights":
            st.caption(step_descriptions["insights"])


def render_planning_agent_view(state):
    """Render the Planning Agent activity view"""
    st.markdown("### 🧠 Planning Agent")
    st.markdown("The Planning Agent is decomposing the problem and creating an analysis plan.")
    
    # Show input alert
    alert_container = st.container(border=True)
    with alert_container:
        st.markdown("**Input Alert:**")
        st.info(state.get("alert", ""))
    
    # Thinking animation or finished plan
    if state.get("plan") is None:
        # Show thinking animation
        thinking_messages = [
            "Identifying required data sources...",
            "Determining appropriate analytical approaches...",
            "Creating task dependency graph...",
            "Designing validation strategy..."
        ]
        
        thinking_placeholder = st.empty()
        for i in range(len(thinking_messages)):
            thinking_placeholder.info(thinking_messages[i % len(thinking_messages)])
            time.sleep(0.5)
    else:
        # Show the created plan
        plan = state.get("plan")
        
        # Show problem statement
        st.markdown("#### 🎯 Problem Statement")
        st.markdown(plan.problem_statement)
        
        # Show required data sources in a table
        st.markdown("#### 📊 Required Data Sources")
        data_sources_df = pd.DataFrame(plan.required_data_sources)
        st.table(data_sources_df)
        
        # Show analysis approaches
        st.markdown("#### 🔍 Analysis Approaches")
        approaches_df = pd.DataFrame(plan.analysis_approaches)
        st.table(approaches_df)
        
        # Show expected insights
        st.markdown("#### 💡 Expected Insights")
        for i, insight in enumerate(plan.expected_insights):
            st.markdown(f"{i+1}. {insight}")


def render_data_agent_view(state):
    """Render the Data Agent activity view"""
    st.markdown("### 🗃️ Data Agent")
    st.markdown("The Data Agent is translating requirements into SQL queries and collecting data.")
    
    # Show data requests
    st.markdown("#### 📋 Data Requests")
    data_requests = state.get("data_requests", [])
    
    if not data_requests:
        st.info("Preparing data requests...")
    else:
        for i, request in enumerate(data_requests):
            with st.expander(f"Request {i+1}: {request.description}", expanded=True):
                st.markdown(f"**Tables:** {', '.join(request.tables)}")
                st.markdown(f"**Purpose:** {request.purpose}")
    
    # Show SQL generation and results
    data_sources = state.get("data_sources", {})
    if data_sources:
        st.markdown("#### 📊 Generated Data Sources")
        
        for source_id, source in data_sources.items():
            with st.expander(f"Data Source: {source.name}", expanded=True):
                # Show SQL query (mocked here - in real implementation you'd extract from the pipeline)
                if st.session_state.show_sql:
                    mock_sql = f"""
                    -- Query for {source.name}
                    SELECT 
                        r.region_name,
                        p.product_name,
                        strftime('%Y-%m', s.sale_date) as month,
                        SUM(s.units_sold) as total_units,
                        SUM(s.revenue) as total_revenue,
                        SUM(s.margin) as total_margin
                    FROM 
                        sales s
                    JOIN 
                        regions r ON s.region_id = r.region_id
                    JOIN
                        products p ON s.product_id = p.product_id
                    WHERE 
                        p.product_id = 'DRX'
                        AND s.sale_date >= date('now', '-90 days')
                    GROUP BY 
                        r.region_name, p.product_name, month
                    ORDER BY 
                        r.region_name, month;
                    """
                    st.code(mock_sql, language="sql")
                
                # Show data preview
                st.markdown("**Data Preview:**")
                st.dataframe(source.content.head(5), use_container_width=True)
                st.caption(f"Shape: {source.content.shape[0]} rows, {source.content.shape[1]} columns")


def render_analytics_agent_view(state):
    """Render the Analytics Agent activity view"""
    st.markdown("### 📊 Analytics Agent")
    st.markdown("The Analytics Agent is performing statistical analysis and modeling.")
    
    # Show analytics requests
    analysis_requests = state.get("analysis_requests", [])
    if not analysis_requests:
        st.info("Preparing analysis requests...")
    else:
        st.markdown("#### 📋 Analysis Requests")
        for i, request in enumerate(analysis_requests):
            st.markdown(f"**Request {i+1}:** {request.description} ({request.analysis_type})")
    
    # Show analysis results
    analysis_results = state.get("analysis_results", {})
    if analysis_results:
        st.markdown("#### 📈 Analysis Results")
        
        for analysis_id, result in analysis_results.items():
            with st.expander(f"Analysis: {result.name}", expanded=True):
                # Show generated Python code
                if st.session_state.show_code and result.code:
                    st.markdown("**Generated Python Code:**")
                    st.code(result.code, language="python")
                
                # Show insights
                if result.insights:
                    st.markdown("**Key Findings:**")
                    for insight in result.insights:
                        st.markdown(f"- **{insight.get('finding', '')}**: {insight.get('details', '')}")
                
                # Show metrics
                if result.metrics:
                    st.markdown("**Metrics:**")
                    metrics_df = pd.DataFrame([result.metrics])
                    st.table(metrics_df.T)
                
                # Show attribution
                if result.attribution:
                    st.markdown("**Attribution Analysis:**")
                    fig = go.Figure([
                        go.Bar(
                            x=list(result.attribution.values()),
                            y=list(result.attribution.keys()),
                            orientation='h'
                        )
                    ])
                    fig.update_layout(
                        title="Factor Attribution",
                        xaxis_title="Attribution (%)",
                        height=300
                    )
                    st.plotly_chart(fig, use_container_width=True)


def render_qa_agent_view(state):
    """Render the QA Agent activity view"""
    st.markdown("### 🔍 QA Agent")
    st.markdown("The QA Agent is validating the analysis results for accuracy and completeness.")
    
    # Show validation requests
    validation_requests = state.get("validation_requests", [])
    if not validation_requests:
        st.info("Preparing validation requests...")
    
    # Show validation results
    validation_results = state.get("validation_results", {})
    if validation_results:
        st.markdown("#### 🧪 Validation Results")
        
        for validation_id, validation in validation_results.items():
            with st.expander(f"Validation: {validation_id}", expanded=True):
                # Show validation scores
                scores_col1, scores_col2, scores_col3 = st.columns(3)
                
                with scores_col1:
                    st.metric("Data Quality", f"{validation.data_quality_score:.0%}")
                
                with scores_col2:
                    st.metric("Analysis Quality", f"{validation.analysis_quality_score:.0%}")
                
                with scores_col3:
                    st.metric("Insight Quality", f"{validation.insight_quality_score:.0%}")
                
                # Show validation checks as a table
                st.markdown("**Validation Checks:**")
                checks_df = pd.DataFrame(validation.validation_checks)
                checks_df = checks_df[["check", "result", "details", "score"]]
                st.table(checks_df)
                
                # Show recommendations
                if validation.recommendations:
                    st.markdown("**Recommendations:**")
                    for rec in validation.recommendations:
                        st.markdown(f"- {rec}")
                
                # Show critical issues if any
                if validation.critical_issues:
                    st.error("**Critical Issues:**")
                    for issue in validation.critical_issues:
                        st.markdown(f"- {issue}")


def render_insights_agent_view(state):
    """Render the Insights Agent activity view"""
    st.markdown("### 💡 Insights Agent")
    st.markdown("The Insights Agent is generating visualizations and actionable recommendations.")
    
    # Show insight requests
    insight_requests = state.get("insight_requests", [])
    if not insight_requests:
        st.info("Preparing insight requests...")
    
    # Show insight cards
    insight_cards = state.get("insight_cards", {})
    visualizations = state.get("visualizations", [])
    
    if insight_cards:
        st.markdown("#### 📊 Generated Insights")
        
        for card_id, card in insight_cards.items():
            with st.expander(f"{card.title}", expanded=True):
                st.markdown(card.description)
                
                # Display visualizations if available
                if card.charts and visualizations:
                    # Create sample visualizations based on chart names
                    st.markdown("**Visualizations:**")
                    cols = st.columns(min(len(card.charts), 2))
                    
                    for i, chart_name in enumerate(card.charts[:2]):
                        with cols[i % 2]:
                            if "trend" in chart_name.lower():
                                # Create sample time series for demonstration
                                months = pd.date_range(start='2023-01-01', periods=12, freq='M')
                                sales = [1200, 1250, 1300, 1400, 1500, 1600, 1550, 1500, 1450, 1300, 1200, 1150]
                                targets = [1200, 1250, 1300, 1350, 1400, 1450, 1500, 1550, 1600, 1650, 1700, 1750]
                                
                                fig = go.Figure()
                                fig.add_trace(go.Scatter(
                                    x=months, y=sales, mode='lines+markers', name='Actual Sales',
                                    line=dict(color='blue', width=2)
                                ))
                                fig.add_trace(go.Scatter(
                                    x=months, y=targets, mode='lines+markers', name='Target',
                                    line=dict(color='green', width=2, dash='dash')
                                ))
                                
                                # Add annotation
                                fig.add_vline(
                                    x=months[9], line_dash="dash", line_color="red",
                                    annotation_text="Competitor Launch"
                                )
                                
                                fig.update_layout(
                                    title="Sales Trend Analysis",
                                    height=300
                                )
                                
                                st.plotly_chart(fig, use_container_width=True)
                            
                            elif "competitor" in chart_name.lower():
                                # Create sample competitor visualization
                                st.markdown(f"**{chart_name}**")
                                st.text("Visualization being generated...")
                            
                            else:
                                st.markdown(f"**{chart_name}**")
                                st.text("Visualization being generated...")
                
                # Display key findings
                if card.key_findings:
                    st.markdown("**Key Findings:**")
                    for finding in card.key_findings:
                        st.markdown(f"- **{finding.get('finding', '')}**: {finding.get('details', '')}")
                        if 'impact' in finding:
                            st.markdown(f"  *Impact: {finding.get('impact', '')}*")
                
                # Display action items
                if card.action_items:
                    st.markdown("**Recommended Actions:**")
                    action_df = []
                    for action in card.action_items:
                        action_df.append({
                            "Action": action.get("action", ""),
                            "Owner": action.get("owner", ""),
                            "Timeline": action.get("timeline", ""),
                            "Priority": action.get("priority", "Medium")
                        })
                    st.table(pd.DataFrame(action_df))


def render_complete_analysis_view(state):
    """Render a complete view of the analysis results"""
    st.markdown("### ✅ Analysis Complete")
    
    # Display the alert
    st.markdown("#### 📱 Alert Analyzed")
    st.info(state["alert"])
    
    # Display the problem statement
    if state.get("plan"):
        st.markdown("#### 🎯 Problem Analysis")
        st.markdown(state["plan"].problem_statement)
    
    # Display insight cards
    st.markdown("---")
    st.markdown("### 💡 Key Insights")
    
    insight_cards = state.get("insight_cards", {})
    visualizations = state.get("visualizations", [])
    
    if not insight_cards:
        st.warning("No insights were generated. Please check the logs for errors.")
    else:
        for card_id, card in insight_cards.items():
            # Create a card-like container
            with st.container():
                st.subheader(card.title)
                st.markdown(card.description)
                
                # Display visualizations if available
                if card.charts and len(visualizations) > 0:
                    # Create sample visualizations
                    cols = st.columns(min(len(card.charts), 2))
                    
                    for i, chart_name in enumerate(card.charts[:2]):
                        with cols[i % 2]:
                            # Create a sample chart based on chart name
                            if "trend" in chart_name.lower():
                                # Create sample time series
                                months = pd.date_range(start='2023-01-01', periods=12, freq='M')
                                sales = [1200, 1250, 1300, 1400, 1500, 1600, 1550, 1500, 1450, 1300, 1200, 1150]
                                targets = [1200, 1250, 1300, 1350, 1400, 1450, 1500, 1550, 1600, 1650, 1700, 1750]
                                
                                fig = go.Figure()
                                fig.add_trace(go.Scatter(
                                    x=months, y=sales, mode='lines+markers', name='Actual Sales',
                                    line=dict(color='blue', width=2)
                                ))
                                fig.add_trace(go.Scatter(
                                    x=months, y=targets, mode='lines+markers', name='Target',
                                    line=dict(color='green', width=2, dash='dash')
                                ))
                                
                                # Add competitor launch annotation
                                fig.add_vline(
                                    x=months[9], line_dash="dash", line_color="red",
                                    annotation_text="Competitor Launch"
                                )
                                
                                fig.update_layout(
                                    title="DrugX Sales Trend",
                                    xaxis_title="Month",
                                    yaxis_title="Sales ($K)",
                                    height=300
                                )
                                
                                st.plotly_chart(fig, use_container_width=True)
                            
                            elif "competitor" in chart_name.lower():
                                # Create sample competitor comparison
                                fig = go.Figure()
                                
                                # Market share data
                                months = pd.date_range(start='2023-01-01', periods=12, freq='M')
                                drugx_share = [65, 64, 66, 67, 68, 67, 66, 65, 64, 58, 54, 50]
                                competitor_share = [0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 10, 15]
                                other_share = [35, 36, 34, 33, 32, 33, 34, 35, 36, 37, 36, 35]
                                
                                fig.add_trace(go.Bar(
                                    x=months, y=drugx_share, name='DrugX',
                                    marker_color='blue'
                                ))
                                fig.add_trace(go.Bar(
                                    x=months, y=competitor_share, name='CompDrug2',
                                    marker_color='red'
                                ))
                                fig.add_trace(go.Bar(
                                    x=months, y=other_share, name='Others',
                                    marker_color='gray'
                                ))
                                
                                fig.update_layout(
                                    title="Market Share Comparison",
                                    xaxis_title="Month",
                                    yaxis_title="Market Share (%)",
                                    barmode='stack',
                                    height=300
                                )
                                
                                st.plotly_chart(fig, use_container_width=True)
                            
                            elif "supply" in chart_name.lower():
                                # Create sample supply chain visualization
                                fig = go.Figure()
                                
                                # Inventory data
                                months = pd.date_range(start='2023-01-01', periods=12, freq='M')
                                inventory = [40, 38, 42, 45, 43, 41, 39, 37, 35, 25, 20, 18]
                                stockouts = [0, 0, 0, 0, 0, 0, 0, 0, 5, 15, 22, 25]
                                
                                fig.add_trace(go.Scatter(
                                    x=months, y=inventory, mode='lines+markers', name='Inventory Days',
                                    line=dict(color='blue', width=2)
                                ))
                                
                                fig.add_trace(go.Bar(
                                    x=months, y=stockouts, name='Stockout %',
                                    marker_color='red'
                                ))
                                
                                fig.update_layout(
                                    title="Supply Chain Metrics",
                                    xaxis_title="Month",
                                    yaxis_title="Inventory Days / Stockout %",
                                    height=300
                                )
                                
                                st.plotly_chart(fig, use_container_width=True)
                            
                            else:
                                # Create a generic chart placeholder
                                st.image("https://img.icons8.com/fluency/240/000000/graph.png", width=100)
                                st.markdown(f"*{chart_name}*")
                
                # Display key findings
                if card.key_findings:
                    st.markdown("#### Key Findings")
                    for i, finding in enumerate(card.key_findings):
                        expander = st.expander(f"{finding.get('finding', 'Finding')}")
                        with expander:
                            st.markdown(f"**Details:** {finding.get('details', '')}")
                            if 'evidence' in finding:
                                st.markdown(f"**Evidence:** {finding.get('evidence', '')}")
                            if 'impact' in finding:
                                st.markdown(f"**Impact:** {finding.get('impact', '')}")
                
                # Display metrics
                if card.metrics:
                    st.markdown("#### Key Metrics")
                    metric_cols = st.columns(min(len(card.metrics), 4))
                    for i, (metric_name, metric_value) in enumerate(card.metrics.items()):
                        with metric_cols[i % len(metric_cols)]:
                            st.metric(
                                label=metric_name.replace('_', ' ').title(),
                                value=metric_value
                            )
                
                # Display action items
                if card.action_items:
                    st.markdown("#### Recommended Actions")
                    for i, action in enumerate(card.action_items):
                        priority = action.get('priority', 'Medium')
                        priority_color = {
                            'High': 'red',
                            'Medium': 'orange',
                            'Low': 'blue'
                        }.get(priority, 'gray')
                        
                        st.markdown(f"**{i+1}. {action.get('action', 'Action')}** "
                                   f"<span style='color:{priority_color};'>[{priority}]</span>", 
                                   unsafe_allow_html=True)
                        
                        action_cols = st.columns(3)
                        with action_cols[0]:
                            st.markdown(f"**Owner:** {action.get('owner', 'TBD')}")
                        with action_cols[1]:
                            st.markdown(f"**Timeline:** {action.get('timeline', 'TBD')}")
                        with action_cols[2]:
                            st.markdown(f"**Expected Impact:** {action.get('expected_impact', 'TBD')}")
                
                st.markdown("---")


def render_debug_info(state):
    """Render debug information about the workflow state"""
    with st.expander("🔧 Debug Information", expanded=False):
        # Show workflow status
        st.markdown(f"**Current Status:** {state.get('status', 'Unknown')}")
        
        # Show state summary
        state_summary = {
            "alert": state.get("alert", "None"),
            "data_requests": len(state.get("data_requests", [])),
            "data_sources": len(state.get("data_sources", {})),
            "analysis_requests": len(state.get("analysis_requests", [])),
            "analysis_results": len(state.get("analysis_results", {})),
            "validation_requests": len(state.get("validation_requests", [])),
            "validation_results": len(state.get("validation_results", {})),
            "insight_requests": len(state.get("insight_requests", [])),
            "insight_cards": len(state.get("insight_cards", {})),
            "visualizations": len(state.get("visualizations", [])),
        }
        
        st.json(state_summary)
        
        # Show logs
        if "logs" in state and state["logs"]:
            st.markdown("**Workflow Logs:**")
            for log in state["logs"]:
                timestamp = log.get("timestamp", "")
                message = log.get("message", "")
                log_type = log.get("type", "info")
                
                if log_type == "error":
                    st.error(f"{timestamp}: {message}")
                else:
                    st.text(f"{timestamp}: {message}")


def render_analysis_viewer(state=None):
    """Main function to render the analysis viewer"""
    # If no state is provided, use the session state
    if state is None:
        state = st.session_state.workflow_state
    
    # If no workflow state, show a message
    if state is None:
        st.info("No analysis is currently running. Start an analysis from the main tab.")
        return
    
    # Display header with current status
    current_step = state.get("status", "planning")
    st.markdown(f"## 🔬 Analysis Progress: {current_step.upper()}")
    
    # Show progress bar
    render_progress_bar(current_step)
    
    # Show page tabs for different views
    tab_names = ["Live View", "Planning", "Data", "Analysis", "Validation", "Insights", "Debug"]
    tabs = st.tabs(tab_names)
    
    # Live View Tab - shows current active agent
    with tabs[0]:
        if current_step == "planning":
            render_planning_agent_view(state)
        elif current_step == "data_collection":
            render_data_agent_view(state)
        elif current_step == "analysis":
            render_analytics_agent_view(state)
        elif current_step == "validation":
            render_qa_agent_view(state)
        elif current_step == "insights":
            render_insights_agent_view(state)
        elif current_step == "complete":
            render_complete_analysis_view(state)
        elif current_step == "error":
            st.error(f"Analysis failed: {state.get('error', 'Unknown error')}")
            if "error_details" in st.session_state:
                for i, error in enumerate(st.session_state.error_details):
                    with st.expander(f"Error Detail {i+1}", expanded=True):
                        st.code(error)
    
    # Planning Tab
    with tabs[1]:
        render_planning_agent_view(state)
    
    # Data Tab
    with tabs[2]:
        render_data_agent_view(state)
    
    # Analysis Tab
    with tabs[3]:
        render_analytics_agent_view(state)
    
    # Validation Tab
    with tabs[4]:
        render_qa_agent_view(state)
    
    # Insights Tab
    with tabs[5]:
        render_insights_agent_view(state)
    
    # Debug Tab
    with tabs[6]:
        render_debug_info(state)
    
    # Auto refresh if analysis is not complete
    if current_step not in ["complete", "error"]:
        time.sleep(1)
        st.rerun()


if __name__ == "__main__":
    # For testing standalone
    st.title("Analysis Viewer (Test Mode)")
    st.markdown("This is a standalone test of the Analysis Viewer component.")
    
    # Create mock state
    mock_state = {
        "alert": "Sales of DrugX down 15% in Northeast region over past 30 days compared to forecast.",
        "status": "planning",
        "data_requests": [],
        "data_sources": {},
        "logs": [
            {"timestamp": datetime.now().isoformat(), "message": "Workflow initialized", "type": "info"}
        ]
    }
    
    render_analysis_viewer(mock_state)