"""
dashboard_helpers.py  —  Map-building and chart functions used by app.py
"""

import numpy as np
import pandas as pd
import geopandas as gpd
import leafmap
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.colors as mcolors


def _classify_terciles(series):
    out = pd.Series(-1, index=series.index, dtype=int)
    mask = series.notna()
    if mask.sum() < 3:
        return out
    labels = pd.qcut(series[mask], q=3, labels=[0, 1, 2], duplicates="drop")
    out[mask] = labels.astype(int)
    return out


# Bivariate palette
# X axis = environmental variable (NDVI or TES): color (red=low, yellow=mid, green=high)
# Y axis = demographic variable: shade (dark=high % community of concern, light=low % community of concern)

BIVARIATE_COLORS = [
    "#8B0000",
    "#B8860B",
    "#1B5E20",
    "#CD5C5C",
    "#DAA520",
    "#4CAF50",
    "#F4CCCC",
    "#FFF9C4",
    "#C8E6C9",
]

CMAP_NAME = "bivariate_9"


# Bivariate choropleth map

def build_bivariate_map(con, x_col, x_label, y_col, y_label):
    df = con.sql(f"""
        SELECT GEOID, geometry_wkt, {x_col}, {y_col}
        FROM block_groups
        WHERE total_pop > 0
          AND {x_col} IS NOT NULL
          AND {y_col} IS NOT NULL
    """).df()

    gdf = gpd.GeoDataFrame(
        df,
        geometry=gpd.GeoSeries.from_wkt(df["geometry_wkt"]),
        crs="EPSG:4269"
    ).drop(columns="geometry_wkt")

    gdf["x_class"] = _classify_terciles(gdf[x_col])
    # For minority %, invert so that high minority = dark
    if y_col == "pct_minority":
        gdf["y_class"] = _classify_terciles(gdf[y_col]).map({2: 0, 1: 1, 0: 2, -1: -1})
    else:
        gdf["y_class"] = _classify_terciles(gdf[y_col])
    
    gdf["bv_class"] = gdf.apply(
        lambda r: int(r["y_class"]) * 3 + int(r["x_class"])
        if r["x_class"] >= 0 and r["y_class"] >= 0 else float("nan"),
        axis=1
    )

    gdf_valid = gdf[gdf["bv_class"].notna()].copy()
    gdf_valid["bv_class"] = gdf_valid["bv_class"].astype(float)

    cmap = mcolors.ListedColormap(BIVARIATE_COLORS, name=CMAP_NAME)
    try:
        matplotlib.colormaps.register(cmap, name=CMAP_NAME)
    except Exception:
        pass

    gdf_valid["hex_color"] = gdf_valid["bv_class"].apply(
        lambda c: BIVARIATE_COLORS[int(c)] if not np.isnan(c) else "#cccccc"
    )

    # Save to a temp file and load via leafmap
    import tempfile, os
    tmp = tempfile.NamedTemporaryFile(suffix=".geojson", delete=False)
    tmp.close()
    gdf_valid[["geometry", "bv_class", "hex_color", "GEOID"]].to_file(tmp.name, driver="GeoJSON")

    import folium
    m = folium.Map(location=[33.45, -112.07], zoom_start=10,
                   tiles="CartoDB dark_matter")

    for _, row in gdf_valid.iterrows():
        folium.GeoJson(
            row.geometry.__geo_interface__,
            style_function=lambda f, c=row["hex_color"]: {
                "fillColor": c,
                "fillOpacity": 0.75,
                "color": "#444",
                "weight": 0.4,
            }
        ).add_to(m)

    return m


# Bivariate legend

def build_bivariate_legend(x_label, y_label, y_col):
    palette = [
        ["#8B0000", "#B8860B", "#1B5E20"],  # low Y:  dark shades
        ["#CD5C5C", "#DAA520", "#4CAF50"],  # mid Y:  medium shades
        ["#F4CCCC", "#FFF9C4", "#C8E6C9"],  # high Y: light shades
    ]

    fig, ax = plt.subplots(figsize=(3.5, 3.5), facecolor="#1a1a2e")
    ax.set_facecolor("#1a1a2e")

    for row in range(3):
        for col in range(3):
            rect = mpatches.FancyBboxPatch(
                (col, row), 1, 1,
                boxstyle="square,pad=0",
                facecolor=palette[row][col],
                edgecolor="#333", linewidth=0.5
            )
            ax.add_patch(rect)

    ax.set_xlim(0, 3)
    ax.set_ylim(0, 3)

    ax.set_xticks([0.5, 1.5, 2.5])
    ax.set_xticklabels(["Low", "Med", "High"], color="#c9d1d9", fontsize=8)

    ax.set_yticks([0.5, 1.5, 2.5])
    if y_col == "pct_minority":
        ax.set_yticklabels(["High", "Med", "Low"], color="#c9d1d9", fontsize=8)
    else:
        ax.set_yticklabels(["Low", "Med", "High"], color="#c9d1d9", fontsize=8)

    ax.set_xlabel(x_label, color="#c9d1d9", fontsize=9, labelpad=8)
    ax.set_ylabel(y_label, color="#c9d1d9", fontsize=9, labelpad=8)
    ax.tick_params(length=0)
    for spine in ax.spines.values():
        spine.set_visible(False)

    fig.subplots_adjust(left=0.22, bottom=0.18, right=0.95, top=0.95)
    return fig

# Summary stats table

