import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import time
import random

# Import from other modules
from ui_components import (
    set_custom_css, 
    metric_card, 
    create_sidebar,
    header_section,
    create_tabs
)
from data_generators import (
    generate_physician_segments,
    generate_prescription_data,
    generate_key_drivers,
    generate_regional_data,
    generate_formulary_scenario_data,
    generate_message_testing_data
)
from api_simulator import simulate_claude_api_call

# Main application
def main():
    # Set page configuration
    st.set_page_config(
        page_title="PhysicianTwin™ - Digital Twin Simulator",
        page_icon="🧬",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Apply custom CSS
    set_custom_css()
    
    # Create sidebar and get parameters
    sidebar_params = create_sidebar()
    brand = sidebar_params["brand"]
    therapeutic_area = sidebar_params["therapeutic_area"]
    
    # Main content header
    header_section(brand, therapeutic_area)
    
    # Create tabs for different analyses
    tabs = create_tabs()
    
    # Generate simulated data
    df_segments = generate_physician_segments()
    df_prescriptions = generate_prescription_data()
    df_drivers = generate_key_drivers()
    df_regional = generate_regional_data()
    df_formulary = generate_formulary_scenario_data()
    df_messages = generate_message_testing_data()

    # Tab 1: Market Overview
    with tabs[0]:
        display_market_overview(df_prescriptions, df_regional, df_drivers)
    
    # Tab 2: Physician Segments
    with tabs[1]:
        display_physician_segments(df_segments)
    
    # Tab 3: Message Testing
    with tabs[2]:
        display_message_testing(df_messages)
    
    # Tab 4: Predictive Scenarios
    with tabs[3]:
        display_predictive_scenarios(df_formulary)
    
    # Tab 5: Interactive Twin
    with tabs[4]:
        display_interactive_twin()

def display_market_overview(df_prescriptions, df_regional, df_drivers):
    """Display Market Overview tab content"""
    st.markdown("## Market Performance Dashboard")
    st.markdown("Analyze current market position and trends for XenoGlip in the diabetes market")
    
    # Metrics row
    col1, col2, col3, col4 = st.columns(4)
    
    with col1:
        st.metric(label="Market Share", value="11.8%")
        
    with col2:
        st.metric(label="YoY Growth", value="+3.2%")
        
    with col3:
        st.metric(label="Formulary Access", value="42.3%")
        
    with col4:
        st.metric(label="Physician Awareness", value="68.7%")
    
    st.markdown("---")
    
    # Prescription trends chart
    st.markdown("### Prescription Trends")
    
    # Chart for prescription trends
    fig_rx = px.line(
        df_prescriptions, 
        x='Date', 
        y='Prescriptions', 
        color='Product',
        title='Weekly Prescription Volume',
        color_discrete_sequence=px.colors.qualitative.Bold
    )
    fig_rx.update_layout(
        xaxis_title="Date",
        yaxis_title="Total Prescriptions",
        legend_title="Product",
        height=450
    )
    st.plotly_chart(fig_rx, use_container_width=True)
    
    # Regional breakdown and key drivers in two columns
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("### Regional Performance")
        fig_regional = px.bar(
            df_regional,
            x='Region',
            y='Market Share',
            color='Growth Rate',
            color_continuous_scale='RdYlGn',
            title='Market Share by Region',
            text_auto='.1%'
        )
        fig_regional.update_layout(height=400)
        st.plotly_chart(fig_regional, use_container_width=True)
        
    with col2:
        st.markdown("### Key Performance Drivers")
        fig_drivers = px.bar(
            df_drivers[df_drivers['Segment'] == 'PCP'].sort_values('Importance', ascending=False).head(6),
            x='Importance',
            y='Driver',
            orientation='h',
            title='Key Drivers for PCPs',
            color='Importance',
            color_continuous_scale='Blues'
        )
        fig_drivers.update_traces(texttemplate='%{x:.0%}', textposition='outside')
        fig_drivers.update_layout(height=400)
        st.plotly_chart(fig_drivers, use_container_width=True)

def display_physician_segments(df_segments):
    """Display Physician Segments tab content"""
    st.markdown("## Physician Segment Analysis")
    st.markdown("Analyze behaviors and preferences across physician microsegments")
    
    # Segment selection
    selected_segments = st.multiselect(
        "Select segments to compare:",
        df_segments['Segment'].tolist(),
        default=df_segments['Segment'].tolist()[:3]
    )
    
    if not selected_segments:
        st.warning("Please select at least one segment to display.")
    else:
        filtered_df = df_segments[df_segments['Segment'].isin(selected_segments)]
        
        # Segment comparison chart
        st.markdown("### Segment Comparison")
        
        # Normalize prescribing volume to 0-1 range for radar chart
        radar_df = filtered_df.copy()
        max_rx = radar_df['Prescribing Volume'].max()
        min_rx = radar_df['Prescribing Volume'].min()
        radar_df['Normalized Prescribing'] = (radar_df['Prescribing Volume'] - min_rx) / (max_rx - min_rx) if max_rx > min_rx else 0.5
        
        # Radar chart for segment comparison
        categories = ['Normalized Prescribing', 'Digital Engagement', 'XenoGlip Affinity', 'Message Receptivity']
        
        fig_radar = go.Figure()
        
        for segment in selected_segments:
            segment_data = radar_df[radar_df['Segment'] == segment]
            values = segment_data[categories].values.flatten().tolist()
            # Add the first value again to close the loop
            values.append(values[0])
            
            fig_radar.add_trace(go.Scatterpolar(
                r=values,
                theta=categories + [categories[0]],
                fill='toself',
                name=segment
            ))
        
        fig_radar.update_layout(
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, 1]
                )),
            showlegend=True,
            height=500
        )
        
        st.plotly_chart(fig_radar, use_container_width=True)
        
        # Segment details table
        st.markdown("### Segment Details")
        
        # Format the dataframe for display
        display_df = filtered_df.copy()
        display_df['Digital Engagement'] = display_df['Digital Engagement'].map('{:.1%}'.format)
        display_df['XenoGlip Affinity'] = display_df['XenoGlip Affinity'].map('{:.1%}'.format)
        display_df['Message Receptivity'] = display_df['Message Receptivity'].map('{:.1%}'.format)
        
        st.dataframe(display_df, use_container_width=True)
        
        # Segment profiles
        st.markdown("### Detailed Segment Profiles")
        
        for segment in selected_segments:
            with st.expander(f"Profile: {segment}"):
                segment_data = df_segments[df_segments['Segment'] == segment].iloc[0]
                
                st.markdown(f"**Size:** {segment_data['Size']:,} physicians")
                st.markdown(f"**Prescribing Volume:** {segment_data['Prescribing Volume']} Rx/month (average)")
                st.markdown(f"**XenoGlip Affinity:** {segment_data['XenoGlip Affinity']:.1%}")
                
                st.markdown("#### Behavioral Characteristics")
                
                if "High Volume" in segment:
                    st.markdown("- Treats large numbers of diabetes patients")
                    st.markdown("- Typically relies on established treatment algorithms")
                    st.markdown("- Highly sensitive to formulary positioning")
                    st.markdown("- Values patient affordability and adherence")
                elif "Early Adopter" in segment:
                    st.markdown("- Quickly evaluates and adopts new treatments")
                    st.markdown("- Actively follows clinical trial data")
                    st.markdown("- Influenced by scientific evidence more than commercial messaging")
                    st.markdown("- Often participates in speaker bureaus and advisory boards")
                else:
                    st.markdown("- Moderate adoption timeline for new treatments")
                    st.markdown("- Evidence-based decision making")
                    st.markdown("- Balance of clinical and practical considerations")
                    st.markdown("- Responsive to peer influence and educational programs")

