import streamlit as st
import pandas as pd
import plotly.express as px
import altair as alt
import folium
from folium.plugins import HeatMap, MarkerCluster
from streamlit_folium import st_folium
from streamlit_plotly_events import plotly_events  # Ensure this is installed

@st.cache_data
def load_and_preprocess_data(file_path):
    # Read the data
    df = pd.read_csv(file_path)
    
    # Basic preprocessing
    df = df.drop(['X', 'Y'], axis=1)
    df.dropna(subset=['Incidentid', 'DateTime', 'Year', 'Latitude', 'Longitude'], inplace=True)

    # Convert Year to int 
    df['Year'] = df['Year'].astype(int)

    # Fill missing values
    numeric = ['Age_Drv1', 'Age_Drv2']
    for col in numeric:
        df[col].fillna(df[col].median(), inplace=True)
        
    categorical = [
        'Gender_Drv1', 'Violation1_Drv1', 'AlcoholUse_Drv1', 'DrugUse_Drv1',
        'Gender_Drv2', 'Violation1_Drv2', 'AlcoholUse_Drv2', 'DrugUse_Drv2',
        'Unittype_Two', 'Traveldirection_Two', 'Unitaction_Two', 'CrossStreet'
    ]
    for col in categorical:
        df[col].fillna('Unknown', inplace=True)
    
    # Remove invalid ages
    df = df[
        (df['Age_Drv1'] <= 90) & 
        (df['Age_Drv2'] <= 90) & 
        (df['Age_Drv1'] >= 16) & 
        (df['Age_Drv2'] >= 16)
    ]
    
    # Create age groups
    bins = [15, 25, 35, 45, 55, 65, 90]
    labels = ['16-25', '26-35', '36-45', '46-55', '56-65', '65+']
    
    df['Age_Group_Drv1'] = pd.cut(df['Age_Drv1'], bins=bins, labels=labels)
    df['Age_Group_Drv2'] = pd.cut(df['Age_Drv2'], bins=bins, labels=labels)
    
    if 'Weather' not in df.columns:
        df['Weather'] = 'Unknown'

    return df

def create_violation_distribution_chart(df, selected_age='All Ages'):
    # Filter by age group if needed
    if selected_age != 'All Ages':
        df = df[(df['Age_Group_Drv1'] == selected_age) | (df['Age_Group_Drv2'] == selected_age)]

    # Combine violations
    violations = pd.concat([
        df['Violation1_Drv1'].value_counts(),
        df['Violation1_Drv2'].value_counts()
    ]).groupby(level=0).sum().reset_index()
    violations.columns = ['Violation', 'Count']

    fig = px.bar(
        violations,
        x='Violation',
        y='Count',
        title=f'Number of Incidents per Violation Type - {selected_age}',
        labels={'Count': 'Number of Incidents', 'Violation': 'Violation Type'},
        height=600
    )
    fig.update_layout(clickmode='event+select', xaxis_tickangle=-45)
    
    return fig, violations

def create_severity_distribution_for_violation(df, violation):
    # Filter for the selected violation
    filtered_df = df[(df['Violation1_Drv1'] == violation) | (df['Violation1_Drv2'] == violation)]
    severity_count = filtered_df['Injuryseverity'].value_counts().reset_index()
    severity_count.columns = ['Severity', 'Count']

    fig = px.bar(
        severity_count,
        x='Severity',
        y='Count',
        title=f'Severity Distribution for {violation}',
        labels={'Count': 'Number of Incidents', 'Severity': 'Injury Severity'},
        height=400
    )
    fig.update_layout(xaxis_tickangle=-45)
    
    return fig

@st.cache_data
def create_map(df, selected_year):
    filtered_df = df[df['Year'] == selected_year]
        
    m = folium.Map(
        location=[33.4255, -111.9400],
        zoom_start=12,
        control_scale=True,
        tiles='CartoDB positron'
    )
    
    marker_cluster = MarkerCluster().add_to(m)
        
    for _, row in filtered_df.iterrows():
        folium.Marker(
            location=[row['Latitude'], row['Longitude']],
            popup=f"Accident at {row['Longitude']}, {row['Latitude']}<br>Date: {row['DateTime']}<br>Severity: {row['Injuryseverity']}",
            icon=folium.Icon(color='red')
        ).add_to(marker_cluster)
    
    heat_data = filtered_df[['Latitude', 'Longitude']].values.tolist()
    HeatMap(heat_data, radius=15, max_zoom=13, min_opacity=0.3).add_to(m)
    
    return m

