import pandas as pd
import altair as alt
import matplotlib.pyplot as plt
import geopandas as gpd
from vega_datasets import data
import streamlit as st
import json
import folium
from streamlit_folium import folium_static
st.set_page_config(layout="wide")
st.markdown('# Plane Travel')

tsa_df = pd.read_excel('./data/tsa.xlsx')
tsa_df_date = tsa_df.set_index('Date')
tsa_df_year = tsa_df_date.stack().reset_index().rename(columns={'level_1':'year'})
tsa_df_stack = tsa_df_year.join(tsa_df_date)
tsa_df_stack = tsa_df_stack[['Date', 'year', 0]]
tsa_df_stack = tsa_df_stack.rename(columns={0:'data'})

line = alt.Chart(tsa_df_stack).mark_line().encode(
        x = alt.X('Date'),
        y = alt.Y('data', title = 'Passenger volumes'),
        color = 'year:N',
        strokeWidth = alt.value(1)
)

# 添加交互式选择器
selector = alt.selection_single(
    on='mouseover',
    nearest=True,
    empty='none',
    fields=['Date'],
    init={'Date': '2022-01-01'}
)

# 创建一个点图层，用于显示选择器位置的信息
points = line.mark_point().encode(
    opacity=alt.condition(selector, alt.value(1), alt.value(0))
).add_selection(selector)

# 创建一个文本图层，用于显示选择器位置对应的y值
text = line.mark_text(align='left', dx=5, dy=-5).encode(
    text=alt.condition(selector, 'data:Q', alt.value(' '))
)

# 添加一条垂直线
vline = alt.Chart(tsa_df_stack).mark_rule(color='red').encode(
    x='Date',
    size=alt.value(0.1)
).transform_filter(
    selector
)

# 将图表和交互式元素组合在一起
chart = alt.layer(line, points, text, vline)

# 显示图表
chart = chart.properties(
    width=800, 
    title='The amount of passengers passing through TSA checkpoints across recent years'
)

st.altair_chart(chart, use_container_width=False)

st.markdown("""
    A look into the most popular destinations by the amount of passengers on flights to specific destinations post-COVID
""")

df_2022 = pd.read_excel('./data/US-Outbound-to-World-Regions_2022.xlsx')
df_2022.columns = df_2022.iloc[2]
df_2022_select = df_2022.iloc[3:10].set_index('Regions')
df_2022_select.columns.name = None
df_2022_select = df_2022_select.iloc[:, : 12]
df_2022_final = df_2022_select.reset_index()
df_2022_stack = df_2022_final.iloc[:, 1:12].stack().reset_index().set_index('level_0').join(df_2022_final).iloc[:, 0:3].rename(columns = {'level_1':'Month', 0:'Data'})
bar_2022 = alt.Chart(df_2022_stack).mark_bar().encode(
        x = alt.X('Month'),
        y = alt.Y('Data'),
        color = 'Regions:N',
)
bar_2022 = bar_2022.properties(title='2022 US citizens travel to international regions')

df_2021 = pd.read_excel('./data/US-Outbound-to-World-Regions_2021.xlsx')
df_2021.columns = df_2021.iloc[2]
df_2021_select = df_2021.iloc[3:10].set_index('Regions')
df_2021_select.columns.name = None
df_2021_select = df_2021_select.iloc[:, : 12]
df_2021_final = df_2021_select.reset_index()
df_2021_stack = df_2021_final.iloc[:, 1:12].stack().reset_index().set_index('level_0').join(df_2021_final).iloc[:, 0:3].rename(columns = {'level_1':'Month', 0:'Data'})
# 创建一个选择器，用于捕获用户在2022年数据的月份柱状图上的鼠标悬停操作
month_selector_2022 = alt.selection_single(fields=['Month'], empty='none', on='mouseover')

# 创建一个选择器，用于捕获用户在2021年数据的月份柱状图上的鼠标悬停操作
month_selector_2021 = alt.selection_single(fields=['Month'], empty='none', on='mouseover')

# Create the bar chart for 2022 data
bar_2022 = alt.Chart(df_2022_stack).mark_bar().encode(
        x = alt.X('Month'),
        y = alt.Y('Data'),
        color = 'Regions:N',
).properties(title='2022 US citizens travel to international regions', width=300).add_selection(
    month_selector_2022
)

pie_2022 = alt.Chart(df_2022_stack).mark_arc(innerRadius=50, outerRadius=100).encode(
    theta='sum(Data)',
    color='Regions:N',
    tooltip='Regions'
).transform_filter(
    month_selector_2022
).properties(width=300)

# Create the bar chart for 2021 data
bar_2021 = alt.Chart(df_2021_stack).mark_bar().encode(
        x = alt.X('Month'),
        y = alt.Y('Data'),
        color = 'Regions:N',
).properties(title='2021 US citizens travel to international regions', width=300).add_selection(
    month_selector_2021
)

pie_2021 = alt.Chart(df_2021_stack).mark_arc(innerRadius=50, outerRadius=100).encode(
    theta='sum(Data)',
    color='Regions:N',
    tooltip='Regions'
).transform_filter(
    month_selector_2021
).properties(width=300)

# Combine the charts
combined_charts = ((bar_2022 & pie_2022) | (bar_2021 & pie_2021)).resolve_scale(y='shared')

# Display the combined charts
combined_charts


alt.data_transformers.disable_max_rows()

# Load data
destination_df = pd.read_csv('./data/air-passengers-carried.csv')
destination_df['Entity'] = destination_df['Entity'].astype(str)
destination_df['Code'] = destination_df['Code'].astype(str)

# Load world data from GeoPandas
world = gpd.read_file(gpd.datasets.get_path('naturalearth_lowres'))

# Add a slider to select the year
selected_year = st.slider("Select Year", min_value=2018, max_value=2020, value=2018, step=1)

# Filter the data based on the selected year
destination_df_year = destination_df[destination_df['Year'] == selected_year]

# Merge the data
merged_data = world.set_index('iso_a3').join(destination_df_year.set_index('Code')).reset_index()

# Convert the merged data to a GeoJSON object
merged_data_geojson = json.loads(merged_data.to_json())

# Create a folium map
m = folium.Map(location=[0, 0], zoom_start=2)

# Define a function to determine the color based on the number of passengers carried
def get_color(feature):
    passengers = feature['properties']['Air transport, passengers carried']
    if passengers is None:
        return '#FFFFFF'
    elif passengers < 1e4:
        return '#fee0b6'
    elif passengers < 1e5:
        return '#f1a340'
    elif passengers < 1e6:
        return '#d73027'
    else:
        return '#a50026'

# Add the choropleth layer to the folium map
folium.GeoJson(
    merged_data_geojson,
    style_function=lambda feature: {
        'fillColor': get_color(feature),
        'color': 'black',
        'weight': 1,
        'fillOpacity': 0.7,
    },
    tooltip=folium.features.GeoJsonTooltip(fields=['Entity', 'Air transport, passengers carried'],
                                           aliases=['Country', 'Passengers Carried'])
).add_to(m)
st.markdown("""
    A look into the most popular destinations by the amount of passengers on flights to specific destinations pre and during COVID
""")
# Display the folium map in Streamlit
folium_static(m)