def display_message_testing(df_messages):
    """Display Message Testing tab content"""
    st.markdown("## Message Testing & Optimization")
    st.markdown("Evaluate message effectiveness across different physician segments")
    
    # Message testing chart
    st.markdown("### Message Effectiveness by Segment")
    
    messages_df = df_messages.pivot(index='Message', columns='Segment', values='Receptivity')
    
    fig_heatmap = px.imshow(
        messages_df,
        text_auto='.0%',
        aspect="auto",
        color_continuous_scale='Blues',
        title="Message Receptivity Heatmap"
    )
    fig_heatmap.update_layout(height=500)
    st.plotly_chart(fig_heatmap, use_container_width=True)
    
    # Message impact chart
    st.markdown("### Message Impact Analysis")
    
    # Calculate average receptivity across segments
    message_impact = df_messages.groupby('Message').agg({
        'Receptivity': 'mean',
        'Impact Score': 'mean'
    }).reset_index()
    
    fig_impact = px.scatter(
        message_impact,
        x='Receptivity',
        y='Impact Score',
        text='Message',
        size='Impact Score',
        color='Receptivity',
        color_continuous_scale='Viridis',
        title="Message Impact vs. Receptivity"
    )
    
    fig_impact.update_traces(textposition='top center')
    fig_impact.update_layout(height=500)
    st.plotly_chart(fig_impact, use_container_width=True)
    
    # Message optimization tool
    st.markdown("### Message Optimization Tool")
    st.markdown("Generate optimized messaging for specific physician segments")
    
    col1, col2 = st.columns(2)
    
    with col1:
        target_segment = st.selectbox(
            "Target Physician Segment",
            ["Primary Care Physicians", "Endocrinologists", "Cardiologists"]
        )
        
        message_focus = st.multiselect(
            "Message Focus Areas",
            ["Efficacy", "Safety", "Convenience", "Cost", "Guidelines", "Mechanism of Action"],
            default=["Efficacy", "Safety"]
        )
    
    with col2:
        clinical_context = st.selectbox(
            "Clinical Context",
            ["Second-line after metformin", "Add-on to current therapy", "Special populations", "Switch from competitor"]
        )
        
        message_tone = st.select_slider(
            "Message Tone",
            options=["Very Clinical", "Balanced", "Practice-Oriented"],
            value="Balanced"
        )
    
    if st.button("Generate Optimized Messages"):
        with st.spinner("Analyzing digital twin data and generating messages..."):
            time.sleep(2)  # Simulate processing time
            
            st.markdown("### Optimized Message Recommendations")
            
            if target_segment == "Primary Care Physicians":
                st.markdown("#### Primary Message (92% predicted effectiveness)")
                st.markdown("**XenoGlip delivers reliable A1C reduction with once-daily dosing and a proven safety profile, making it an ideal choice for your patients after metformin.**")
                
                st.markdown("#### Supporting Points:")
                st.markdown("• A1C reductions of 0.7-0.9% in clinical trials when added to metformin")
                st.markdown("• Established cardiovascular safety profile with no increased risk of major adverse cardiac events")
                st.markdown("• Simple once-daily dosing that doesn't require titration")
            
            elif target_segment == "Endocrinologists":
                st.markdown("#### Primary Message (87% predicted effectiveness)")
                st.markdown("**XenoGlip provides a flexible treatment option with established efficacy and minimal drug interactions, complementing your individualized approach to glycemic control.**")
                
                st.markdown("#### Supporting Points:")
                st.markdown("• Demonstrated efficacy across diverse patient populations in 25+ clinical trials")
                st.markdown("• Compatible with multiple antihyperglycemic agents for tailored therapy")
                st.markdown("• Minimal risk of drug-drug interactions with commonly prescribed medications")

