import dash
from dash import dcc, html, Input, Output, State, callback, ctx, dash_table
import dash_bootstrap_components as dbc
import duckdb

from dashboard_helpers import (
    get_neighborhood_geojson,
    get_all_neighborhoods,
    build_trip_choropleth,
    build_demo_choropleth,
    build_rr_bar_chart,
    build_rr_heatmap,
    build_neighborhood_profile,
    build_comparison_map,
    get_trip_stats_df,
    get_download_csv,
    _GRAPH_CONFIG,
)

DB_PATH = "sf_dashboard.db"

def get_con():
    c = duckdb.connect(DB_PATH, read_only=True)
    c.install_extension("spatial")
    c.load_extension("spatial")
    return c

_init_con = get_con()
GEOJSON = get_neighborhood_geojson(_init_con)
NEIGHBORHOODS = get_all_neighborhoods(_init_con)
MONTHS = _init_con.sql(
    "SELECT DISTINCT month FROM trip_counts_pu ORDER BY month"
).df()["month"].tolist()

baseline_df = _init_con.sql("SELECT * FROM city_baselines").df()
BASELINE_WHITE = float(baseline_df["baseline_white_pct"].iloc[0])
BASELINE_ASIAN = float(baseline_df["baseline_asian_pct"].iloc[0])
_init_con.close()

app = dash.Dash(
    __name__,
    external_stylesheets=[dbc.themes.DARKLY],
    meta_tags=[
        {"name": "viewport", "content": "width=device-width, initial-scale=1"}
    ],
    title="SF Taxi Mobility Equity Dashboard",
)
server = app.server

sidebar = dbc.Card(
    [
        html.H4("Controls", className="text-center mb-3"),
        html.Hr(),
        dbc.Label("Month"),
        dcc.Dropdown(
            id="month-selector",
            options=[{"label": m, "value": m} for m in MONTHS],
            value="Jan2024",
            clearable=False,
            className="mb-3",
        ),
        dbc.Label("Hail Type"),
        dbc.Checklist(
            id="hail-type-filter",
            options=[
                {"label": " Street-Hail", "value": "Street"},
                {"label": " App-Based", "value": "App"},
            ],
            value=["Street", "App"],
            className="mb-3",
        ),
        html.Hr(),
        html.Div(id="selected-nhood-display", className="mb-3"),
        dbc.Button(
            "Reset Selection",
            id="reset-selection-btn",
            color="secondary",
            size="sm",
            className="w-100 mb-2",
        ),
        dbc.Button(
            "Download CSV",
            id="download-btn",
            color="info",
            size="sm",
            className="w-100 mb-2",
        ),
        dbc.Button(
            "Download GeoJSON",
            id="download-geojson-btn",
            color="success",
            size="sm",
            className="w-100 mb-2",
        ),
        dcc.Download(id="csv-download"),
        dcc.Download(id="geojson-download"),
    ],
    body=True,
    id="sidebar",
    className="bg-dark",
    style={
        "position": "sticky",
        "top": "10px",
        "height": "calc(100vh - 20px)",
        "overflowY": "auto",
    },
)

_STAT_CARD_STYLE = {"height": "100%"}

insights_banner = dbc.Row(
    [
        dbc.Col(
            dbc.Card(
                [
                    html.H3(id="stat-total-trips", className="text-center mb-0"),
                    html.P("Total Trips", className="text-center text-muted small"),
                ],
                body=True,
                className="bg-dark border-secondary",
                style=_STAT_CARD_STYLE,
            ),
            md=3,
        ),
        dbc.Col(
            dbc.Card(
                [
                    html.H3(id="stat-top-nhood", className="text-center mb-0",
                             style={"fontSize": "1.6rem"}),
                    html.P("Most Served", className="text-center text-muted small"),
                ],
                body=True,
                className="bg-dark border-secondary",
                style=_STAT_CARD_STYLE,
            ),
            md=3,
        ),
        dbc.Col(
            dbc.Card(
                [
                    html.H3(id="stat-rr-white", className="text-center mb-0"),
                    html.P("White RR", className="text-center text-muted small"),
                ],
                body=True,
                className="bg-dark border-secondary",
                style=_STAT_CARD_STYLE,
            ),
            md=3,
        ),
        dbc.Col(
            dbc.Card(
                [
                    html.H3(id="stat-rr-asian", className="text-center mb-0"),
                    html.P("Asian RR", className="text-center text-muted small"),
                ],
                body=True,
                className="bg-dark border-secondary",
                style=_STAT_CARD_STYLE,
            ),
            md=3,
        ),
    ],
    id="insights-banner",
    className="mb-2 g-2",
)