def build_summary_stats(con, x_col, y_col):
    df = con.sql(f"""
        SELECT {x_col}, {y_col}, tree_equity_score, ndvi_mean, total_pop
        FROM block_groups
        WHERE total_pop > 0
          AND {x_col} IS NOT NULL
          AND {y_col} IS NOT NULL
    """).df()

    df["x_tercile"] = _classify_terciles(df[x_col]).map(
        {-1: "No data", 0: f"Low", 1: f"Mid", 2: f"High"}
    )

    summary = (
        df.groupby("x_tercile")
        .agg(
            Block_Groups=(x_col, "count"),
            **{f"Avg_{x_col}": (x_col, lambda x: round(x.mean(), 2))},
            **{f"Avg_{y_col}": (y_col, lambda x: round(x.mean(), 1))},
            Avg_TES=("tree_equity_score", lambda x: round(x.mean(), 1) if x.notna().any() else np.nan),
            Avg_NDVI=("ndvi_mean", lambda x: round(x.mean(), 3) if x.notna().any() else np.nan),
        )
        .reindex(["Low", "Mid", "High"])
        .reset_index()
        .rename(columns={"x_tercile": "Group"})
    )

    styled = summary.style.background_gradient(
        cmap="RdYlGn", subset=[f"Avg_{x_col}"]
    ).hide(axis="index")

    return styled


# Tree-need bar chart

def build_need_chart(con, y_col):
    df = con.sql(f"""
        SELECT GEOID, ndvi_mean, tree_equity_score, {y_col}
        FROM block_groups
        WHERE ndvi_mean IS NOT NULL AND tree_equity_score IS NOT NULL
          AND {y_col} IS NOT NULL AND total_pop > 0
    """).df()

    if df.empty:
        fig, ax = plt.subplots(facecolor="#161b22")
        ax.text(0.5, 0.5, "No data yet", ha="center", va="center", color="#8b949e")
        return fig

    def minmax(s):
        mn, mx = s.min(), s.max()
        return (s - mn) / (mx - mn) if mx > mn else s * 0

    df["ndvi_norm"] = minmax(df["ndvi_mean"])
    df["tes_norm"]  = minmax(df["tree_equity_score"])
    df["need"]      = (1 - df["ndvi_norm"]) * 0.5 + (1 - df["tes_norm"]) * 0.5

    if y_col == "pct_minority":
        df["shade_norm"] = minmax(df[y_col])         # high = darker
    else:
        df["shade_norm"] = 1 - minmax(df[y_col])     # low income = darker

    top15 = df.nlargest(15, "need").reset_index(drop=True)
    labels = [str(g)[-6:] for g in top15["GEOID"]]
    x = np.arange(len(top15))

    fig, ax = plt.subplots(figsize=(10, 3.2), facecolor="#161b22")
    ax.set_facecolor("#161b22")

    for i, row in top15.iterrows():
        alpha = 0.3 + 0.7 * row["shade_norm"]
        ax.bar(i, row["need"], color="#e05c5c", alpha=alpha, width=0.6)

    # Shade legend
    y_label = "% Minority" if y_col == "pct_minority" else "Median Income"
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor="#e05c5c", alpha=1.0,
              label=f"High concern ({('high' if y_col == 'pct_minority' else 'low')} {y_label})"),
        Patch(facecolor="#e05c5c", alpha=0.3,
              label=f"Lower concern ({('low' if y_col == 'pct_minority' else 'high')} {y_label})"),
    ]

    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=45, ha="right", color="#8b949e", fontsize=8)
    ax.set_ylabel("Tree Need Score", color="#8b949e", fontsize=8)
    ax.set_title("Top 15 Block Groups by Tree Program Need  (GEOID suffix)",
                 color="#c9d1d9", fontsize=9, pad=6)
    ax.tick_params(colors="#8b949e", labelsize=8)
    for sp in ax.spines.values():
        sp.set_color("#30363d")
    ax.legend(handles=legend_elements, frameon=False,
              labelcolor="#8b949e", fontsize=8, loc="upper right")
    fig.tight_layout(pad=0.5)
    return fig

def build_need_map(con, y_col):
    df = con.sql(f"""
        SELECT GEOID, geometry_wkt, ndvi_mean, tree_equity_score, {y_col}
        FROM block_groups
        WHERE ndvi_mean IS NOT NULL AND tree_equity_score IS NOT NULL
          AND {y_col} IS NOT NULL AND total_pop > 0
    """).df()

    def minmax(s):
        mn, mx = s.min(), s.max()
        return (s - mn) / (mx - mn) if mx > mn else s * 0

    df["ndvi_norm"] = minmax(df["ndvi_mean"])
    df["tes_norm"]  = minmax(df["tree_equity_score"])
    df["need"]      = (1 - df["ndvi_norm"]) * 0.5 + (1 - df["tes_norm"]) * 0.5

    top15_geoids = set(df.nlargest(15, "need")["GEOID"].tolist())

    gdf = gpd.GeoDataFrame(
        df,
        geometry=gpd.GeoSeries.from_wkt(df["geometry_wkt"]),
        crs="EPSG:4269"
    ).drop(columns="geometry_wkt")

    import folium
    m = folium.Map(location=[33.45, -112.07], zoom_start=9,
                   tiles="CartoDB dark_matter")

    for _, row in gdf.iterrows():
        is_top15 = row["GEOID"] in top15_geoids
        folium.GeoJson(
            row.geometry.__geo_interface__,
            style_function=lambda f, top=is_top15: {
                "fillColor":   "#e05c5c" if top else "#444444",
                "fillOpacity": 0.85 if top else 0.3,
                "color":       "#e05c5c" if top else "#333333",
                "weight":      1.5 if top else 0.3,
            }
        ).add_to(m)

    return m