import streamlit as st
import pandas as pd
import numpy as np
import folium
from folium.plugins import MarkerCluster
from streamlit_folium import st_folium
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import calendar

st.title("Khai phá dữ liệu giao thông tại thành phố New York")
st.markdown("Halo. Minh sumita")

DATA_URL = 'src/nypd-motor-vehicle-collisions.csv'


@st.cache_data
def load_data(nrows):
    data = pd.read_csv(DATA_URL, nrows=nrows, parse_dates=[['ACCIDENT DATE', 'ACCIDENT TIME']])
    data.dropna(subset=['LATITUDE', 'LONGITUDE'], inplace=True)
    lowercase = lambda x: str(x).lower()
    data.rename(lowercase, axis='columns', inplace=True)
    data.rename(columns={'accident date_accident time': 'date/time'}, inplace=True)
    return data


data = load_data(10000)
original_data = data.copy()

if st.checkbox('Show Raw data'):
    st.subheader('Raw Data')
    st.write(data)

st.header('Đâu là nơi có nhiều người bị thương nhất ở thành phố New York?')
injured_people = st.slider('Số người bị thương trong các vụ va chạm xe cộ', 0, 20)
st.map(data.query('`number of persons injured` >= @injured_people')[['latitude', 'longitude']].dropna(how='any'))

st.header('Có bao nhiêu va chạm xảy ra trong khoảng thời gian?')

min_date = data['date/time'].dt.date.min()
max_date = data['date/time'].dt.date.max()

start_date = st.date_input(
    "Chọn ngày bắt đầu",
    value=min_date,
    min_value=min_date,
    max_value=max_date
)

end_date = st.date_input(
    "Chọn ngày kết thúc",
    value=max_date,
    min_value=min_date,
    max_value=max_date
)

if start_date > end_date:
    st.error("Ngày bắt đầu phải nhỏ hơn hoặc bằng ngày kết thúc!")
else:
    st.markdown(f'Xem dữ liệu từ ngày **{start_date}** đến ngày **{end_date}**')

    filtered_data = data[(data['date/time'].dt.date >= start_date) & (data['date/time'].dt.date <= end_date)]

    st.markdown(f"Tổng số va chạm trong khoảng thời gian: **{len(filtered_data)}**")

    if not filtered_data.empty:
        midpoint = (np.average(filtered_data['latitude']), np.average(filtered_data['longitude']))

        m = folium.Map(location=midpoint, tiles='OpenStreetMap', zoom_start=11, control_scale=True)

        mc = MarkerCluster()
        for idx, row in filtered_data.iterrows():
            mc.add_child(folium.Marker(location=[row['latitude'], row['longitude']]))
        m.add_child(mc)

        st_folium(m, width=700, height=500)
    else:
        st.write("Không có dữ liệu tai nạn trong khoảng thời gian bạn chọn.")

# --- PHẦN BIỂU ĐỒ THỐNG KÊ ---

st.header("Biểu đồ thống kê tai nạn giao thông")

data['date'] = data['date/time'].dt.date
if 'borough' in data.columns:
    data['borough'] = data['borough'].astype(str).str.upper()
else:
    data['borough'] = 'UNKNOWN'

boroughs = ['QUEENS', 'BROOKLYN', 'MANHATTAN', 'BRONX', 'STATEN ISLAND']

# Thêm cột tháng
data['month'] = data['date/time'].dt.month
by_month = data.groupby('month')[['number of persons injured', 'number of persons killed']].sum().reset_index()
by_month['month_name'] = by_month['month'].apply(lambda x: calendar.month_name[x])

chart_option = st.selectbox("Chọn biểu đồ muốn xem:", [
    "Số người bị thương và tử vong theo ngày và quận",
    "Số vụ tai nạn trong 365 ngày gần nhất theo quận",
    "Tổng số người bị thương và tử vong theo năm",
    "Số người bị thương và tử vong theo loại nạn nhân (Người đi bộ, Xe đạp, Người lái xe)",
    "Số người bị thương và tử vong theo tháng"
])

