import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import numpy as np

# Page config
st.set_page_config(
    page_title="Uber NYC Pickups",
    page_icon="🚕",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Header with required link
st.title("🚕 Uber NYC Pickups")
st.markdown("[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")

# Load data (simulate as if from the repo; in practice, download or use cached)
@st.cache_data
def load_data():
    # For demo purposes, generate sample data similar to uber-raw-data-sep14.csv
    # In a real app, load from 'https://github.com/streamlit/demo-uber-nyc-pickups/raw/master/uber-raw-data-sep14.csv'
    np.random.seed(42)
    n_points = 10000
    dates = pd.date_range('2014-09-01', periods=n_points, freq='min')
    lats = np.random.normal(40.75, 0.1, n_points)
    lons = np.random.normal(-73.97, 0.1, n_points)
    df = pd.DataFrame({
        'Date/Time': dates,
        'Lat': lats,
        'Lon': lons,
        'Base': np.random.choice(['B02512', 'B02598'], n_points)
    })
    return df

df = load_data()

# Sidebar filters
st.sidebar.header("Filters")
date_range = st.sidebar.slider(
    "Select Date Range",
    min_value=df['Date/Time'].min(),
    max_value=df['Date/Time'].max(),
    value=(df['Date/Time'].min(), df['Date/Time'].max()),
    format="YYYY-MM-DD HH:MM"
)
hour_filter = st.sidebar.slider("Select Hour of Day", 0, 23, (0, 23))

base_filter = st.sidebar.multiselect(
    "Select Base",
    options=df['Base'].unique(),
    default=df['Base'].unique()
)

# Filter data
filtered_df = df[
    (df['Date/Time'] >= date_range[0]) &
    (df['Date/Time'] <= date_range[1]) &
    (df['Date/Time'].dt.hour.between(hour_filter[0], hour_filter[1])) &
    (df['Base'].isin(base_filter))
].copy()

# Metrics
col1, col2, col3 = st.columns(3)
with col1:
    st.metric("Total Pickups", len(filtered_df))
with col2:
    avg_hour = filtered_df['Date/Time'].dt.hour.mean()
    st.metric("Average Hour", f"{avg_hour:.0f}")
with col3:
    unique_bases = filtered_df['Base'].nunique()
    st.metric("Unique Bases", unique_bases)

# Visualizations
if len(filtered_df) > 0:
    # Map
    st.subheader("Pickups Map")
    fig_map = px.scatter_mapbox(
        filtered_df,
        lat="Lat",
        lon="Lon",
        color="Base",
        hover_data=["Date/Time"],
        zoom=10,
        height=500,
        mapbox_style="carto-positron"
    )
    st.plotly_chart(fig_map, use_container_width=True)

    # Time series
    st.subheader("Pickups Over Time")
    filtered_df['Date'] = filtered_df['Date/Time'].dt.date
    hourly_data = filtered_df.groupby(filtered_df['Date/Time'].dt.hour).size().reset_index(name='Count')
    fig_time = px.line(hourly_data, x='Date/Time', y='Count', title="Pickups by Hour")
    st.plotly_chart(fig_time, use_container_width=True)

    # Hourly heatmap
    st.subheader("Hourly Distribution")
    hourly_dist = filtered_df.groupby('Date/Time').size().reset_index(name='Count')
    fig_heatmap = px.density_heatmap(
        filtered_df,
        x=filtered_df['Date/Time'].dt.hour,
        y=filtered_df['Date/Time'].dt.date,
        z=filtered_df.groupby([filtered_df['Date/Time'].dt.date, filtered_df['Date/Time'].dt.hour]).size().values.reshape(-1, 24),
        color_continuous_scale="Viridis"
    )
    st.plotly_chart(fig_heatmap, use_container_width=True)
else:
    st.info("No data matches the selected filters.")