narrative_row = dbc.Row(
    dbc.Col(
        html.P(
            id="narrative-text",
            className="text-center fst-italic",
            style={"color": "#adb5bd", "fontSize": "0.95rem"},
        ),
    ),
    className="mb-3",
)

trip_maps = dbc.Row(
    [
        dbc.Col(
            dcc.Graph(id="street-pu-map", config=_GRAPH_CONFIG),
            md=6,
        ),
        dbc.Col(
            dcc.Graph(id="app-pu-map", config=_GRAPH_CONFIG),
            md=6,
        ),
        dbc.Col(
            dcc.Graph(id="street-do-map", config=_GRAPH_CONFIG),
            md=6,
            className="mt-2",
        ),
        dbc.Col(
            dcc.Graph(id="app-do-map", config=_GRAPH_CONFIG),
            md=6,
            className="mt-2",
        ),
    ],
    className="mb-3",
)

demographics_tab = dbc.Row(
    [
        dbc.Col(dcc.Graph(id="white-deviation-map", config=_GRAPH_CONFIG), md=6),
        dbc.Col(dcc.Graph(id="asian-deviation-map", config=_GRAPH_CONFIG), md=6),
    ]
)

rr_tab = html.Div(
    [
        dbc.Row(
            [
                dbc.Col(dcc.Graph(id="rr-bar-chart", config=_GRAPH_CONFIG), md=7),
                dbc.Col(dcc.Graph(id="rr-heatmap", config=_GRAPH_CONFIG), md=5),
            ]
        ),
    ],
    id="rr-section",
)

comparison_tab = html.Div(
    [
        dbc.Row(
            [
                dbc.Col(
                    [
                        dbc.Label("Month A"),
                        dcc.Dropdown(
                            id="comp-month-a",
                            options=[{"label": m, "value": m} for m in MONTHS],
                            value="Jan2024",
                            clearable=False,
                        ),
                    ],
                    md=3,
                ),
                dbc.Col(
                    [
                        dbc.Label("Month B"),
                        dcc.Dropdown(
                            id="comp-month-b",
                            options=[{"label": m, "value": m} for m in MONTHS],
                            value="Mar2024",
                            clearable=False,
                        ),
                    ],
                    md=3,
                ),
                dbc.Col(
                    [
                        dbc.Label("Hail Type"),
                        dcc.Dropdown(
                            id="comp-hail",
                            options=[
                                {"label": "Street-Hail", "value": "Street"},
                                {"label": "App-Based", "value": "App"},
                            ],
                            value="Street",
                            clearable=False,
                        ),
                    ],
                    md=3,
                ),
                dbc.Col(
                    [
                        dbc.Label("Metric"),
                        dcc.Dropdown(
                            id="comp-metric",
                            options=[
                                {"label": "Pickups", "value": "pu"},
                                {"label": "Drop-offs", "value": "do"},
                            ],
                            value="pu",
                            clearable=False,
                        ),
                    ],
                    md=3,
                ),
            ],
            className="mb-3",
        ),
        dbc.Row(dbc.Col(dcc.Graph(id="comparison-map", config=_GRAPH_CONFIG))),
    ]
)

