# Knoebels Queue Dashboard v1.1 - Scaled wait times fix
import streamlit as st
import pandas as pd
import psycopg2
from psycopg2.extras import RealDictCursor
import plotly.express as px
import plotly.graph_objects as go
import os
from datetime import datetime, timedelta
from dotenv import load_dotenv

# Load environment variables from .env if it exists
load_dotenv(os.path.join(os.path.dirname(__file__), '.env'))

# --- Configuration ---
st.set_page_config(
    page_title="Knoebels Queue Dashboard",
    page_icon="🎢",
    layout="wide",
)

# --- Database Connection ---
def get_db_connection():
    try:
        conn = psycopg2.connect(
            host=os.environ.get('DB_HOST', '192.168.88.9'),
            port=os.environ.get('DB_PORT', '5432'),
            database=os.environ.get('DB_NAME', 'gp0'),
            user=os.environ.get('DB_USER', 'kuhnobowls'),
            password=os.environ.get('DB_PASSWORD'),
            options="-c search_path=knoebels"
        )
        return conn
    except Exception as e:
        st.error(f"Database connection failed: {e}")
        return None

@st.cache_data(ttl=300)
def fetch_latest_status():
    conn = get_db_connection()
    if not conn: return pd.DataFrame()
    query = "SELECT * FROM dm_latest_ride_info"
    df = pd.read_sql(query, conn)
    conn.close()
    if not df.empty:
        # Safely convert and scale wait times to minutes
        df['queue_time'] = pd.to_numeric(df['queue_time'], errors='coerce')
        df['queue_time'] = (df['queue_time'] / 60.0).round(0)
        # Ensure coordinates are numeric
        df['lattitude'] = pd.to_numeric(df['lattitude'], errors='coerce')
        df['longitude'] = pd.to_numeric(df['longitude'], errors='coerce')
    return df

@st.cache_data(ttl=600)
def fetch_historical_trends(days=2):
    conn = get_db_connection()
    if not conn: return pd.DataFrame()
    query = f"""
    SELECT time_stamp, name, queue_time, is_operational, category
    FROM dm_knb_live_data 
    WHERE time_stamp > NOW() - INTERVAL '{days} days'
    ORDER BY time_stamp ASC
    """
    df = pd.read_sql(query, conn)
    conn.close()
    if not df.empty:
        # Safely convert and scale wait times to minutes
        df['queue_time'] = pd.to_numeric(df['queue_time'], errors='coerce')
        df['queue_time'] = (df['queue_time'] / 60.0).round(1)
    return df

@st.cache_data(ttl=3600)
def fetch_categories():
    conn = get_db_connection()
    if not conn: return {}
    query = 'SELECT _id, name FROM "LZ_attractions_io_categories"'
    with conn.cursor() as cur:
        cur.execute(query)
        cats = {row[0]: row[1] for row in cur.fetchall()}
    conn.close()
    return cats

# --- PCA: "reading the park's mind" ----------------------------------------
# Reproduces the analysis from the blog post (Reading a Theme Park's Mind with
# PCA). Standardize the time x ride wait matrix, take the top 3 components, and
# read them as: PC1 = overall busyness, PC2 = Kiddieland vs. main midway
# (spatial), PC3 = daytime school-group rides vs. the evening crowd.

PC_LABELS = {
    "PC1": "Overall busyness",
    "PC2": "Kiddieland ↔ main midway",
    "PC3": "Daytime crowd ↔ evening crowd",
}
PC_BLURB = {
    "PC1": "How many people are in the park. Every ride loads the same sign — "
           "when it's packed, everything spikes together.",
    "PC2": "Which half of the park they're in. A contrast: Kiddieland corner "
           "(negative) vs. the main midway (positive). It's geography, not the "
           "park's own category tags.",
    "PC3": "What time of day it is. Daytime family flat-rides the school/daycare "
           "groups hit at open (positive) vs. the evening coaster & date-night "
           "crowd (negative).",
}
CATEGORY_COLORS = {"Family": "#4C9BE8", "Kiddie": "#E8A33D", "Thrill": "#E0524E"}


