import streamlit as st
import pandas as pd
import altair as alt

largest_hospitals = [
    {
        'name': 'Florida Hospital Orlando',
        'city': 'Orlando',
        'state': 'FL',
        'zip_code': '32803',
        'bed_count': 2411,
        'lat': 28.555149,
        'lng': -81.362244
    },
    {
        'name': 'Cleveland Clinic',
        'city': 'Cleveland',
        'state': 'OH',
        'zip_code': '44195',
        'bed_count': 1730,
        'lat': 41.501642,
        'lng': -81.621223
    },
    {
        'name': 'Mayo Clinic',
        'city': 'Rochester',
        'state': 'MN',
        'zip_code': '55905',
        'bed_count': 1372,
        'lat': 44.019126,
        'lng': -92.463362
    },
    {
        'name': 'NewYork-Presbyterian Hospital-Columbia and Cornell',
        'city': 'New York',
        'state': 'NY',
        'zip_code': '10032',
        'bed_count': 2332,
        'lat': 40.841708,
        'lng': -73.942635
    },
    {
        'name': 'UCHealth University of Colorado Hospital',
        'city': 'Aurora',
        'state': 'CO',
        'zip_code': '80045',
        'bed_count': 672,
        'lat': 39.743943,
        'lng': -104.834322
    },
    {
        'name': 'Houston Methodist Hospital',
        'city': 'Houston',
        'state': 'TX',
        'zip_code': '77030',
        'bed_count': 1063,
        'lat': 29.710773,
        'lng': -95.399676
    },
    {
        'name': 'Johns Hopkins Hospital',
        'city': 'Baltimore',
        'state': 'MD',
        'zip_code': '21287',
        'bed_count': 1293,
        'lat': 39.296154,
        'lng': -76.591972
    },
    {
        'name': 'Massachusetts General Hospital',
        'city': 'Boston',
        'state': 'MA',
        'zip_code': '02114',
        'bed_count': 1032,
        'lat': 42.362251,
        'lng': -71.069405
    },
    {
        'name': 'University of Michigan Hospitals-Michigan Medicine',
        'city': 'Ann Arbor',
        'state': 'MI',
        'zip_code': '48109',
        'bed_count': 1145,
        'lat': 42.285932,
        'lng': -83.730833
    },
    {
        'name': 'Mount Sinai Hospital',
        'city': 'New York',
        'state': 'NY',
        'zip_code': '10029',
        'bed_count': 1168,
        'lat': 40.788127,
        'lng': -73.952826
    }
]

largest_hospitals_df = pd.DataFrame(largest_hospitals)

def stacked_bar_chart_with_text_overlay():
    chart = alt.Chart(largest_hospitals_df).mark_bar().encode(
        y=alt.Y('state:N', sort='-x'),
        x=alt.X('bed_count:Q', stack='normalize'),
        color=alt.Color('name:N'),
        tooltip=['name', 'city', 'state', 'bed_count']
    ).properties(
        width=700,
        height=500,
        title='Largest Hospitals by State (Stacked Bar Chart with Text Overlay)'
    ).configure_axisX(
        labelAngle=-45
    )
    text = chart.mark_text(align='left', baseline='middle', dx=3).encode(
        text=alt.Text('bed_count:Q', format='.1f')
    )
    st.altair_chart(chart + text)

def bump_chart():
    chart = alt.Chart(largest_hospitals_df).transform_joinaggregate(
        rank='rank(bed_count)',
        groupby=['state']
    ).transform_filter(
        alt.datum.rank <= 3
    ).transform_window(
        y='row_number()',
        sort=[alt.SortField('bed_count', order='descending')]
    ).mark_line().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('y:O', axis=None),
        color=alt.Color('name:N'),
        tooltip=['name', 'city', 'state', 'bed_count']
    ).properties(
        width=700,
        height=500,
        title='Largest Hospitals by State (Bump Chart)'
    )
    st.altair_chart(chart)

def radial_chart():
    chart = alt.Chart(largest_hospitals_df).mark_bar().encode(
        x=alt.X('count()', title='Count'),
        y=alt.Y('state:N', sort='-x'),
        color=alt.Color('bed_count:Q', legend=None),
        column=alt.Column('bed_count:Q', bin=alt.Bin(maxbins=10)),
        tooltip=['name', 'city', 'state', 'bed_count']
    ).properties(
        width=700,
        height=500,
        title='Largest Hospitals by State (Radial Chart)'
    )
    st.altair_chart(chart)