top10_section = dbc.Row(
    [
        dbc.Col(
            dbc.Card(
                [
                    html.H5(
                        id="street-stats-title",
                        className="text-center",
                    ),
                    dbc.Row(
                        [
                            dbc.Col(
                                [
                                    html.H6("Pickups", className="text-center text-muted"),
                                    html.Div(id="street-pu-table"),
                                ],
                                md=6,
                            ),
                            dbc.Col(
                                [
                                    html.H6("Drop-offs", className="text-center text-muted"),
                                    html.Div(id="street-do-table"),
                                ],
                                md=6,
                            ),
                        ]
                    ),
                ],
                body=True,
                className="bg-dark border-secondary",
            ),
            md=6,
        ),
        dbc.Col(
            dbc.Card(
                [
                    html.H5(
                        id="app-stats-title",
                        className="text-center",
                    ),
                    dbc.Row(
                        [
                            dbc.Col(
                                [
                                    html.H6("Pickups", className="text-center text-muted"),
                                    html.Div(id="app-pu-table"),
                                ],
                                md=6,
                            ),
                            dbc.Col(
                                [
                                    html.H6("Drop-offs", className="text-center text-muted"),
                                    html.Div(id="app-do-table"),
                                ],
                                md=6,
                            ),
                        ]
                    ),
                ],
                body=True,
                className="bg-dark border-secondary",
            ),
            md=6,
        ),
    ],
    className="mt-3",
)

analysis_tabs = dbc.Tabs(
    [
        dbc.Tab(demographics_tab, label="Demographics", tab_id="tab-demo"),
        dbc.Tab(rr_tab, label="Representative Ratios", tab_id="tab-rr"),
        dbc.Tab(comparison_tab, label="Monthly Comparison", tab_id="tab-comp"),
    ],
    id="analysis-tabs",
    active_tab="tab-demo",
    className="mb-3",
)

nhood_offcanvas = dbc.Offcanvas(
    html.Div(id="nhood-detail-content"),
    id="nhood-offcanvas",
    title="Neighborhood Profile",
    placement="end",
    is_open=False,
    style={"width": "400px", "backgroundColor": "#303030"},
)

app.layout = dbc.Container(
    [
        dcc.Store(id="selected-neighborhood", data=None),
        nhood_offcanvas,
        # Header
        dbc.Row(
            dbc.Col(
                [
                    html.H2(
                        "SF Taxi Mobility Equity Dashboard",
                        className="text-center mt-3 mb-1",
                    ),
                    html.P(
                        "Analyzing whether Street-Hail and App-Based taxi services "
                        "in San Francisco are equitably distributed across "
                        "neighborhoods with different demographic compositions.",
                        className="text-center text-muted mb-3",
                        style={"maxWidth": "700px", "margin": "0 auto"},
                    ),
                ]
            )
        ),
        # Body: sidebar + main
        dbc.Row(
            [
                dbc.Col(sidebar, md=2, className="pe-1"),
                dbc.Col(
                    [
                        insights_banner,
                        narrative_row,
                        html.H5("Trip Distribution Maps", className="mb-2"),
                        trip_maps,
                        html.H5("Analysis", className="mb-2"),
                        analysis_tabs,
                        html.H5("Top 10 Neighborhoods", className="mb-2"),
                        top10_section,
                    ],
                    md=10,
                ),
            ]
        ),
        # Footer
        dbc.Row(
            dbc.Col(
                html.P(
                    [
                        "Data: ",
                        html.A("DataSF Taxi Trips", href="https://data.sfgov.org/Transportation/Taxi-Trips/m8hk-2ipk/", target="_blank"),
                        " | ",
                        html.A("ACS 2022", href="https://api.census.gov/data/2022/acs/acs5.html", target="_blank"),
                        " | Debayan Mandal",
                    ],
                    className="text-center text-muted small mt-4 mb-3",
                )
            )
        ),
    ],
    fluid=True,
)

@callback(
    Output("street-pu-map", "figure"),
    Output("app-pu-map", "figure"),
    Output("street-do-map", "figure"),
    Output("app-do-map", "figure"),
    Input("month-selector", "value"),
    Input("hail-type-filter", "value"),
    Input("selected-neighborhood", "data"),
)
def update_trip_maps(month, hail_types, sel_nhood):
    hail_types = hail_types or ["Street", "App"]
    con = get_con()
    figs = []
    for ht, metric in [
        ("Street", "pu"),
        ("App", "pu"),
        ("Street", "do"),
        ("App", "do"),
    ]:
        if ht in hail_types:
            figs.append(
                build_trip_choropleth(con, GEOJSON, ht, month, metric, sel_nhood)
            )
        else:
            import plotly.graph_objects as go
            fig = go.Figure()
            fig.update_layout(
                title=f"{ht} {'Pickups' if metric == 'pu' else 'Drop-offs'} (filtered out)",
                paper_bgcolor="rgba(0,0,0,0)",
                font_color="#e0e0e0",
                height=420,
            )
            figs.append(fig)
    return figs[0], figs[1], figs[2], figs[3]