@st.cache_data(ttl=21600)  # 6h — underlying history grows slowly, PCA is heavy
def compute_pca():
    """Pull the full wait history, build the standardized time x ride matrix,
    and return (loadings, scores, explained_variance_ratio, ride_meta).

    ride_meta carries each ride's category tag (trailing space trimmed) and GPS
    coordinates, so the same DataFrame drives the behavioral map."""
    from sklearn.preprocessing import StandardScaler
    from sklearn.decomposition import PCA

    conn = get_db_connection()
    if not conn:
        return None

    waits = pd.read_sql(
        """
        SELECT a.time_stamp, r.ride AS ride_name, a.queue_time AS queue_time_sec
          FROM knoebels."LZ_attractions_io" a
          JOIN knoebels.ride_master r ON r.attractions_dot_io_id = a._id
         WHERE a.queue_time IS NOT NULL
        """,
        conn,
    )
    meta = pd.read_sql(
        """
        SELECT r.ride AS ride_name,
               trim(c.name) AS category,
               (p.location)[0]::float AS lat,
               (p.location)[1]::float AS lon
          FROM knoebels.ride_master r
          JOIN knoebels."LZ_attractions_io_poi" p ON p._id = r.attractions_dot_io_id
          LEFT JOIN knoebels."LZ_attractions_io_categories" c ON c._id = p.category
        """,
        conn,
    ).set_index("ride_name")
    conn.close()

    if waits.empty:
        return None

    waits["time_stamp"] = pd.to_datetime(waits["time_stamp"])
    waits["queue_time_min"] = pd.to_numeric(waits["queue_time_sec"], errors="coerce") / 60.0

    # long -> (time x ride) matrix: mean wait per 5-min bin, ffill then zero
    matrix = (
        waits.assign(ts=waits["time_stamp"].dt.floor("5min"))
        .pivot_table(index="ts", columns="ride_name",
                     values="queue_time_min", aggfunc="mean")
        .sort_index()
        .ffill()
        .fillna(0)
        .drop(columns=["powersurge"], errors="ignore")  # ~90% imputed, noise
    )

    scaled = StandardScaler().fit_transform(matrix)  # each ride mean 0, var 1
    pca = PCA(n_components=3)
    sc = pca.fit_transform(scaled)

    scores = pd.DataFrame(sc, index=matrix.index, columns=["PC1", "PC2", "PC3"])
    loadings = pd.DataFrame(pca.components_.T, index=matrix.columns,
                            columns=["PC1", "PC2", "PC3"])
    loadings = loadings.join(meta)

    return {
        "loadings": loadings,
        "scores": scores,
        "evr": pca.explained_variance_ratio_,
    }

# --- Header ---
st.title("🎢 Knoebels Queue Time Dashboard")
st.markdown("---")

# --- Sidebar Filters ---
st.sidebar.header("Dashboard Filters")
latest_df = fetch_latest_status()
categories = fetch_categories()

if not latest_df.empty:
    all_cats = sorted(latest_df['category'].dropna().unique())
    cat_options = {c: categories.get(c, f"Category {c}") for c in all_cats}
    selected_cats = st.sidebar.multiselect(
        "Filter by Category", 
        options=all_cats, 
        format_func=lambda x: cat_options.get(x),
        default=all_cats
    )
    
    filtered_latest = latest_df[latest_df['category'].isin(selected_cats)]
else:
    st.warning("No data found in the database. Please check your connection and .env file.")
    st.stop()

# --- Tabs ---
tab1, tab2, tab3, tab4 = st.tabs(
    ["📍 Live Status", "📈 Historical Trends", "🔍 Ride Explorer", "🧠 Park Mind (PCA)"]
)

