import gradio as gr
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import numpy as np
import math
import random

# Load and preprocess data
df = pd.read_csv('Hurun_Global_Rich_List_2025_COORD.csv')
df['Wealth $B'] = pd.to_numeric(df['Wealth $B'], errors='coerce')
df = df.dropna(subset=['lat', 'lon'])
df['Log Wealth'] = np.log10(df['Wealth $B'])

# Ensure "Age" is an integer
df['Age'] = pd.to_numeric(df['Age'], errors='coerce').astype('Int64')

# Extract country from location
df['Country'] = df['Location'].str.split(',').str[-1].str.strip()

# Add jitter to coordinates to prevent overlapping
def add_jitter(coords, jitter_range=0.2):
    """Add random jitter to coordinates to prevent overlapping markers"""
    coord_counts = {}
    jittered_coords = []
    
    for coord in coords:
        rounded = (round(coord[0], 2), round(coord[1], 2))
        count = coord_counts.get(rounded, 0)
        coord_counts[rounded] = count + 1
        if count > 0:
            angle = 2 * math.pi * (count / 8)
            distance = jitter_range * min(0.5 + count/10, 1.5)
            jitter_lat = math.sin(angle) * distance
            jitter_lon = math.cos(angle) * distance
            jittered_coords.append((coord[0] + jitter_lat, coord[1] + jitter_lon))
        else:
            jittered_coords.append(coord)
    
    return jittered_coords

# Create logarithmic scale for visualization
def create_globe():
    df['Size'] = df['Log Wealth'] * 8
    df['Color Value'] = df['Log Wealth']
    coords = list(zip(df['lat'], df['lon']))
    jittered_coords = add_jitter(coords)
    df['lat_jitter'] = [c[0] for c in jittered_coords]
    df['lon_jitter'] = [c[1] for c in jittered_coords]
    
    min_log = df['Log Wealth'].min()
    max_log = df['Log Wealth'].max()
    wealth_ticks = [1, 10, 30, 100, 300, 1000]
    log_ticks = [math.log10(x) for x in wealth_ticks if x <= 10**max_log and x >= 10**min_log]
    tick_labels = [f"${x}B" for x in wealth_ticks if x <= 10**max_log and x >= 10**min_log]
    
    fig = go.Figure(go.Scattergeo(
        lon = df['lon_jitter'],
        lat = df['lat_jitter'],
        text = df.apply(lambda row: (
            f"<b>{row['Name']}</b><br>"
            f"Rank: {row['Rank']}<br>"
            f"Wealth: ${row['Wealth $B']}B<br>"
            f"Enterprise: {row['Enterprise']}<br>"
            f"Sector: {row['Sector']}<br>"
            f"Age: {row['Age']}<br>"
            f"Location: {row['Location']}"
        ), axis=1),
        marker = dict(
            size = df['Size'],
            color = df['Color Value'],
            colorscale = 'Turbo',
            opacity = 0.7,
            line_color='rgb(40,40,40)',
            line_width=0.5,
            sizemode = 'diameter',
            colorbar = dict(
                title="Wealth (Billion USD)",
                thickness=15,
                len=0.35,
                bgcolor='rgba(255,255,255,0.5)',
                tickvals = log_ticks,
                ticktext = tick_labels
            )
        ),
        hoverinfo = 'text',
        name = ''
    ))
    
    fig.update_geos(
        projection_type="orthographic",
        showcoastlines=True,
        coastlinecolor="RebeccaPurple",
        showland=True,
        landcolor="rgb(243, 243, 243)",
        showocean=True,
        oceancolor="rgb(217, 244, 252)"
    )
    
    fig.update_layout(
        height=500,
        margin={"r":0,"t":0,"l":0,"b":0},
        paper_bgcolor='rgba(0,0,0,0)',
        geo=dict(bgcolor='rgba(0,0,0,0)')
    )
    return fig

# Create distribution plots
def create_age_plot():
    fig = px.histogram(df, x='Age', nbins=20, 
                      title='Age Distribution',
                      color_discrete_sequence=['#636EFA'])
    fig.update_layout(height=300)
    return fig