@callback(
    Output("selected-neighborhood", "data"),
    Input("street-pu-map", "clickData"),
    Input("app-pu-map", "clickData"),
    Input("street-do-map", "clickData"),
    Input("app-do-map", "clickData"),
    Input("white-deviation-map", "clickData"),
    Input("asian-deviation-map", "clickData"),
    Input("reset-selection-btn", "n_clicks"),
    prevent_initial_call=True,
)
def sync_map_selection(c1, c2, c3, c4, c5, c6, reset):
    trigger = ctx.triggered_id
    if trigger == "reset-selection-btn":
        return None
    for click in [c1, c2, c3, c4, c5, c6]:
        if click and trigger in [
            "street-pu-map", "app-pu-map", "street-do-map", "app-do-map",
            "white-deviation-map", "asian-deviation-map",
        ]:
            try:
                return click["points"][0]["customdata"][0]
            except (KeyError, IndexError, TypeError):
                try:
                    return click["points"][0]["location"]
                except (KeyError, IndexError, TypeError):
                    pass
    return dash.no_update

@callback(
    Output("selected-nhood-display", "children"),
    Input("selected-neighborhood", "data"),
)
def display_selected_nhood(nhood):
    if nhood:
        return dbc.Alert(
            [html.Strong("Selected: "), nhood],
            color="info",
            className="py-2 mb-0",
        )
    return html.P("Click a neighborhood on any map", className="text-muted small")

@callback(
    Output("stat-total-trips", "children"),
    Output("stat-top-nhood", "children"),
    Output("stat-rr-white", "children"),
    Output("stat-rr-asian", "children"),
    Output("narrative-text", "children"),
    Input("month-selector", "value"),
    Input("hail-type-filter", "value"),
)
def update_insights(month, hail_types):
    hail_types = hail_types or ["Street", "App"]
    ht_filter = ", ".join(f"'{h}'" for h in hail_types)
    con = get_con()

    total = con.sql(f"""
        SELECT SUM(trips_pu) AS total
        FROM trip_counts_pu
        WHERE month = '{month}' AND hail_type IN ({ht_filter})
    """).df()["total"].iloc[0]
    total = int(total) if total else 0

    top = con.sql(f"""
        SELECT nhood, SUM(trips_pu) AS t
        FROM trip_counts_pu
        WHERE month = '{month}' AND hail_type IN ({ht_filter})
        GROUP BY nhood ORDER BY t DESC LIMIT 1
    """).df()
    top_nhood = top["nhood"].iloc[0] if not top.empty else "N/A"

    rr = con.sql(f"""
        SELECT AVG(RR_white_PU) AS rr_w, AVG(RR_asian_PU) AS rr_a
        FROM representative_ratios
        WHERE month = '{month}' AND hail_type IN ({ht_filter})
    """).df()
    rr_w = round(float(rr["rr_w"].iloc[0]), 2) if not rr.empty and rr["rr_w"].iloc[0] else 0
    rr_a = round(float(rr["rr_a"].iloc[0]), 2) if not rr.empty and rr["rr_a"].iloc[0] else 0

    # Dynamic narrative
    parts = [f"In {month}, {total:,} taxi trips were recorded across SF."]
    parts.append(f"{top_nhood} was the most served neighborhood.")
    if rr_w > 1.0:
        parts.append(
            f"White-majority neighborhoods received {rr_w:.2f}x their expected "
            f"share of service,"
        )
    if rr_a < 1.0:
        parts.append(
            f"and Asian-majority neighborhoods received {rr_a:.2f}x."
        )
    narrative = " ".join(parts)

    return f"{total:,}", top_nhood, f"{rr_w:.2f}x", f"{rr_a:.2f}x", narrative

