import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from collections import Counter

def create_world_map(docs_df):
    """Create interactive world map showing study distribution"""
    if docs_df.empty:
        return None
    
    # Count studies by country
    country_counts = Counter()
    for countries_str in docs_df['study_countries'].dropna():
        if isinstance(countries_str, str) and countries_str.lower() != 'nan':
            # Split multiple countries
            countries = [c.strip() for c in countries_str.split(',')]
            for country in countries:
                country_counts[country] += 1
    
    if not country_counts:
        return None
    
    # Create choropleth map
    countries = list(country_counts.keys())
    counts = list(country_counts.values())
    
    fig = go.Figure(data=go.Choropleth(
        locations=countries,
        z=counts,
        locationmode='country names',
        colorscale='Viridis',
        text=countries,
        hovertemplate='<b>%{text}</b><br>Studies: %{z}<extra></extra>',
        colorbar_title="Number of Studies"
    ))
    
    fig.update_layout(
        title={
            'text': '🌍 Global Research Coverage',
            'x': 0.5,
            'font': {'size': 20}
        },
        geo=dict(
            showframe=False,
            showcoastlines=True,
            projection_type='equirectangular'
        ),
        height=500
    )
    
    return fig

def create_sector_analysis(docs_df):
    """Create sector distribution charts"""
    if docs_df.empty:
        return None, None
    
    # Sector distribution
    sector_counts = docs_df['world_bank_sector'].value_counts().head(10)
    
    fig1 = px.bar(
        x=sector_counts.values, 
        y=sector_counts.index,
        orientation='h',
        title="📊 Research by World Bank Sector",
        labels={'x': 'Number of Studies', 'y': 'Sector'},
        color=sector_counts.values,
        color_continuous_scale='viridis'
    )
    fig1.update_layout(height=400, showlegend=False)
    
    # Research design pie chart
    design_counts = docs_df['research_design'].value_counts().head(8)
    
    fig2 = px.pie(
        values=design_counts.values,
        names=design_counts.index,
        title="🔬 Research Design Distribution",
        color_discrete_sequence=px.colors.qualitative.Set3
    )
    fig2.update_traces(textposition='inside', textinfo='percent+label')
    fig2.update_layout(height=400)
    
    return fig1, fig2

def create_methodology_dashboard(docs_df):
    """Create methodology analysis dashboard"""
    if docs_df.empty:
        return None
    
    # Create subplot figure
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=('Sample Size Distribution', 'Rigor Scores', 
                       'Data Collection Methods', 'Quality Indicators'),
        specs=[[{"secondary_y": False}, {"secondary_y": False}],
               [{"secondary_y": False}, {"secondary_y": False}]]
    )
    
    # Sample size histogram
    sample_sizes = pd.to_numeric(docs_df['sample_size'], errors='coerce').dropna()
    if not sample_sizes.empty:
        fig.add_trace(
            go.Histogram(x=sample_sizes, name="Sample Size", nbinsx=20),
            row=1, col=1
        )
    
    # Rigor scores
    rigor_scores = pd.to_numeric(docs_df['rigor_score'], errors='coerce').dropna()
    if not rigor_scores.empty:
        fig.add_trace(
            go.Histogram(x=rigor_scores, name="Rigor Score", nbinsx=10),
            row=1, col=2
        )
    
    # Data collection methods
    data_methods = docs_df['data_collection_method'].value_counts().head(8)
    if not data_methods.empty:
        fig.add_trace(
            go.Bar(x=data_methods.values, y=data_methods.index, 
                   orientation='h', name="Data Methods"),
            row=2, col=1
        )
    
    # Quality indicators (RCT, Validation, etc.)
    quality_data = []
    for col in ['has_randomization', 'has_validation', 'has_mixed_methods']:
        if col in docs_df.columns:
            true_count = (docs_df[col] == 'true').sum()
            quality_data.append((col.replace('has_', '').title(), true_count))
    
    if quality_data:
        labels, values = zip(*quality_data)
        fig.add_trace(
            go.Bar(x=list(labels), y=list(values), name="Quality Features"),
            row=2, col=2
        )
    
    fig.update_layout(
        height=800,
        title_text="📈 Methodology Dashboard",
        title_x=0.5,
        showlegend=False
    )
    
    return fig

def filter_studies(docs_df, countries, sectors, min_year, max_year, has_rct, min_sample_size):
    """Filter and display studies based on criteria"""
    if docs_df.empty:
        return "No data available"
    
    filtered = docs_df.copy()
    
    # Apply filters
    if countries:
        country_mask = filtered['study_countries'].str.contains('|'.join(countries), case=False, na=False)
        filtered = filtered[country_mask]
    
    if sectors:
        sector_mask = filtered['world_bank_sector'].isin(sectors)
        filtered = filtered[sector_mask]
    
    if min_year:
        year_col = pd.to_numeric(filtered['publication_year'], errors='coerce')
        filtered = filtered[year_col >= min_year]
    
    if max_year:
        year_col = pd.to_numeric(filtered['publication_year'], errors='coerce')
        filtered = filtered[year_col <= max_year]
    
    if has_rct:
        filtered = filtered[filtered['has_randomization'] == 'true']
    
    if min_sample_size:
        sample_col = pd.to_numeric(filtered['sample_size'], errors='coerce')
        filtered = filtered[sample_col >= min_sample_size]
    
    # Display results
    if filtered.empty:
        return "No studies match the selected criteria."
    
    # Create summary
    result = f"## 🔍 Filtered Results: {len(filtered)} studies\n\n"
    
    # Show sample of results
    display_cols = ['title', 'authors', 'publication_year', 'study_countries', 
                   'world_bank_sector', 'research_design', 'sample_size']
    available_cols = [col for col in display_cols if col in filtered.columns]
    
    sample_df = filtered[available_cols].head(10)
    result += sample_df.to_markdown(index=False)
    
    if len(filtered) > 10:
        result += f"\n\n*... and {len(filtered) - 10} more studies*"
    
    return result

def get_unique_values(docs_df):
    """Extract unique countries and sectors for dropdowns"""
    countries_list = []
    sectors_list = []
    
    if not docs_df.empty:
        # Extract unique countries
        for countries_str in docs_df['study_countries'].dropna():
            if isinstance(countries_str, str) and countries_str.lower() != 'nan':
                countries = [c.strip() for c in countries_str.split(',')]
                countries_list.extend(countries)
        countries_list = sorted(list(set(countries_list)))
        
        # Extract unique sectors
        sectors_list = sorted(docs_df['world_bank_sector'].dropna().unique().tolist())
    
    return countries_list, sectors_list