if chart_option == "Số người bị thương và tử vong theo ngày và borough":
    grouped = data.groupby(['date', 'borough'])[['number of persons injured', 'number of persons killed']].sum().reset_index()
    fig = make_subplots(rows=2, cols=1, subplot_titles=('Số người bị thương mỗi ngày', 'Số người tử vong mỗi ngày'))
    for bor in boroughs:
        df_bor = grouped[grouped['borough'] == bor]
        fig.add_trace(go.Scatter(x=df_bor['date'], y=df_bor['number of persons injured'], name=bor), row=1, col=1)
        fig.add_trace(go.Scatter(x=df_bor['date'], y=df_bor['number of persons killed'], name=bor), row=2, col=1)
    fig.update_layout(height=700, template='plotly_dark')
    st.plotly_chart(fig, use_container_width=True)

elif chart_option == "Số vụ tai nạn trong 365 ngày gần nhất theo borough":
    recent_data = data[data['date'] >= (data['date'].max() - pd.Timedelta(days=365))]
    feat_cols = ['number of persons injured', 'number of persons killed']
    fig = make_subplots(rows=2, cols=1, subplot_titles=('Số người bị thương mỗi ngày (365 ngày gần nhất)', 'Số người tử vong mỗi ngày (365 ngày gần nhất)'))
    for bor in boroughs:
        df_bor = recent_data[recent_data['borough'] == bor]
        df_grouped = df_bor.groupby('date')[feat_cols].sum()
        fig.add_trace(go.Scatter(x=df_grouped.index, y=df_grouped['number of persons injured'], name=bor), row=1, col=1)
        fig.add_trace(go.Scatter(x=df_grouped.index, y=df_grouped['number of persons killed'], name=bor), row=2, col=1)
    fig.update_layout(height=700, template='plotly_dark')
    st.plotly_chart(fig, use_container_width=True)

elif chart_option == "Tổng số người bị thương và tử vong theo năm":
    data['year'] = data['date/time'].dt.year
    by_year = data.groupby('year')[['number of persons injured', 'number of persons killed']].sum().reset_index()
    fig = make_subplots(rows=1, cols=2, subplot_titles=('Tổng số người bị thương theo năm', 'Tổng số người tử vong theo năm'))
    fig.add_trace(go.Bar(x=by_year['year'], y=by_year['number of persons injured'], name='Bị thương'), row=1, col=1)
    fig.add_trace(go.Bar(x=by_year['year'], y=by_year['number of persons killed'], name='Tử vong'), row=1, col=2)
    fig.update_layout(height=400, template='plotly_dark')
    st.plotly_chart(fig, use_container_width=True)

elif chart_option == "Số người bị thương và tử vong theo loại nạn nhân (Người đi bộ, Xe đạp, Người lái xe)":
    inj_cols = ['number of pedestrians injured', 'number of cyclist injured', 'number of motorist injured']
    kill_cols = ['number of pedestrians killed', 'number of cyclist killed', 'number of motorist killed']
    labels = ['Người đi bộ', 'Xe đạp', 'Người lái xe']

    fig = make_subplots(rows=3, cols=1, subplot_titles=labels)
    for i, (inj_col, kill_col) in enumerate(zip(inj_cols, kill_cols)):
        daily_inj = data.groupby('date')[inj_col].sum()
        daily_kill = data.groupby('date')[kill_col].sum()
        fig.add_trace(go.Scatter(x=daily_inj.index, y=daily_inj.values, name='Bị thương'), row=i+1, col=1)
        fig.add_trace(go.Scatter(x=daily_kill.index, y=daily_kill.values, name='Tử vong'), row=i+1, col=1)
    fig.update_layout(height=900, template='plotly_dark')
    st.plotly_chart(fig, use_container_width=True)

elif chart_option == "Số người bị thương và tử vong theo tháng":
    fig = make_subplots(rows=1, cols=2, subplot_titles=('Số người bị thương theo tháng', 'Số người tử vong theo tháng'))
    
    fig.add_trace(go.Bar(x=by_month['month_name'], y=by_month['number of persons injured'], name='Bị thương'), row=1, col=1)
    fig.add_trace(go.Bar(x=by_month['month_name'], y=by_month['number of persons killed'], name='Tử vong'), row=1, col=2)
    
    fig.update_layout(template='plotly_dark', width=900, height=450)
    st.plotly_chart(fig, use_container_width=True)