def create_injuries_fatalities_chart(crash_data, unit_type):
    crash_data = crash_data[['DateTime', 'Totalinjuries', 'Totalfatalities', 'Unittype_One', 'Unittype_Two']].dropna()
    crash_data['DateTime'] = pd.to_datetime(crash_data['DateTime'], errors='coerce')
    crash_data['Month'] = crash_data['DateTime'].dt.month_name()

    month_order = ['January', 'February', 'March', 'April', 'May', 'June', 
                   'July', 'August', 'September', 'October', 'November', 'December']
    crash_data['Month'] = pd.Categorical(crash_data['Month'], categories=month_order, ordered=True)

    if unit_type == 'Total':
        filtered_data = crash_data
    else:
        unit_one, unit_two = unit_type.split(' vs ')
        filtered_data = crash_data[((crash_data['Unittype_One'] == unit_one) & (crash_data['Unittype_Two'] == unit_two)) |
                                   ((crash_data['Unittype_One'] == unit_two) & (crash_data['Unittype_Two'] == unit_one))]

    monthly_sum = filtered_data.groupby('Month').agg({'Totalinjuries': 'sum', 'Totalfatalities': 'sum'}).reset_index()

    injuries = monthly_sum[['Month', 'Totalinjuries']].rename(columns={'Totalinjuries': 'Value'})
    injuries['Measure'] = 'Total Injuries'
    fatalities = monthly_sum[['Month', 'Totalfatalities']].rename(columns={'Totalfatalities': 'Value'})
    fatalities['Measure'] = 'Total Fatalities'

    combined_data = pd.concat([injuries, fatalities])

    line_chart = alt.Chart(combined_data).mark_line(point=True).encode(
        x=alt.X('Month:N', sort=month_order, title='Month'),
        y=alt.Y('Value:Q', title='Total Injuries & Fatalities'),
        color=alt.Color('Measure:N', title='', scale=alt.Scale(domain=['Total Injuries', 'Total Fatalities'], range=['blue', 'red'])),
        tooltip=['Month', 'Measure:N', 'Value:Q']
    ).properties(
        title=f'Total Injuries and Fatalities by Month for Unit Type Pair: {unit_type}',
        width=600,
        height=400
    )

    return line_chart

def create_crash_trend_chart(df, weather=None):
    if weather and weather != 'All Conditions':
        df = df[df['Weather'] == weather]
    
    trend_data = df.groupby('Year')['Incidentid'].nunique().reset_index()
    trend_data.columns = ['Year', 'Crash Count']
    
    fig = px.line(
        trend_data,
        x='Year',
        y='Crash Count',
        title=f'Crash Trend Over Time ({weather})',
        labels={'Year': 'Year', 'Crash Count': 'Number of Unique Crashes'},
        markers=True,
        height=600
    )
    
    fig.update_traces(line=dict(width=2), marker=dict(size=8))
    fig.update_layout(legend_title_text='Trend')
    
    return fig

def create_category_distribution_chart(df, selected_category, selected_year):
    if selected_year != 'All Years':
        df = df[df['Year'] == int(selected_year)]

    grouped_data = df.groupby([selected_category, 'Injuryseverity']).size().reset_index(name='Count')
    total_counts = grouped_data.groupby(selected_category)['Count'].transform('sum')
    grouped_data['Percentage'] = (grouped_data['Count'] / total_counts * 100).round(2)

    fig = px.bar(
        grouped_data,
        x=selected_category,
        y='Count',
        color='Injuryseverity',
        text='Percentage',
        title=f'Distribution of Incidents by {selected_category} ({selected_year})',
        labels={'Count': 'Number of Incidents', selected_category: 'Category'},
        height=600,
    )

    fig.update_traces(texttemplate='%{text}%', textposition='inside')
    fig.update_layout(
        barmode='stack',
        xaxis_tickangle=-45,
        legend_title='Injury Severity',
        margin=dict(t=50, b=100),
    )

    return fig