def trellis_area_sort_chart():
    chart = alt.Chart(largest_hospitals_df).mark_area(opacity=0.8).encode(
        x=alt.X('yearmonth(date):T', title='Date'),
        y=alt.Y('sum(revenue):Q', stack='center', axis=None),
        color=alt.Color('product:N', scale=alt.Scale(scheme='category10')),
        row=alt.Row('market:N', header=alt.Header(title='Market')),
        column=alt.Column('product:N', header=alt.Header(title='Product')),
        tooltip=[alt.Tooltip('product:N'), alt.Tooltip('revenue:Q', format='$,.0f')]
    ).properties(
        width=300,
        height=200,
        title='Trellis Area Sort Chart'
    ).configure_facet(
        spacing=0
    ).configure_view(
        stroke=None
    )
    st.altair_chart(chart)

def wind_vector_map():
    source = pd.DataFrame({
        'lat': largest_hospitals_df['lat'],
        'lon': largest_hospitals_df['lng'],
        'u': [10, 20, 30, 40, 50, -10, -20, -30, -40, -50],
        'v': [-10, -20, -30, -40, -50, 10, 20, 30, 40, 50],
        'names': largest_hospitals_df['name']
    })
    max_speed = 60

    # Create a layer of the world map
    background = alt.Chart(
        data=topo_feature('world-110m')
    ).mark_geoshape(
        fill='white',
        stroke='lightgray'
    ).properties(
        width=700,
        height=400
    ).project('naturalEarth1')

    # Add the wind vectors as arrows
    vectors = background.mark_arrow(
        length=300,
        stroke='black',
        strokeWidth=0.5
    ).encode(
        longitude='lon:Q',
        latitude='lat:Q',
        angle=alt.Angle('atan2(v, u):Q'),
        size=alt.Size(alt.Color('length:Q', legend=None), scale=alt.Scale(range=[0, 0.08]), title='Wind speed'),
        opacity=alt.Opacity(alt.Color('length:Q', legend=None), scale=alt.Scale(range=[0, 1]), title='Wind speed'),
        tooltip=['names:N', alt.Tooltip('length:Q', format='.1f')]
    ).transform_calculate(
        # Cartographic rotation for arrows
        angle=calc_wind_angle('u', 'v'),
        # Vector length
        length=calc_wind_speed('u', 'v'),
        # Limit vector length
        length=alt.datum.length > max_speed ? max_speed : alt.datum.length
    )

    st.altair_chart(background + vectors)

def table_bubble_plot():
    chart = alt.Chart(largest_hospitals_df).mark_circle().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('state:N', sort='-x'),
        size=alt.Size('bed_count:Q', title='Bed Count'),
        color=alt.Color('bed_count:Q', legend=None),
        tooltip=['name', 'city', 'state', 'bed_count']
    ).properties(
        width=700,
        height=500,
        title='Largest Hospitals by State (Table Bubble Plot)'
    )
    st.altair_chart(chart)

def locations_of_us_airports():
    airports = data.airports.url

    states = alt.topo_feature(data.us_10m.url, 'states')
    lookup = {'New York City': 'New York', 'Chicago': 'Illinois', 'Los Angeles': 'California', 'San Francisco': 'California', 'Houston': 'Texas'}

    chart = alt.Chart(states).mark_geoshape(
        fill='lightgray',
        stroke='white'
    ).encode(
        color=alt.Color('count()', scale=alt.Scale(scheme='yelloworangered')),
        tooltip=[alt.Tooltip('state:N'), alt.Tooltip('count():Q')]
    ).transform_lookup(
        lookup='state',
        from_=alt.LookupData(airports, 'state', ['latitude', 'longitude'])
    ).transform_fold(
        ['latitude', 'longitude'],
        as_=['key', 'value']
    ).transform_filter(
        (alt.datum.value[0] != 'NaN') & (alt.datum.value[1] != 'NaN')
    ).mark_circle(
        size=10
    ).encode(
        longitude='value:Q',
        latitude='key:Q',
        color=alt.Color('count()', scale=alt.Scale(scheme='yelloworangered')),
        tooltip=[alt.Tooltip('state:N'), alt.Tooltip('count():Q')]
    ).properties(
        width=700,
        height=500,
        title='Locations of US Airports'
    )
    st.altair_chart(chart)