with tab1:
    # KPI Row
    col1, col2, col3, col4 = st.columns(4)
    open_rides = filtered_latest[filtered_latest['is_open'] == True]
    avg_wait = open_rides['queue_time'].mean()
    max_wait_row = open_rides.loc[open_rides['queue_time'].idxmax()] if not open_rides.empty else None
    
    col1.metric("Rides Open", len(open_rides))
    col2.metric("Avg Wait Time", f"{int(avg_wait)} min" if not pd.isna(avg_wait) else "N/A")
    if max_wait_row is not None:
        col3.metric("Longest Queue", f"{int(max_wait_row['queue_time'])} min", f"on {max_wait_row['name']}")
    else:
        col3.metric("Longest Queue", "N/A")
    col4.metric("Operational %", f"{int(len(open_rides)/len(filtered_latest)*100)}%" if len(filtered_latest) > 0 else "0%")

    # Map and Table
    mcol1, mcol2 = st.columns([2, 1])
    
    with mcol1:
        st.subheader("Park Map")
        
        # Prepare data for map: drop rows without valid coordinates
        map_df = filtered_latest.dropna(subset=['lattitude', 'longitude']).copy()
        
        if not map_df.empty:
            # Handle NaN wait times for sizing: fill with 0 and set minimum size
            map_df['wait_size'] = map_df['queue_time'].fillna(0).apply(lambda x: max(float(x), 5.0))
            
            # Plotly map for better control than st.map
            fig_map = px.scatter_mapbox(
                map_df,
                lat="lattitude",
                lon="longitude",
                color="queue_time",
                size="wait_size",
                hover_name="name",
                hover_data=["queue_status_message", "price"],
                color_continuous_scale="Viridis",
                zoom=15,
                height=600,
                mapbox_style="carto-darkmatter"
            )
            fig_map.update_layout(margin={"r":0,"t":0,"l":0,"b":0})
            st.plotly_chart(fig_map, use_container_width=True)
        else:
            st.info("No location data available to display map.")

    with mcol2:
        st.subheader("Live Board")
        board_df = filtered_latest[['name', 'queue_time', 'is_open', 'price']].sort_values('queue_time', ascending=False)
        st.dataframe(
            board_df, 
            column_config={
                "name": "Ride",
                "queue_time": st.column_config.NumberColumn("Wait (min)", format="%d ⏱️"),
                "is_open": "Status",
                "price": st.column_config.NumberColumn("Price", format="$%.2f")
            },
            hide_index=True,
            use_container_width=True
        )

with tab2:
    st.subheader("Wait Time Trends (Last 48 Hours)")
    hist_df = fetch_historical_trends(days=2)
    
    if not hist_df.empty:
        # Filter historical data by selected categories
        hist_df = hist_df[hist_df['category'].isin(selected_cats)]
        
        # Aggregate by timestamp
        avg_hist = hist_df.groupby('time_stamp')['queue_time'].mean().reset_index()
        
        fig_trend = px.line(
            avg_hist, 
            x='time_stamp', 
            y='queue_time', 
            title="Park-wide Average Wait Time",
            labels={'queue_time': 'Avg Wait (min)', 'time_stamp': 'Time'}
        )
        fig_trend.update_traces(line_color="#2E8B57")
        st.plotly_chart(fig_trend, use_container_width=True)
        
        # Heatmap
        st.subheader("Busiest Times of Day")
        hist_df['hour'] = pd.to_datetime(hist_df['time_stamp']).dt.hour
        hist_df['day'] = pd.to_datetime(hist_df['time_stamp']).dt.day_name()
        
        heatmap_data = hist_df.groupby(['day', 'hour'])['queue_time'].mean().unstack().fillna(0)
        # Order days
        days_order = ['Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday', 'Sunday']
        heatmap_data = heatmap_data.reindex(days_order)
        
        fig_heat = px.imshow(
            heatmap_data,
            labels=dict(x="Hour of Day", y="Day of Week", color="Avg Wait (min)"),
            x=heatmap_data.columns,
            y=heatmap_data.index,
            color_continuous_scale="YlGn"
        )
        st.plotly_chart(fig_heat, use_container_width=True)
    else:
        st.info("No historical data available for the selected period.")