@callback(
    Output("white-deviation-map", "figure"),
    Output("asian-deviation-map", "figure"),
    Input("selected-neighborhood", "data"),
)
def update_demo_maps(sel_nhood):
    con = get_con()
    w = build_demo_choropleth(
        con, GEOJSON, "white_pct", BASELINE_WHITE,
        f"White Pop. Deviation ({BASELINE_WHITE:.1f}% baseline)",
        sel_nhood,
    )
    a = build_demo_choropleth(
        con, GEOJSON, "asian_pct", BASELINE_ASIAN,
        f"Asian Pop. Deviation ({BASELINE_ASIAN:.1f}% baseline)",
        sel_nhood,
    )
    return w, a

@callback(
    Output("rr-bar-chart", "figure"),
    Output("rr-heatmap", "figure"),
    Input("month-selector", "value"),
)
def update_rr(month):
    con = get_con()
    return build_rr_bar_chart(con, month), build_rr_heatmap(con)

@callback(
    Output("comparison-map", "figure"),
    Input("comp-month-a", "value"),
    Input("comp-month-b", "value"),
    Input("comp-hail", "value"),
    Input("comp-metric", "value"),
)
def update_comparison(month_a, month_b, hail, metric):
    con = get_con()
    return build_comparison_map(con, GEOJSON, hail, metric, month_a, month_b)

@callback(
    Output("nhood-offcanvas", "is_open"),
    Output("nhood-detail-content", "children"),
    Input("selected-neighborhood", "data"),
    State("month-selector", "value"),
)
def update_nhood_panel(nhood, month):
    if not nhood:
        return False, []

    con = get_con()
    profile = build_neighborhood_profile(con, nhood, month)
    demo = profile["demographics"]

    children = [
        html.H4(profile["name"]),
        html.Hr(),
    ]

    if demo:
        children.extend(
            [
                html.H6("Demographics"),
                html.P(f"Population: {demo['total_pop']:,}"),
                html.P(f"White: {demo['white_pct']}% (deviation: {demo['white_dev']:+.1f} pp)"),
                html.P(f"Black: {demo['black_pct']}%"),
                html.P(f"Asian: {demo['asian_pct']}% (deviation: {demo['asian_dev']:+.1f} pp)"),
                html.Hr(),
            ]
        )

    if profile["trips"]:
        children.append(html.H6(f"Trips ({month})"))
        for key, val in sorted(profile["trips"].items()):
            if month in key:
                label = key.replace(f"_{month}", "").replace("_", " ")
                children.append(html.P(f"{label}: {val:,}"))
        children.append(html.Hr())

    if profile["trend_fig"]:
        children.append(
            dcc.Graph(
                figure=profile["trend_fig"],
                config={"displayModeBar": False},
                style={"height": "280px"},
            )
        )

    return True, children

def _make_table(df):
    if df.empty:
        return html.P("No data", className="text-muted")
    return dash_table.DataTable(
        data=df.to_dict("records"),
        columns=[{"name": c, "id": c} for c in df.columns],
        style_table={"overflowX": "auto"},
        style_header={
            "backgroundColor": "#375a7f",
            "color": "white",
            "fontWeight": "bold",
            "textAlign": "center",
        },
        style_cell={
            "backgroundColor": "#303030",
            "color": "#e0e0e0",
            "textAlign": "center",
            "padding": "6px",
            "fontSize": "0.85rem",
        },
        style_data_conditional=[
            {
                "if": {"row_index": 0},
                "backgroundColor": "#3a506b",
                "fontWeight": "bold",
            }
        ],
        page_size=10,
    )


