Py.Cafe

"""
Hofstede Cultural Dimensions - Agent-Based Model with Cultural Distance Calculator
Using Mesa 3.0 and Solara for interactive world map visualization
"""
import pandas as pd
import mesa
import solara
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

# Try to import geopandas
try:
    import geopandas as gpd
    GEOPANDAS_AVAILABLE = True
except ImportError:
    GEOPANDAS_AVAILABLE = False


class CountryAgent(mesa.Agent):
    """An agent representing a country with Hofstede cultural dimensions."""
    def __init__(self, model, country_name, country_code, pdi, idv, mas, uai, ltowvs, ivr):
        super().__init__(model)
        self.country_name = country_name
        self.country_code = country_code  # Standard ISO 3166-1 alpha-3
        self.pdi = pdi
        self.idv = idv
        self.mas = mas
        self.uai = uai
        self.ltowvs = ltowvs
        self.ivr = ivr

    def step(self):
        pass

    def get_dimension_value(self, dimension):
        return {
            'pdi': self.pdi,
            'idv': self.idv,
            'mas': self.mas,
            'uai': self.uai,
            'ltowvs': self.ltowvs,
            'ivr': self.ivr
        }.get(dimension)


class HofstedeModel(mesa.Model):
    """Model containing country agents with cultural dimensions."""
    def __init__(self, data_path):
        super().__init__()  # Initializes self.agents as AgentSet automatically
        # Load data as strings first to handle blanks/spaces
        df = pd.read_csv(data_path, dtype=str)
        
        # Convert numeric columns safely: blanks become NaN
        numeric_cols = ['pdi', 'idv', 'mas', 'uai', 'ltowvs', 'ivr']
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # Create agents — Mesa auto-adds them to self.agents
        for _, row in df.iterrows():
            if pd.isna(row['ctr']) or pd.isna(row['country']):
                continue
            # Use ISO code directly from 'ctr' column
            CountryAgent(
                self,
                country_name=row['country'],
                country_code=row['ctr'],
                pdi=row['pdi'],
                idv=row['idv'],
                mas=row['mas'],
                uai=row['uai'],
                ltowvs=row['ltowvs'],
                ivr=row['ivr']
            )

    def step(self):
        self.agents.do("step")

    def get_dimension_data(self, dimension):
        data = {}
        for agent in self.agents:
            value = agent.get_dimension_value(dimension)
            if value is not None and not np.isnan(value):
                data[agent.country_code] = value
        return data

    def calculate_cultural_distances(self, ref_country_code, weights):
        """Calculate weighted Euclidean cultural distances."""
        ref_agent = next((a for a in self.agents if a.country_code == ref_country_code), None)
        if ref_agent is None:
            return {}

        dimensions = ['pdi', 'idv', 'mas', 'uai', 'ltowvs', 'ivr']
        distances = {}

        for agent in self.agents:
            if agent.country_code == ref_country_code:
                continue

            valid = True
            squared_diff = 0.0
            for dim in dimensions:
                ref_val = ref_agent.get_dimension_value(dim)
                comp_val = agent.get_dimension_value(dim)
                if ref_val is None or comp_val is None or np.isnan(ref_val) or np.isnan(comp_val):
                    valid = False
                    break
                diff = ref_val - comp_val
                squared_diff += weights[dim] * (diff ** 2)

            if valid:
                distances[agent.country_code] = np.sqrt(squared_diff)

        return distances