with tab3:
    st.subheader("Ride Explorer")
    ride_names = sorted(latest_df['name'].unique())
    selected_ride = st.selectbox("Select a Ride to Explore", options=ride_names)
    
    ride_info = latest_df[latest_df['name'] == selected_ride].iloc[0]
    
    ecol1, ecol2 = st.columns([1, 2])
    
    with ecol1:
        st.markdown(f"### {ride_info['name']}")
        st.write(f"**Price:** ${ride_info['price']:.2f}")
        st.write(f"**Capacity:** {ride_info['capacity']} pph")
        st.write(f"**Duration:** {ride_info['duration']} sec")
        st.write(f"**Handstamp Included:** {'Yes' if ride_info['hand_stamp_included'] else 'No'}")
        st.write(f"**Height Requirement:** {ride_info['minimum_height_requirement']}\"")
        st.markdown("---")
        st.write(f"**Summary:** {ride_info['summary']}")
        if ride_info['restriction_summary']:
            st.warning(f"**Restrictions:** {ride_info['restriction_summary']}")

    with ecol2:
        st.subheader(f"7-Day History: {selected_ride}")
        # Fetch longer history for specific ride
        conn = get_db_connection()
        ride_id = ride_info['_id']
        ride_hist_query = f"SELECT time_stamp, queue_time FROM dm_knb_live_data WHERE _id = {ride_id} AND time_stamp > NOW() - INTERVAL '7 days' ORDER BY time_stamp ASC"
        ride_hist_df = pd.read_sql(ride_hist_query, conn)
        conn.close()
        
        if not ride_hist_df.empty:
            # Convert to minutes
            ride_hist_df['queue_time'] = pd.to_numeric(ride_hist_df['queue_time'], errors='coerce')
            ride_hist_df['queue_time'] = (ride_hist_df['queue_time'] / 60.0).round(1)
            
            fig_ride_hist = px.area(
                ride_hist_df, 
                x='time_stamp', 
                y='queue_time', 
                labels={'queue_time': 'Wait Time (min)', 'time_stamp': 'Time'}
            )
            fig_ride_hist.update_traces(line_color="#2E8B57", fillcolor="rgba(46, 139, 87, 0.3)")
            st.plotly_chart(fig_ride_hist, use_container_width=True)
        else:
            st.info("No historical data found for this ride.")