def display_predictive_scenarios(df_formulary):
    """Display Predictive Scenarios tab content"""
    st.markdown("## Predictive Scenario Simulation")
    st.markdown("Explore the impact of different market events and strategic decisions")
    
    # Scenario selection
    scenario_type = st.radio(
        "Select Scenario Type:",
        ["Formulary Change Impact", "Competitive Launch", "Clinical Data Update", "Campaign Optimization"],
        horizontal=True
    )
    
    # Formulary Change Scenario
    if scenario_type == "Formulary Change Impact":
        # Create a grouped bar chart for formulary scenarios
        formulary_melted = df_formulary.copy()
        
        fig_formulary = px.bar(
            formulary_melted,
            x='Scenario',
            y='Value',
            color='Metric',
            barmode='group',
            title='Projected Impact of Formulary Changes',
            color_discrete_sequence=px.colors.qualitative.Safe
        )
        
        fig_formulary.update_layout(
            xaxis_title="Formulary Scenario",
            yaxis_title="Impact Value",
            yaxis_tickformat='.0%',
            height=500
        )
        
        st.plotly_chart(fig_formulary, use_container_width=True)
    
    # Other scenario types would be implemented similarly
    else:
        st.info(f"Selected scenario: {scenario_type}. Configure settings and run the simulation to see results.")