def create_world_map(model, dimension=None, distance_data=None, ref_country_name=None):
    """Create world map for dimension or cultural distances."""
    dimension_info = {
        'pdi': {'name': 'Power Distance Index', 'description': 'Acceptance of unequal power distribution'},
        'idv': {'name': 'Individualism', 'description': 'Independence vs. interdependence'},
        'mas': {'name': 'Masculinity', 'description': 'Achievement vs. cooperation orientation'},
        'uai': {'name': 'Uncertainty Avoidance Index', 'description': 'Tolerance for ambiguity'},
        'ltowvs': {'name': 'Long Term Orientation', 'description': 'Future vs. present/past focus'},
        'ivr': {'name': 'Indulgence vs Restraint', 'description': 'Gratification vs. control of desires'}
    }
    
    fig = Figure(figsize=(16, 8), dpi=100)
    ax = fig.add_subplot(111)

    if GEOPANDAS_AVAILABLE:
        try:
            # Use built-in dataset (no internet required!)
            world = gpd.read_file(gpd.datasets.get_path('naturalearth_lowres'))

            if distance_data is not None:
                world['dimension_value'] = world['iso_a3'].map(distance_data)
                title = f"Cultural Distance from {ref_country_name}"
                legend_label = "Cultural Distance"
            else:
                dimension_data = model.get_dimension_data(dimension)
                world['dimension_value'] = world['iso_a3'].map(dimension_data)
                dim_info = dimension_info.get(dimension, {'name': dimension.upper(), 'description': ''})
                title = f"Hofstede's {dim_info['name']} ({dimension.upper()})\n{dim_info['description']}"
                legend_label = f"{dim_info['name']} Score"

            valid_vals = world['dimension_value'].dropna()
            if len(valid_vals) > 0:
                vmin, vmax = valid_vals.min(), valid_vals.max()
                cmap = LinearSegmentedColormap.from_list(
                    "blue_red", ["#08306b", "#f0f0f0", "#a50f15"], N=256
                )
                world[world['dimension_value'].notna()].plot(
                    column='dimension_value',
                    ax=ax,
                    cmap=cmap,
                    vmin=vmin,
                    vmax=vmax,
                    edgecolor='#555555',
                    linewidth=0.5,
                    legend=True,
                    legend_kwds={
                        'label': legend_label,
                        'orientation': 'horizontal',
                        'shrink': 0.6,
                        'pad': 0.02
                    }
                )

            # Show missing countries as WHITE
            world[world['dimension_value'].isna()].plot(
                ax=ax,
                color='white',
                edgecolor='#555555', 
                linewidth=0.5
            )

            ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
            ax.set_axis_off()
            fig.text(0.5, 0.02, 'Source: Hofstede Insights | Map: Natural Earth',
                     ha='center', fontsize=9, style='italic', color='#666666')

        except Exception as e:
            ax.text(0.5, 0.5,
                   f"Map error:\n{str(e)}",
                   ha='center', va='center', transform=ax.transAxes,
                   fontsize=12, bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.7))
            ax.set_axis_off() 
    else:
        ax.text(0.5, 0.5,
                "Install geopandas for map:\npip install geopandas",
               ha='center', va='center', transform=ax.transAxes,
               fontsize=14, bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        ax.set_axis_off()

    fig.tight_layout()
    return fig


def create_statistics_table(model, dimension):
    """Create statistics HTML table."""
    dimension_data = model.get_dimension_data(dimension)
    if not dimension_data:
        return "<div style='padding: 10px;'><p>No data available</p></div>"
    
    values = list(dimension_data.values())
    stats = {
        'Countries': len(values),
        'Minimum': f"{min(values):.1f}",
        'Maximum': f"{max(values):.1f}",
        'Average': f"{np.mean(values):.1f}",
        'Median': f"{np.median(values):.1f}",
        'Std Dev': f"{np.std(values):.1f}"
    }

    agents_dict = {agent.country_code: agent.country_name for agent in model.agents}
    sorted_countries = sorted(dimension_data.items(), key=lambda x: x[1], reverse=True)
    top_3 = [(agents_dict.get(code, code), value) for code, value in sorted_countries[:3]]
    bottom_3 = [(agents_dict.get(code, code), value) for code, value in sorted_countries[-3:]]
    bottom_3.reverse()

    html = "<div style='padding: 10px;'>"
    html += "<h3>Statistics</h3>"
    html += "<table style='width: 100%; border-collapse: collapse;'>"
    for key, value in stats.items():
        html += f"<tr><td style='padding: 5px; border-bottom: 1px solid #ddd;'><strong>{key}:</strong></td><td style='padding: 5px; border-bottom: 1px solid #ddd;'>{value}</td></tr>"
    html += "</table>"

    html += "<h3 style='margin-top: 20px;'>Top 3 Countries</h3><ol>"
    for country, value in top_3:
        html += f"<li>{country}: {value:.1f}</li>"
    html += "</ol>"

    html += "<h3>Bottom 3 Countries</h3><ol>"
    for country, value in bottom_3:
        html += f"<li>{country}: {value:.1f}</li>"
    html += "</ol></div>"

    return html


@solara.component
def WorldMapVisualization(model):
    dimension_display = {
        'pdi': 'Power Distance',
        'idv': 'Individualism',
        'mas': 'Masculinity',
        'uai': 'Uncertainty Avoidance',
        'ltowvs': 'Long Term Orientation',
        'ivr': 'Indulgence'
    }
    
    # Get countries with complete data for distance calculation
    complete_countries = []
    dims = ['pdi', 'idv', 'mas', 'uai', 'ltowvs', 'ivr']
    country_code_to_name = {}

    for agent in model.agents:
        country_code_to_name[agent.country_code] = agent.country_name
        has_all_data = True
        for dim in dims:
            val = agent.get_dimension_value(dim)
            if val is None or np.isnan(val):
                has_all_data = False
                break
        if has_all_data:
            complete_countries.append((agent.country_code, agent.country_name))

    complete_countries.sort(key=lambda x: x[1])  # Sort by name

    selected_dimension = solara.use_reactive('pdi')
    show_distance = solara.use_reactive(False)

    # Initialize with first country code (string, not tuple)
    initial_country = complete_countries[0][0] if complete_countries else None
    selected_country = solara.use_reactive(initial_country)

    weights = solara.use_reactive({d: 1/6 for d in dims})

    # Calculate map and stats
    if show_distance.value and selected_country.value:
        # ✅ SAFETY CHECK: Ensure selected_country.value is a string
        ref_country_code = selected_country.value
        if isinstance(ref_country_code, list):
            # If it's a list, take the first element (the code)
            ref_country_code = ref_country_code[0] if ref_country_code else None
        
        if ref_country_code:
            dist_data = model.calculate_cultural_distances(ref_country_code, weights.value)
            ref_name = country_code_to_name.get(ref_country_code, ref_country_code)
            fig = create_world_map(model, distance_data=dist_data, ref_country_name=ref_name)
            stats_html = "<p>Select a dimension to see statistics</p>"
        else:
            fig = create_world_map(model, selected_dimension.value)
            stats_html = create_statistics_table(model, selected_dimension.value)
    else:
        fig = create_world_map(model, selected_dimension.value)
        stats_html = create_statistics_table(model, selected_dimension.value)

    with solara.Column(style={'width': '100%', 'max-width': '1400px', 'margin': '0 auto'}):
        solara.Markdown("# 🌍 Hofstede's Cultural Dimensions")
        
        with solara.Card("Control Panel", style={'margin': '20px 0'}):
            solara.Markdown(f"**Selected: {dimension_display[selected_dimension.value]}**")
            
            with solara.ColumnsResponsive(12, large=[2, 2, 2], style={'gap': '10px'}):
                for code, label in dimension_display.items():
                    solara.Button(
                        label,
                        color='primary' if selected_dimension.value == code else 'default',
                        on_click=lambda c=code: selected_dimension.set(c),
                        style={'width': '100%'}
                    )
            
            solara.Checkbox(label="Show Cultural Distance Map", value=show_distance.value, on_value=show_distance.set)
            
            # Weight inputs (only when showing distance map)
            if show_distance.value:
                if complete_countries:
                    solara.Select(
                        label="Reference Country",
                        values=complete_countries,  # [(code, name), ...]
                        value=selected_country.value,
                        on_value=selected_country.set
                    )
                else:
                    solara.Warning("No countries with complete data available for distance calculation.")
            
                # Weight inputs
                solara.Markdown("### Dimension Weights (sum must equal 1.0)")
                total_weight = sum(weights.value.values())
                if abs(total_weight - 1.0) > 0.01:
                    solara.Error(f"Total weight = {total_weight:.3f} (must equal 1.0)")
            
                new_weights = weights.value.copy()
                for dim in dims:
                    def make_handler(d):
                        def handler(value_str):
                            try:
                                value = float(value_str)
                                if 0 <= value <= 1:
                                    new_weights[d] = value
                                    weights.set(new_weights.copy())
                            except ValueError:
                                pass
                        return handler
            
                    solara.InputFloat(
                        label=dimension_display[dim],
                        value=new_weights[dim],
                        on_value=make_handler(dim)
                    )

        with solara.Card("World Map Visualization", style={'margin': '20px 0'}):
            solara.FigureMatplotlib(fig, dependencies=[
                selected_dimension.value,
                show_distance.value,
                str(selected_country.value),  # Convert to string for dependency tracking
                str(weights.value)
            ])
            plt.close(fig)

        with solara.Columns([2, 1], style={'margin': '20px 0', 'gap': '20px'}):
            with solara.Card("About Hofstede's Cultural Dimensions"):
                solara.Markdown("""
                Geert Hofstede's cultural dimensions theory describes how cultural values influence behavior across societies.
                
                **The Six Dimensions:**
                  
                - **PDI**: Power Distance Index – acceptance of hierarchical power distribution
                - **IDV**: Individualism vs Collectivism – loose vs tight social frameworks
                - **MAS**: Masculinity vs Femininity – competition vs cooperation values
                - **UAI**: Uncertainty Avoidance – tolerance for ambiguity and risk
                - **LTOWVS**: Long-Term Orientation – future rewards vs tradition
                - **IVR**: Indulgence vs Restraint – gratification of desires
                
                ---
                *Data source: Hofstede Insights*
                """)
            
            with solara.Card("Dimension Statistics"):
                solara.HTML(unsafe_innerHTML=stats_html)


@solara.component
def Page():
    # ✅ Use correct filename based on your data
    model = solara.use_memo(lambda: HofstedeModel("hofstede_data_final.csv"), dependencies=[])
    WorldMapVisualization(model)

page = Page