with tab4:
    st.subheader("🧠 What two seasons of wait times reveal about the park")
    st.markdown(
        "Feed every ride's wait-time history into **Principal Component Analysis** "
        "— with *no* labels for what any ride is or where it sits — and three "
        "hidden axes fall out on their own. They turn out to be **how busy the "
        "park is**, **which half of it you're in**, and **what time of day it is**. "
        "[Full write-up here.](https://blog.c0smere.net/reading-a-theme-parks-mindwith-pca/)"
    )

    try:
        pca = compute_pca()
    except ImportError:
        st.error("`scikit-learn` isn't installed in this image yet — rebuild the "
                 "dashboard container to enable the PCA tab.")
        pca = None

    if pca is None:
        st.info("Not enough data to compute the components right now.")
    else:
        loadings, scores, evr = pca["loadings"], pca["scores"], pca["evr"]

        # --- the three axes, with their explained variance ---
        st.markdown("#### The three axes PCA found")
        cols = st.columns(3)
        for col, pc in zip(cols, ["PC1", "PC2", "PC3"]):
            col.metric(f"{pc} · {PC_LABELS[pc]}", f"{evr[int(pc[-1]) - 1]:.1%}",
                       help=PC_BLURB[pc])
        st.caption(f"Just these 3 axes (out of {loadings.shape[0]}) capture "
                   f"**{evr.sum():.1%}** of all the variation in park wait times.")

        st.markdown("---")

        # --- the behavioral map: PC2 vs PC3, colored by the park's own tag ---
        st.markdown("#### The map of the park, drawn from wait times alone")
        st.markdown(
            "Each ride placed by its **PC2** (left↔right = Kiddieland↔midway) and "
            "**PC3** (bottom↔top = evening↔daytime) loading — its *behavioral* "
            "position — and colored by the park's category tag. The position comes "
            "purely from how each queue moves over time, yet **the picture is the "
            "park**: Kiddieland huddles left, evening thrill rides sink bottom-right, "
            "daytime family rides float to the top."
        )
        plot_df = loadings.reset_index().copy()
        plot_df["category"] = plot_df["category"].fillna("Untagged")
        fig_map_pca = px.scatter(
            plot_df, x="PC2", y="PC3", color="category", text="ride_name",
            color_discrete_map=CATEGORY_COLORS,
            labels={"PC2": "PC2  ·  Kiddieland ←→ main midway",
                    "PC3": "PC3  ·  evening crowd ←→ daytime crowd",
                    "category": "Park's tag"},
            hover_name="ride_name", height=620,
        )
        fig_map_pca.update_traces(textposition="top center",
                                  textfont=dict(size=9, color="#9aa0a6"),
                                  marker=dict(size=11, line=dict(width=1, color="#1A1C18")))
        fig_map_pca.add_hline(y=0, line_width=0.5, line_color="grey")
        fig_map_pca.add_vline(x=0, line_width=0.5, line_color="grey")
        fig_map_pca.update_layout(legend=dict(orientation="h", yanchor="bottom", y=1.02))
        st.plotly_chart(fig_map_pca, use_container_width=True)

        st.markdown("---")

        # --- loadings explorer: which rides drive a chosen component ---
        lcol, rcol = st.columns([1, 2])
        with lcol:
            st.markdown("#### Read an axis")
            pick = st.radio("Component", ["PC1", "PC2", "PC3"],
                            format_func=lambda p: f"{p} — {PC_LABELS[p]}")
            st.info(PC_BLURB[pick])
        with rcol:
            ranked = loadings[[pick]].copy().sort_values(pick)
            ranked["ride"] = ranked.index
            lim = ranked[pick].abs().max()  # symmetric so the scale's white sits at 0
            fig_load = px.bar(
                ranked, x=pick, y="ride", orientation="h",
                color=pick, color_continuous_scale="RdBu", range_color=[-lim, lim],
                labels={pick: f"{pick} loading", "ride": ""},
                height=max(420, 16 * len(ranked)),
            )
            fig_load.update_layout(coloraxis_showscale=False,
                                   margin={"l": 0, "r": 0, "t": 10, "b": 0})
            st.plotly_chart(fig_load, use_container_width=True)

        st.markdown("---")

        # --- the clock: each axis' average score by hour of day ---
        st.markdown("#### How each axis moves over a day")
        st.markdown(
            "Average score for each component by hour. **PC1** climbs into the "
            "evening as crowds build; **PC3** humps at midday (the school-group "
            "window) then falls — the fingerprint of *time of day*."
        )
        by_hour = scores.groupby(scores.index.hour).mean()
        by_hour.index.name = "hour"
        hour_long = by_hour.reset_index().melt("hour", var_name="Component",
                                               value_name="Mean score")
        fig_hour = px.line(
            hour_long, x="hour", y="Mean score", color="Component", markers=True,
            labels={"hour": "Hour of day (local)"},
            color_discrete_map={"PC1": "#2E8B57", "PC2": "#E8A33D", "PC3": "#4C9BE8"},
        )
        fig_hour.add_hline(y=0, line_width=0.5, line_color="grey")
        st.plotly_chart(fig_hour, use_container_width=True)

# --- Footer ---
st.sidebar.markdown("---")
st.sidebar.caption(f"Last updated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
if st.sidebar.button("Force Refresh Data"):
    st.cache_data.clear()
    st.rerun()