def connections_among_us_airports_interactive():
    airports = data.airports.url
    routes = data.routes.url

    states = alt.topo_feature(data.us_10m.url, 'states')

    source = pd.read_json(airports)
    lookup = {'New York City': 'New York', 'Chicago': 'Illinois', 'Los Angeles': 'California', 'San Francisco': 'California', 'Houston': 'Texas'}
    source['state'] = source['state'].apply(lambda x: lookup[x] if x in lookup.keys() else x)

    base = alt.Chart(states).mark_geoshape(
        fill='lightgray',
        stroke='white'
    ).properties(
        width=700,
        height=400
    ).project('albersUsa')

    airports = base.mark_circle(size=10).encode(
        longitude='longitude:Q',
        latitude='latitude:Q',
        tooltip=['name:N', 'city:N', 'state:N', 'country:N']
    ).transform_lookup(
        lookup='iata',
        from_=alt.LookupData(airports, 'iata', ['name', 'city', 'state', 'country', 'latitude', 'longitude'])
    ).properties(title='US Airports')

    routes = base.mark_geoshape(
        stroke='black',
        strokeWidth=0.1
    ).encode(
        longitude='start_lon:Q',
        latitude='start_lat:Q',
        longitude2='end_lon:Q',
        latitude2='end_lat:Q'
    ).transform_lookup(
        lookup='start',
        from_=alt.LookupData(source, 'iata', ['state', 'latitude', 'longitude']),
        as_=['start_state', 'start_lat', 'start_lon']
    ).transform_lookup(
        lookup='end',
        from_=alt.LookupData(source, 'iata', ['state', 'latitude', 'longitude']),
        as_=['end_state', 'end_lat', 'end_lon']
    ).transform_filter(
        (alt.datum.start_lat != None) & (alt.datum.start_lon != None) & (alt.datum.end_lat != None) & (alt.datum.end_lon != None)
    ).transform_aggregate(
        count='count()',
        groupby=['start', 'start_state', 'end', 'end_state']
    ).transform_filter(
        (alt.datum['count'] > 10)
    ).transform_calculate(
        start_lon=-alt.datum.start_lon,
        end_lon=-alt.datum.end_lon
    )

    chart = (base + routes + airports).configure_view(
        width=800,
        height=500,
        stroke=None
    )
    st.altair_chart(chart)

def one_dot_per_zipcode():
    chart = alt.Chart(largest_hospitals_df).mark_circle().encode(
        longitude='lng:Q',
        latitude='lat:Q',
        size=alt.Size('bed_count:Q', title='Bed Count'),
        color=alt.Color('bed_count:Q', legend=None),
        tooltip=['name', 'city', 'state', 'zip_code', 'bed_count']
    ).properties(
        width=700,
        height=500,
        title='One Dot Per Zipcode'
    )
    st.altair_chart(chart)

def isotype_visualization_with_emoji():
    chart = alt.Chart(largest_hospitals_df).mark_point().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('state:N', sort='-x'),
        color=alt.Color('state:N'),
        shape=alt.Shape('state:N'),
        tooltip=['name', 'city', 'state', 'zip_code', 'bed_count']
    )

    shape_lookup = {'CO': '🏥', 'FL': '🏥', 'MA': '🏥', 'MD': '🏥', 'MI': '🏥', 'MN': '🏥', 'NY': '🏥', 'OH': '🏥', 'TX': '🏥'}
    chart = chart.transform_calculate(
        shape=f'"{shape_lookup}"[datum.state]'
    )

    chart = chart.mark_text(
        align='center',
        baseline='middle',
        size=30,
        font='Segoe UI Emoji',
        dx=0,
        dy=0,
    ).encode(
        text='shape:N'
    ).properties(
        width=700,
        height=500,
        title='Isotype Visualization with Emoji'
    )
    st.altair_chart(chart)