def display_interactive_twin():
    """Display Interactive Twin tab content"""
    st.markdown("## Interactive Physician Digital Twin")
    st.markdown("Ask questions to the digital twin model and receive AI-powered insights about physician behavior and preferences")
    
    # Create a text input for questions
    query = st.text_area(
        "Ask a question about physician behavior and prescribing patterns:",
        placeholder="Example: What factors drive primary care physicians to prescribe XenoGlip for diabetes?",
        height=100
    )
    
    # Add some example questions
    with st.expander("Example questions to ask the Digital Twin"):
        st.markdown("""
        - What messaging resonates best with endocrinologists for XenoGlip?
        - How would a formulary change from Tier 3 to Tier 2 impact prescribing?
        - What patient profile is most likely to receive XenoGlip from PCPs?
        - How do cardiologists view XenoGlip compared to other diabetes medications?
        """)
    
    # Additional filters for the query
    col1, col2 = st.columns(2)
    
    with col1:
        specialty_filter = st.multiselect(
            "Filter by specialty (optional):",
            ["Primary Care", "Endocrinology", "Cardiology", "NP/PA"],
            default=[]
        )
    
    with col2:
        region_filter = st.multiselect(
            "Filter by region (optional):",
            ["Northeast", "Southeast", "Midwest", "Southwest", "West"],
            default=[]
        )
    
    # Submit button for the query
    if st.button("Ask Digital Twin") and query:
        with st.spinner("Digital twin is analyzing physician data..."):
            # In a real app, this would make an API call to Claude
            # Here we'll simulate it with our function
            time.sleep(2)
            
            # Prepare the prompt with filters if applicable
            prompt = query
            if specialty_filter:
                prompt += f" Focus on {', '.join(specialty_filter)}."
            if region_filter:
                prompt += f" Consider regional factors in {', '.join(region_filter)}."
            
            # Get simulated response from Claude
            response = simulate_claude_api_call(prompt)
            
            # Display response in a nice format
            st.markdown("### Digital Twin Response")
            
            # Use a card-like container for the response
            with st.container():
                st.markdown(f'<div class="highlight">{response}</div>', unsafe_allow_html=True)
            
            # Display confidence metrics
            st.markdown("### Analysis Confidence")
            
            col1, col2, col3 = st.columns(3)
            
            with col1:
                # Simulate confidence scores
                data_confidence = random.uniform(0.7, 0.95)
                st.metric(
                    label="Data Confidence",
                    value=f"{data_confidence:.0%}"
                )
            
            with col2:
                response_consistency = random.uniform(0.75, 0.98)
                st.metric(
                    label="Response Consistency",
                    value=f"{response_consistency:.0%}"
                )
            
            with col3:
                physician_sample = random.randint(1500, 8000)
                st.metric(
                    label="Physician Sample Size",
                    value=f"{physician_sample:,}"
                )

if __name__ == "__main__":
    main()