@callback(
    Output("street-stats-title", "children"),
    Output("street-pu-table", "children"),
    Output("street-do-table", "children"),
    Output("app-stats-title", "children"),
    Output("app-pu-table", "children"),
    Output("app-do-table", "children"),
    Input("month-selector", "value"),
)
def update_top10(month):
    con = get_con()
    s_pu = get_trip_stats_df(con, "Street", month, "pu")
    s_do = get_trip_stats_df(con, "Street", month, "do")
    a_pu = get_trip_stats_df(con, "App", month, "pu")
    a_do = get_trip_stats_df(con, "App", month, "do")

    return (
        f"Street-Hail Top 10 ({month})",
        _make_table(s_pu),
        _make_table(s_do),
        f"App-Based Top 10 ({month})",
        _make_table(a_pu),
        _make_table(a_do),
    )

@callback(
    Output("csv-download", "data"),
    Input("download-btn", "n_clicks"),
    State("month-selector", "value"),
    State("hail-type-filter", "value"),
    State("selected-neighborhood", "data"),
    prevent_initial_call=True,
)
def trigger_download(n_clicks, month, hail_types, nhood):
    con = get_con()
    df = get_download_csv(con, month, hail_types, nhood)
    filename = f"sf_taxi_data_{month}"
    if nhood:
        filename += f"_{nhood.replace(' ', '_').replace('/', '_')}"
    return dcc.send_data_frame(df.to_csv, f"{filename}.csv", index=False)

@callback(
    Output("geojson-download", "data"),
    Input("download-geojson-btn", "n_clicks"),
    State("month-selector", "value"),
    State("hail-type-filter", "value"),
    State("selected-neighborhood", "data"),
    prevent_initial_call=True,
)
def trigger_geojson_download(n_clicks, month, hail_types, nhood):
    import json
    import copy
    con = get_con()
    hail_types = hail_types or ["Street", "App"]
    ht_filter = ", ".join(f"'{h}'" for h in hail_types)
    nhood_clause = f"AND n.nhood = '{nhood}'" if nhood else ""

    # Get trip + demographic data per neighborhood
    data_df = con.sql(f"""
        SELECT n.nhood,
               COALESCE(SUM(pu.trips_pu), 0) AS total_pickups,
               COALESCE(SUM(td.trips_do), 0) AS total_dropoffs,
               nd.total_pop, nd.white_pct, nd.black_pct, nd.asian_pct
        FROM neighborhoods n
        LEFT JOIN trip_counts_pu pu
          ON pu.nhood = n.nhood AND pu.month = '{month}'
          AND pu.hail_type IN ({ht_filter})
        LEFT JOIN trip_counts_do td
          ON td.nhood = n.nhood AND td.month = '{month}'
          AND td.hail_type IN ({ht_filter})
        LEFT JOIN nhood_demographics nd ON nd.nhood = n.nhood
        WHERE 1=1 {nhood_clause}
        GROUP BY n.nhood, nd.total_pop, nd.white_pct, nd.black_pct, nd.asian_pct
    """).df()
    data_map = {row["nhood"]: row.to_dict() for _, row in data_df.iterrows()}

    geojson = copy.deepcopy(GEOJSON)
    # Filter to selected neighborhood if one is chosen
    if nhood:
        geojson["features"] = [
            f for f in geojson["features"]
            if f["properties"]["nhood"] == nhood
        ]
    # Enrich features with data
    for feat in geojson["features"]:
        name = feat["properties"]["nhood"]
        if name in data_map:
            d = data_map[name]
            feat["properties"]["month"] = month
            feat["properties"]["total_pickups"] = int(d["total_pickups"])
            feat["properties"]["total_dropoffs"] = int(d["total_dropoffs"])
            feat["properties"]["total_pop"] = int(d["total_pop"]) if d["total_pop"] else 0
            feat["properties"]["white_pct"] = round(float(d["white_pct"]), 1) if d["white_pct"] else 0
            feat["properties"]["black_pct"] = round(float(d["black_pct"]), 1) if d["black_pct"] else 0
            feat["properties"]["asian_pct"] = round(float(d["asian_pct"]), 1) if d["asian_pct"] else 0

    filename = f"sf_taxi_{month}"
    if nhood:
        filename += f"_{nhood.replace(' ', '_').replace('/', '_')}"
    return dict(
        content=json.dumps(geojson, indent=2),
        filename=f"{filename}.geojson",
        type="application/geo+json",
    )

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)