def binned_heatmap():
    chart = alt.Chart(largest_hospitals_df).mark_rect().encode(
        x=alt.X('bed_count:Q', bin=True),
        y=alt.Y('state:N'),
        color=alt.Color('count()', scale=alt.Scale(scheme='yelloworangered')),
        tooltip=[alt.Tooltip('state:N'), alt.Tooltip('count():Q')]
    ).properties(
        width=700,
        height=500,
        title='Binned Heatmap'
    )
    st.altair_chart(chart)

def facetted_scatterplot_with_marginal_histograms():
    brush = alt.selection(type='interval', encodings=['x'])

    base = alt.Chart(largest_hospitals_df).transform_filter(
        brush
    ).properties(
        width=500,
        height=500
    )

    points = base.mark_point().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('state:N', sort='-x', title=None),
        color=alt.Color('state:N'),
        tooltip=['name', 'city', 'state', 'zip_code', 'bed_count']
    )

    top_hist = base.mark_bar().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('count()', title='Number of Hospitals'),
        color=alt.condition(brush, alt.ColorValue('gray'), alt.ColorValue('lightgray')),
    ).properties(
        title='Bed Count Distribution',
        width=500,
        height=100
    )

    right_hist = base.mark_bar().encode(
        y=alt.X('state:N', title='State'),
        x=alt.X('count()', title='Number of Hospitals'),
        color=alt.condition(brush, alt.ColorValue('gray'), alt.ColorValue('lightgray')),
    ).properties(
        title='Hospital Count by State',
        width=100,
        height=500
    )

    chart = ((points | top_hist) & right_hist).add_selection(
        brush
    ).configure_view(
        stroke=None
    ).properties(
        title='Facetted Scatterplot with Marginal Histograms',
        width=700,
        height=500
    )
    st.altair_chart(chart)
    
def ridgeline_plot():
    base = alt.Chart(largest_hospitals_df).transform_density(
        'bed_count',
        as_=['bed_count', 'density'],
        extent=[0, 3000],
        bandwidth=50,
        groupby=['state']
    )

    chart = base.mark_area().encode(
        x=alt.X('bed_count:Q', title='Bed Count'),
        y=alt.Y('state:N', sort='-x', title=None),
        color=alt.Color('state:N'),
        opacity=alt.Opacity('density:Q', legend=None, scale=alt.Scale(range=[0.3, 1]))
    ).properties(
        title='Ridgeline Plot',
        width=700,
        height=500
    )

    st.altair_chart(chart)


def create_sidebar():
    chart_functions = {
        'Stacked Bar Chart with Text Overlay': stacked_bar_chart,
        'Bump Chart': bump_chart,
        'Radial Chart': radial_chart,
        'Trellis Area Sort Chart': trellis_area_sort_chart,
        'Wind Vector Map': wind_vector_map,
        'Table Bubble Plot': table_bubble_plot,
        'Locations of US Airports': locations_of_us_airports,
        'Connections Among U.S. Airports Interactive': connections_among_us_airports_interactive,
        'One Dot Per Zipcode': one_dot_per_zipcode,
        'Isotype Visualization with Emoji': isotype_visualization_with_emoji,
        'Binned Heatmap': binned_heatmap,
        'Facetted Scatterplot with Marginal Histograms': facetted_scatterplot_with_marginal_histograms,
        'Ridgeline Plot': ridgeline_plot
    }
    
    st.sidebar.title('Charts')
    
    for chart_name, chart_function in chart_functions.items():
        chart_button = st.sidebar.button(f'{chart_name} {emoji(chart_name)}')
        if chart_button:
            chart_function()

def emoji(chart_name):
    emojis = {
        'Stacked Bar Chart with Text Overlay': '📊',
        'Bump Chart': '📈',
        'Radial Chart': '🎡',
        'Trellis Area Sort Chart': '📉',
        'Wind Vector Map': '🌬️',
        'Table Bubble Plot': '💬',
        'Locations of US Airports': '✈️',
        'Connections Among U.S. Airports Interactive': '🛫',
        'One Dot Per Zipcode': '📍',
        'Isotype Visualization with Emoji': '😀',
        'Binned Heatmap': '🗺️',
        'Facetted Scatterplot with Marginal Histograms': '🔳',
        'Ridgeline Plot': '🏔️'
    }
    return emojis.get(chart_name, '')

create_sidebar()