def main():
    st.title('Traffic Accident Dataset')

    st.markdown("""
     **Team Members:**
    - Janhavi Tushar Zarapkar （jzarap2@illinois.edu）
    - Hangyue Zhang (hz85@illinois.edu)
    - Andrew Nam (donginn2@illinois.edu)
    - Nirmal Attarde
    - Maanas Sandeep Agrawa
    """)

    st.markdown("""
    ### Introduction to the Traffic Accident Dataset
    This dataset contains detailed information about traffic accidents in the city of **Tempe**.
    """)

    # Load data
    df = load_and_preprocess_data('1.08_Crash_Data_Report_(detail).csv')

    tab1, tab2, tab3, tab4, tab5 = st.tabs(["Crash Statistics", "Crash Map", "Crash Trend", "Crash Injuries/Fatalities","Distribution by Category"])
    
    with tab1:
        # Age group selection
        age_groups = ['All Ages', '16-25', '26-35', '36-45', '46-55', '56-65', '65+']
        selected_age = st.selectbox('Select Age Group:', age_groups)
    
        # Create and display the main violation distribution chart
        fig, violations = create_violation_distribution_chart(df, selected_age)
    
        # Display the figure using plotly_events only (no extra st.plotly_chart)
        selected_points = plotly_events(
            fig,
            click_event=True,
            hover_event=False,
            select_event=True,
            key="violation_chart"
        )
    
        # If user clicked on a bar, show severity distribution
        if selected_points:
            clicked_violation = violations.iloc[selected_points[0]['pointIndex']]['Violation']
            severity_fig = create_severity_distribution_for_violation(df, clicked_violation)
            st.plotly_chart(severity_fig, use_container_width=True)
    
        # Display total incidents info
        if selected_age == 'All Ages':
            total_incidents = len(df)
        else:
            total_incidents = len(df[
                (df['Age_Group_Drv1'] == selected_age) | 
                (df['Age_Group_Drv2'] == selected_age)
            ])
    
        st.markdown(f"### Total Incidents for {selected_age}")
        st.markdown(f"**{total_incidents:,}** incidents")
    

    with tab2:
        years = sorted(df['Year'].unique())
        selected_year = st.selectbox('Select Year:', years)
        
        st.markdown("### Crash Location Map")
        m = create_map(df, selected_year)
        st_folium(
            m,
            width=800,
            height=600,
            key=f"map_{selected_year}",
            returned_objects=["null_drawing"]
        )

    with tab3:
        weather = ['All Conditions'] + sorted(df['Weather'].unique())
        selected_weather = st.selectbox('Select Weather Condition:', weather)
        
        st.markdown("### Crash Trend Over Time")
        trend_fig = create_crash_trend_chart(df, selected_weather)
        st.plotly_chart(trend_fig, use_container_width=True)

    with tab4:
        unit_type_pairs = set()
        for _, row in df[['Unittype_One', 'Unittype_Two']].dropna().iterrows():
            if row['Unittype_One'] != 'Driverless' or row['Unittype_Two'] != 'Driverless':
                pair = ' vs '.join(sorted([row['Unittype_One'], row['Unittype_Two']]))
                unit_type_pairs.add(pair)
        unit_type_pairs = sorted(list(unit_type_pairs))
        unit_type = st.selectbox("Select Unit Type Pair", options=['Total'] + unit_type_pairs)
    
        injuries_fatalities_chart = create_injuries_fatalities_chart(df, unit_type)
        st.altair_chart(injuries_fatalities_chart, use_container_width=True)

    with tab5:
        categories = [
            'Collisionmanner',
            'Lightcondition',
            'Weather',
            'SurfaceCondition',
            'AlcoholUse_Drv1',
            'Gender_Drv1',
        ]
        selected_category = st.selectbox("Select Category:", categories)
        years = ['All Years'] + sorted(df['Year'].dropna().unique().astype(int).tolist())
        selected_year = st.selectbox("Select Year:", years)
        
        st.markdown(f"### Distribution of Incidents by {selected_category}")
        distribution_chart = create_category_distribution_chart(df, selected_category, selected_year)
        st.plotly_chart(distribution_chart, use_container_width=True)

if __name__ == "__main__":
    main()