def create_country_plot():
    top_countries = df['Country'].value_counts().head(10)
    fig = px.bar(top_countries, 
                title='Top 10 Countries',
                labels={'value': 'Count', 'index': 'Country'},
                color=top_countries.index,
                color_discrete_sequence=px.colors.qualitative.Pastel)
    fig.update_layout(height=300, showlegend=False)
    return fig

def create_gender_plot():
    gender_counts = df['Sex'].value_counts()
    fig = px.pie(gender_counts, 
                names=gender_counts.index, 
                values=gender_counts.values,
                title='Gender Distribution',
                hole=0.4)
    fig.update_layout(height=300)
    return fig

def create_sector_plot():
    top_sectors = df['Sector'].value_counts().head(10)
    fig = px.bar(top_sectors, 
                title='Top 10 Sectors',
                labels={'value': 'Count', 'index': 'Sector'},
                color=top_sectors.index,
                color_discrete_sequence=px.colors.qualitative.Pastel)
    fig.update_layout(height=300, showlegend=False)
    return fig

# Create display dataframe without extra columns
display_columns = ['Rank', 'Rank change', 'Wealth $B', 'Enterprise', 'Sector', 
                   'Name', 'Sex', 'Age', 'Sidekick', 'Location']
display_df = df[display_columns].sort_values('Rank').reset_index(drop=True)

# Function to update table based on search and sort
def update_table(search_term, sort_column, sort_order):
    temp_df = df[display_columns].copy()
    
    # Filter based on search term
    if search_term:
        search_term = search_term.lower()
        temp_df = temp_df[temp_df.apply(lambda row: search_term in ' '.join(map(str, row)).lower(), axis=1)]
    
    # Sort based on selected column and order
    if sort_column:
        ascending = (sort_order == "Ascending")
        temp_df = temp_df.sort_values(by=sort_column, ascending=ascending)
    
    return temp_df.reset_index(drop=True)

# Create Gradio interface
with gr.Blocks(theme='SebastianBravo/simci_css', css="#my_dataframe th { font-size: 12px !important; } #my_dataframe td { font-size: 10px !important; }") as demo:
    gr.Markdown("# Global Billionaires Dashboard 2025")
    gr.Markdown("Interactive visualisation of the world's wealthiest individuals")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🌍 Global Wealth Distribution")
            gr.Markdown("Circle size and color intensity both represent wealth on a logarithmic scale")
            gr.Markdown("Nearby points are slightly randomized to prevent overlapping")
            globe = gr.Plot(create_globe(), show_label=False)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📊 Distribution Analysis")
    
    with gr.Row():
        with gr.Column():
            age_plot = gr.Plot(create_age_plot(), show_label=False)
        with gr.Column():
            country_plot = gr.Plot(create_country_plot(), show_label=False)
    
    with gr.Row():
        with gr.Column():
            gender_plot = gr.Plot(create_gender_plot(), show_label=False)
        with gr.Column():
            sector_plot = gr.Plot(create_sector_plot(), show_label=False)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 📋 Full Billionaire Data")
            # Add search and sort controls
            with gr.Row():
                search_box = gr.Textbox(label="Search", placeholder="Type to search")
                with gr.Column():
                    sort_dropdown = gr.Dropdown(choices=display_columns, label="Sort by", value="Rank")
                    sort_order = gr.Dropdown(choices=["Ascending", "Descending"], label="Sort order", value="Ascending")
            
            # Dataframe with custom elem_id
            data_table = gr.Dataframe(
                value=display_df,
                interactive=False,
                wrap=True,
                col_count=(len(display_columns), "fixed"),
                elem_id="my_dataframe"
            )
    
    # Set up event listeners for search and sort
    search_box.change(update_table, inputs=[search_box, sort_dropdown, sort_order], outputs=data_table)
    sort_dropdown.change(update_table, inputs=[search_box, sort_dropdown, sort_order], outputs=data_table)
    sort_order.change(update_table, inputs=[search_box, sort_dropdown, sort_order], outputs=data_table)

demo.launch()