# app.py
import os
import re
import pandas as pd
import plotly.graph_objs as go
import plotly.express as px
import dash
from dash import dcc, html, Input, Output

# ========= CONFIG =========
EXCEL_MAIN = "WR_PMP_MultiCalib_Policy.xlsx"
EXCEL_MUN  = "crop_surface_subterraneotheft_parsed.xlsx"
BASE_YEAR  = 2015  # se assente, fallback = primo anno del gruppo

SCENARIO_DESCR = {
    "bau": "Business-as-usual: evoluzione osservata, senza nuove misure restrittive.",
    "scenario2": "Scenario 2: restrizioni/progressive o pricing moderato.",
    "scenario5": "Scenario 5: regolazione stringente con segnali di prezzo più forti.",
    "nwt": "No Water Theft: pieno enforcement e azzeramento dei prelievi illegali.",
}

# ========= HELPERS =========
def _first_existing_column(df, candidates, required=True):
    for c in candidates:
        if c in df.columns:
            return c
    if required:
        raise KeyError(f"Nessuna colonna trovata tra: {candidates}")
    return None

def _find_column_regex(df, pattern, required=False):
    """Find first column matching regex (case-insensitive)."""
    pat = re.compile(pattern, re.I)
    for col in df.columns:
        if pat.search(str(col)):
            return col
    if required:
        raise KeyError(f"Nessuna colonna che combacia con regex: {pattern}")
    return None

def municipio_label(code, mapper):
    if pd.isna(code):
        return "NA"
    code = int(code)
    name = mapper.get(code, str(code))
    return f"{name} ({code})"

def is_donana(v):
    return v == "DONANA"

def weighted_mean(group_df, value_col, weights_df):
    g = group_df.merge(weights_df, on="Municipio", how="left")
    g = g.dropna(subset=[value_col, "Hectares"])
    denom = g["Hectares"].sum()
    if denom == 0:
        return None
    return (g[value_col] * g["Hectares"]).sum() / denom

# ========= LOAD WR MAIN =========
def load_wr():
    xf = pd.ExcelFile(EXCEL_MAIN, engine="openpyxl")
    water_sheets = [s for s in xf.sheet_names if s.startswith("WaterSimulations_")]
    crop_sheets  = [s for s in xf.sheet_names if s.startswith("CropWaterSim_")]
    if not water_sheets:
        raise ValueError("Nessun foglio 'WaterSimulations_*' trovato.")
    if not crop_sheets:
        raise ValueError("Nessun foglio 'CropWaterSim_*' trovato.")

    df_gm  = pd.concat([pd.read_excel(xf, s) for s in water_sheets], ignore_index=True)
    df_crp = pd.concat([pd.read_excel(xf, s) for s in crop_sheets],  ignore_index=True)

    # Normalizza colonne base
    df_gm  = df_gm.rename(columns={"Policy Scenario": "Scenario"})
    df_crp = df_crp.rename(columns={"Policy Scenario": "Scenario"})
    for d in (df_gm, df_crp):
        d["Scenario"]  = d["Scenario"].astype(str).str.lower()
        d["Municipio"] = pd.to_numeric(d["Municipio"], errors="coerce").astype("Int64")
        d["Year"]      = pd.to_numeric(d["Year"], errors="coerce").astype("Int64")

    # MainInfo -> ettari per municipio (pesi)
    if "MainInfo" not in xf.sheet_names:
        raise ValueError("Foglio 'MainInfo' non trovato in WR_PMP_MultiCalib_Policy.xlsx")
    df_main = pd.read_excel(xf, sheet_name="MainInfo").rename(columns={"Policy Scenario": "Scenario"})
    df_main["Scenario"]  = df_main["Scenario"].astype(str).str.lower()
    df_main["Municipio"] = pd.to_numeric(df_main["Municipio"], errors="coerce").astype("Int64")
    hect_col = _first_existing_column(df_main, ["Total hectares", "Total Hectares", "Hectares"])
    weights = (df_main.groupby("Municipio", as_index=False)[hect_col]
                      .first()
                      .rename(columns={hect_col: "Hectares"}))
    return df_gm, df_crp, weights

df_gm, df_crop, w_hect = load_wr()
df_crop["Crop"] = df_crop["Crop"].astype(str).str.strip()
KNOWN_CROPS = set(df_crop["Crop"].dropna().astype(str).str.strip().unique())

# ========= MUNICIPIO NAMES + WATER REQUIREMENT (robusto) =========
def load_municipio_names_and_wr(known_crops):
    """
    Dal file 'crop_surface_subterraneotheft_parsed.xlsx':
    - Mapping codice -> nome municipio
    - Water requirement per crop×municipio = total water / total superficie  [m3/ha]
    Robusto a header diversi: cerca con regex e, se manca la colonna 'crop',
    la deduce confrontando valori con i crop di df_crop.
    """
    try:
        mxf = pd.ExcelFile(EXCEL_MUN, engine="openpyxl")
        df = pd.read_excel(mxf, sheet_name=mxf.sheet_names[0])
    except Exception as e:
        print(f"[WR-LOAD] Impossibile leggere {EXCEL_MUN}: {e}")
        return {}, pd.DataFrame(columns=["Municipio","Crop","wr_m3_ha"])

    # municipio code col (regex per 'municip')
    code_col = _find_column_regex(df, r"(id|cod|codigo).*(munic|muni)", required=False)
    if code_col is None:
        # fallback: prova una lista ampia
        try:
            code_col = _first_existing_column(df, [
                "ID PROVINCIA/MUNICIPIO","ID_PROVINCIA/MUNICIPIO","ID MUNICIPIO","ID_MUNICIPIO",
                "MUNICIPIO_ID","COD_MUNICIPIO","CODIGO_MUNICIPIO","COD MUNICIPIO","CODIGO MUNICIPIO"
            ])
        except Exception:
            # ultima spiaggia: qualunque colonna int-like con pochi NaN
            int_candidates = [c for c in df.columns if pd.api.types.is_integer_dtype(pd.to_numeric(df[c], errors="ignore"))]
            code_col = int_candidates[0] if int_candidates else df.columns[0]

    # municipio name col
    name_col = _find_column_regex(df, r"(munic|municipios|nombre.*munic)", required=False)
    if name_col is None:
        name_col = _first_existing_column(df, [
            "MUNICIPIOS","Municipio","NOMBRE_MUNICIPIO","NOMBRE MUNICIPIO","MUNICIPIO"
        ], required=False)

    # crop col (regex ampia: crop|cult|variedad|especie)
    crop_col = _find_column_regex(df, r"(crop|cultiv|coltur|variedad|especie|cult)", required=False)
    if crop_col is None:
        # prova elenco esteso
        for c in ["Crop","CROP","Cultivo","CULTIVO","CULTURE","Cultivar","Variedad","Especie","Nombre cultivo","Nombre Cultivo","CULTIVO NOMBRE"]:
            if c in df.columns:
                crop_col = c
                break
    if crop_col is None:
        # detezione per matching con KNOWN_CROPS
        best_col, best_overlap = None, -1
        cropy_cols = [c for c in df.columns if df[c].dtype == object]
        for c in cropy_cols:
            vals = set(df[c].dropna().astype(str).str.strip().unique())
            overlap = len(vals & known_crops)
            if overlap > best_overlap:
                best_overlap, best_col = overlap, c
        if best_col is not None and best_overlap > 0:
            crop_col = best_col
        else:
            print("[WR-LOAD] Colonna 'Crop' non trovata. Disabilito il grafico 'Total water consumption'.")
            # costruiamo comunque il mapper dei municipi se possibile:
            mapper = {}
            try:
                dd = df[[code_col, name_col]].copy() if name_col else df[[code_col]].copy()
                dd[code_col] = pd.to_numeric(dd[code_col], errors="coerce").astype("Int64")
                if name_col:
                    mapper = dict(zip(dd[code_col], dd[name_col]))
            except Exception as _:
                pass
            return mapper, pd.DataFrame(columns=["Municipio","Crop","wr_m3_ha"])

    # superficie & water col
    sup_col = _find_column_regex(df, r"(total\s*)?(superf|ha)", required=False)
    wat_col = _find_column_regex(df, r"(total\s*)?(agua|water|m3)", required=False)
    if sup_col is None:
        sup_col = _first_existing_column(df, [
            "total superficie","Total Superficie","TOTAL_SUPERFICIE","Superficie total","SUPERFICIE TOTAL",
            "Superficie","SUPERFICIE","Hectareas","Hectáreas","ha","HA"
        ], required=False)
    if wat_col is None:
        wat_col = _first_existing_column(df, [
            "total water","Total Water","TOTAL_WATER","Agua total","Total Agua","TOTAL AGUA",
            "m3","M3","m³","M³","Volumen","VOLUMEN"
        ], required=False)

    if sup_col is None or wat_col is None:
        print("[WR-LOAD] Colonne 'total superficie' o 'total water' non trovate. Disabilito il grafico consumo.")
        sup_col = sup_col or df.columns[0]
        wat_col = wat_col or df.columns[1]

    # mapper municipio (se possibile)
    mapper = {}
    try:
        dd = df[[code_col, name_col]].copy() if name_col else df[[code_col]].copy()
        dd[code_col] = pd.to_numeric(dd[code_col], errors="coerce").astype("Int64")
        if name_col:
            mapper = dict(zip(dd[code_col], dd[name_col]))
    except Exception as e:
        print(f"[WR-LOAD] Impossibile costruire il mapper municipio: {e}")

    # build wr table
    tmp = df[[code_col, crop_col, sup_col, wat_col]].copy()
    tmp[code_col] = pd.to_numeric(tmp[code_col], errors="coerce").astype("Int64")
    tmp = tmp.dropna(subset=[code_col, crop_col, sup_col, wat_col])
    tmp = tmp[pd.to_numeric(tmp[sup_col], errors="coerce") > 0]
    tmp["wr_m3_ha"] = pd.to_numeric(tmp[wat_col], errors="coerce") / pd.to_numeric(tmp[sup_col], errors="coerce")
    wr = tmp[[code_col, crop_col, "wr_m3_ha"]].rename(columns={code_col:"Municipio", crop_col:"Crop"})
    wr["Crop"] = wr["Crop"].astype(str).str.strip()
    return mapper, wr

mun_code_to_name, df_wr = load_municipio_names_and_wr(KNOWN_CROPS)

# ========= GM % vs base =========
group_keys = ["Scenario", "Municipio", "Method", "AM", "CM", "RCP"]
gm_base_2015 = (df_gm[df_gm["Year"].eq(BASE_YEAR)]
                .groupby(group_keys)["GM_PMP"].first().rename("GM_base_2015"))
first_year_base = (df_gm.sort_values("Year")
                   .groupby(group_keys)["GM_PMP"].first().rename("GM_base_fallback"))
gm_merged = (df_gm.merge(gm_base_2015, on=group_keys, how="left")
                 .merge(first_year_base, on=group_keys, how="left"))
gm_merged["GM_base"] = gm_merged["GM_base_2015"].fillna(gm_merged["GM_base_fallback"])
gm_merged["GM_perc"] = 100.0 * gm_merged["GM_PMP"] / gm_merged["GM_base"]

# ========= Water availability vs BAU baseline (100%) =========
is_bau = df_gm["Scenario"].eq("bau")
bau_first = (df_gm[is_bau].sort_values("Year")
             .groupby(["Municipio","Method","AM","CM","RCP"], as_index=False)
             .first()[["Municipio","Method","AM","CM","RCP","Water_PMP"]]
             .rename(columns={"Water_PMP":"Water_BAU_base"}))
water_all = df_gm.merge(bau_first, on=["Municipio","Method","AM","CM","RCP"], how="left")
water_all["Water_rel"] = 100.0 * water_all["Water_PMP"] / water_all["Water_BAU_base"]

# ========= DASH APP =========
external_stylesheets = ["https://fonts.googleapis.com/css2?family=Montserrat:wght@500;700&display=swap"]
app = dash.Dash(__name__, external_stylesheets=external_stylesheets)
server = app.server

# Dropdown values
scenarios = sorted(gm_merged["Scenario"].dropna().unique())
rcps      = sorted(gm_merged["RCP"].dropna().unique())
methods   = sorted(gm_merged["Method"].dropna().unique())
mun_codes = sorted(gm_merged["Municipio"].dropna().unique())
mun_options = [{"label": "Doñana (media pesata)", "value": "DONANA"}] + [
    {"label": municipio_label(c, mun_code_to_name), "value": int(c)} for c in mun_codes
]
methods_options = [{"label": "Tutti", "value": "Tutti"}] + [{"label": m, "value": m} for m in methods]

app.layout = html.Div(
    style={"fontFamily": "'Montserrat', sans-serif", "backgroundColor": "#f5f7fa", "minHeight": "100vh"},
    children=[
        html.Div(
            style={"background": "linear-gradient(90deg,#0a3967 0,#009ee0 100%)","padding": "22px 0 14px 0",
                   "marginBottom": "16px","textAlign": "center","color": "white","boxShadow": "0 6px 20px rgba(10,57,103,.19)"},
            children=[
                html.H1("🌱 Water Reallocation PMP Dashboard 👨‍🌾",
                        style={"fontWeight": "700", "fontSize": "2.3rem", "margin": "0"}),
                html.H4("WWF Doñana & IMDEA Water", style={"fontWeight": "500", "marginTop": "6px"})
            ]
        ),

        # Controls + Scenario description
        html.Div([
            html.Div([
                html.Label("Policy Scenario:", style={"marginRight": "8px", "fontWeight": "600"}),
                dcc.Dropdown(id='scenario', options=[{"label": s.upper(), "value": s} for s in scenarios],
                             value=scenarios[0] if scenarios else None, clearable=False,
                             style={'width': '180px','display': 'inline-block'}),

                html.Label("Municipio:", style={"marginLeft": "18px","marginRight":"8px","fontWeight":"600"}),
                dcc.Dropdown(id='municipio', options=mun_options, value="DONANA", clearable=False,
                             style={'width': '300px','display': 'inline-block'}),

                html.Label("RCP:", style={"marginLeft": "18px","marginRight":"8px","fontWeight":"600"}),
                dcc.Dropdown(id='rcp', options=[{"label": r, "value": r} for r in rcps],
                             value=rcps[0] if rcps else None, clearable=False,
                             style={'width': '150px','display': 'inline-block'}),

                html.Label("Method:", style={"marginLeft": "18px","marginRight":"8px","fontWeight":"600"}),
                dcc.Dropdown(id='method', options=methods_options, value="Tutti", clearable=False,
                             style={'width': '160px','display': 'inline-block'}),
            ], style={'padding': '16px 18px 10px 18px', "background":"#fff","borderRadius":"10px",
                      "boxShadow":"0 2px 12px rgba(10,57,103,.09)", "margin":"0 14px"}),

            html.Div(id="scenario-desc",
                     style={"margin":"8px 14px 0 14px","background":"#fff","borderRadius":"10px",
                            "boxShadow":"0 2px 12px rgba(10,57,103,.08)","padding":"10px 14px",
                            "fontSize":"0.95rem","color":"#113"})
        ]),

        # Row 1: GM trend / Crop stacked
        html.Div([
            html.Div([
                html.Div([
                    html.H4("Gross Margin Trend (% vs base)"),
                    dcc.Graph(id='gm-graph', config={"displayModeBar": False})
                ], style={"background":"#fff","borderRadius":"12px","boxShadow":"0 2px 10px rgba(10,57,103,.10)",
                          "padding":"16px","margin":"0 10px"})
            ], style={'width':'49%','display':'inline-block','verticalAlign':'top'}),

            html.Div([
                html.Div([
                    html.H4("Allocated Area by Crop (stacked)"),
                    dcc.Graph(id='crop-graph', config={"displayModeBar": False})
                ], style={"background":"#fff","borderRadius":"12px","boxShadow":"0 2px 10px rgba(10,57,103,.10)",
                          "padding":"16px","margin":"0 10px"})
            ], style={'width':'49%','display':'inline-block','verticalAlign':'top'}),
        ], style={'width':'100%','paddingTop':'6px'}),

        # Row 2: Water availability vs BAU baseline (100%) + Total water consumption
        html.Div([
            html.Div([
                html.H4("Water availability vs BAU baseline (100%)"),
                dcc.Graph(id='water-graph', config={"displayModeBar": False})
            ], style={"background":"#fff","borderRadius":"12px","boxShadow":"0 2px 10px rgba(10,57,103,.10)",
                      "padding":"16px","margin":"10px 24px"}),

            html.Div([
                html.H4("Total water consumption (m³)"),
                dcc.Graph(id='cons-graph', config={"displayModeBar": False})
            ], style={"background":"#fff","borderRadius":"12px","boxShadow":"0 2px 10px rgba(10,57,103,.10)",
                      "padding":"16px","margin":"10px 24px"})
        ])
    ]
)

# ========= CORE: compute total water consumption =========
def compute_total_water_timeseries(selected_scenario, selected_mun, selected_rcp, selected_method):
    """
    Calcola (per anno) il consumo idrico assoluto:
    sum_{municipio,crop} [ share(PMP Area) * Hectares(Municipio) * wr_m3_ha(Crop,Municipio) ].
    Media su AM (mean) + banda min/max.
    Se df_wr è vuoto (non rilevato), ritorna None.
    """
    if df_wr is None or df_wr.empty:
        return None

    # 1) filtro crop portfolio
    crp = df_crop[(df_crop["Scenario"]==selected_scenario) &
                  (df_crop["RCP"]==selected_rcp)].copy()
    if selected_method != "Tutti":
        crp = crp[crp["Method"]==selected_method]

    # per Doñana includo tutti i municipi; altrimenti 1 municipio
    if not is_donana(selected_mun):
        code = int(selected_mun)
        crp = crp[crp["Municipio"]==code].copy()

    if crp.empty:
        return None

    # 2) media su CM (e su Method se "Tutti") → manteniamo AM per banda min-max
    crp_mean = (crp.groupby(["Year","AM","Municipio","Crop"], as_index=False)["PMP Area"]
                   .mean())
    # 3) merge con ettari e water requirement
    crp_mean = crp_mean.merge(w_hect, on="Municipio", how="left")
    crp_mean = crp_mean.merge(df_wr, on=["Municipio","Crop"], how="left")

    crp_mean = crp_mean.dropna(subset=["Hectares","wr_m3_ha","PMP Area"])
    if crp_mean.empty:
        return None

    # 4) acqua assoluta = share * Ha * m3/ha
    crp_mean["water_m3"] = crp_mean["PMP Area"] * crp_mean["Hectares"] * crp_mean["wr_m3_ha"]

    # 5) aggregazione per anno & AM (somma su municipi e crop)
    agg = (crp_mean.groupby(["Year","AM"], as_index=False)["water_m3"].sum())

    # 6) riassunto su AM → mean / min / max per banda
    summary = agg.groupby("Year")["water_m3"].agg(["mean","min","max"]).reset_index()
    return summary

# ========= CALLBACK =========
@app.callback(
    Output("scenario-desc","children"),
    Output("gm-graph","figure"),
    Output("crop-graph","figure"),
    Output("water-graph","figure"),
    Output("cons-graph","figure"),
    Input("scenario","value"),
    Input("municipio","value"),
    Input("rcp","value"),
    Input("method","value")
)
def update_all(selected_scenario, selected_mun, selected_rcp, selected_method):
    # ---------- 1) Descrizione scenario ----------
    desc = SCENARIO_DESCR.get(selected_scenario, "Scenario non documentato.")
    desc_div = html.Div([html.B(selected_scenario.upper()+": "), html.Span(desc)])

    # ---------- 2) Filtri base ----------
    gmp = gm_merged[(gm_merged["Scenario"]==selected_scenario) &
                    (gm_merged["RCP"]==selected_rcp)].copy()
    crp = df_crop[(df_crop["Scenario"]==selected_scenario) &
                  (df_crop["RCP"]==selected_rcp)].copy()
    wdf = water_all[(water_all["Scenario"]==selected_scenario) &
                    (water_all["RCP"]==selected_rcp)].copy()
    if selected_method != "Tutti":
        gmp = gmp[gmp["Method"]==selected_method]
        crp = crp[crp["Method"]==selected_method]
        wdf = wdf[wdf["Method"]==selected_method]

    # ---------- 3) GM Trend (% vs base) — y-range [0,100] ----------
    if is_donana(selected_mun):
        tmp = gmp.dropna(subset=["GM_perc"]).copy()
        if tmp.empty:
            fig_gm = go.Figure().update_layout(title="No data", template="plotly_white",
                                               yaxis=dict(range=[0,100]))
        else:
            wm = (tmp.groupby(["Year","AM"])
                    .apply(lambda x: weighted_mean(x, "GM_perc", w_hect))
                    .reset_index(name="GM_perc_w")
                    .dropna(subset=["GM_perc_w"]))
            summary = wm.groupby("Year")["GM_perc_w"].agg(["mean","min","max"]).reset_index()
            fig_gm = go.Figure()
            fig_gm.add_trace(go.Scatter(x=summary["Year"], y=summary["mean"], mode="lines+markers",
                                        name="GM Trend", line=dict(width=3)))
            fig_gm.add_trace(go.Scatter(
                x=summary["Year"].tolist()+summary["Year"][::-1].tolist(),
                y=summary["max"].tolist()+summary["min"][::-1].tolist(),
                fill="toself", fillcolor="rgba(9,103,174,0.13)",
                line=dict(color="rgba(255,255,255,0)"), hoverinfo="skip",
                showlegend=True, name="Range (Min–Max)"
            ))
            fig_gm.add_hline(y=100, line_dash="dash", line_color="#8e8e8e",
                             annotation_text="Baseline", annotation_position="top left")
            fig_gm.update_layout(yaxis_title="Gross Margin [% vs base]", xaxis_title="Year",
                                 template="plotly_white", hovermode="x unified",
                                 legend=dict(bgcolor="rgba(255,255,255,0.88)"),
                                 yaxis=dict(range=[0,100]))
    else:
        code = int(selected_mun)
        tmp = gmp[gmp["Municipio"]==code].copy().sort_values(["AM","Year"])
        if tmp.empty:
            fig_gm = go.Figure().update_layout(title="No data", template="plotly_white",
                                               yaxis=dict(range=[0,100]))
        else:
            tmp["GM_perc_MA"] = tmp.groupby("AM")["GM_perc"].transform(lambda x: x.rolling(3, min_periods=1).mean())
            summary = tmp.groupby("Year")["GM_perc_MA"].agg(["mean","min","max"]).reset_index()
            fig_gm = go.Figure()
            fig_gm.add_trace(go.Scatter(x=summary["Year"], y=summary["mean"], mode="lines+markers",
                                        name="GM Trend", line=dict(width=3)))
            fig_gm.add_trace(go.Scatter(
                x=summary["Year"].tolist()+summary["Year"][::-1].tolist(),
                y=summary["max"].tolist()+summary["min"][::-1].tolist(),
                fill="toself", fillcolor="rgba(9,103,174,0.13)",
                line=dict(color="rgba(255,255,255,0)"), hoverinfo="skip",
                showlegend=True, name="Range (Min–Max)"
            ))
            fig_gm.add_hline(y=100, line_dash="dash", line_color="#8e8e8e",
                             annotation_text="Baseline", annotation_position="top left")
            fig_gm.update_layout(yaxis_title="Gross Margin [% vs base]", xaxis_title="Year",
                                 template="plotly_white", hovermode="x unified",
                                 legend=dict(bgcolor="rgba(255,255,255,0.88)"),
                                 title=dict(text=municipio_label(code, mun_code_to_name),
                                            y=0.98, x=0.02, xanchor='left', yanchor='top'),
                                 yaxis=dict(range=[0,100]))

    # ---------- 4) Crop stacked (share 0–1 → 0–100%) ----------
    if is_donana(selected_mun):
        tc = crp.dropna(subset=["PMP Area"]).copy()
        if tc.empty:
            fig_crop = go.Figure().update_layout(title="No crop data", template="plotly_white",
                                                 yaxis=dict(range=[0,1], tickformat=".0%"))
        else:
            tc = tc.merge(w_hect, on="Municipio", how="left").dropna(subset=["Hectares"])
            tc["w_area"] = tc["PMP Area"] * tc["Hectares"]
            wm = (tc.groupby(["Year","Crop"], as_index=False)
                    .agg(w_share=("w_area","sum"), H=("Hectares","sum")))
            wm["share"] = wm["w_share"] / wm["H"]
            pivot = wm.pivot(index="Year", columns="Crop", values="share").fillna(0)
            cols = [c for c in pivot.columns if c != "Secano"] + (["Secano"] if "Secano" in pivot.columns else [])
            pivot = pivot[cols]
            fig_crop = go.Figure()
            palette = px.colors.qualitative.Plotly + px.colors.qualitative.Pastel
            for i, crop in enumerate(pivot.columns):
                fig_crop.add_trace(go.Scatter(
                    x=pivot.index, y=pivot[crop], mode="lines", stackgroup="one", name=crop,
                    line=dict(width=0.8), opacity=0.98, fillcolor=palette[i % len(palette)]
                ))
            fig_crop.update_layout(yaxis_title="Allocated area (share)",
                                   xaxis_title="Year", template="plotly_white",
                                   hovermode="x unified", legend_title_text="Crop",
                                   yaxis=dict(range=[0,1], tickformat=".0%"))
    else:
        code = int(selected_mun)
        tc = crp[crp["Municipio"]==code].copy()
        if tc.empty:
            fig_crop = go.Figure().update_layout(title="No crop data", template="plotly_white",
                                                 yaxis=dict(range=[0,1], tickformat=".0%"))
        else:
            pivot = tc.pivot_table(index="Year", columns="Crop", values="PMP Area", aggfunc="mean").fillna(0)
            cols = [c for c in pivot.columns if c != "Secano"] + (["Secano"] if "Secano" in pivot.columns else [])
            pivot = pivot[cols]
            fig_crop = go.Figure()
            palette = px.colors.qualitative.Plotly + px.colors.qualitative.Pastel
            for i, crop in enumerate(pivot.columns):
                fig_crop.add_trace(go.Scatter(
                    x=pivot.index, y=pivot[crop], mode="lines", stackgroup="one", name=crop,
                    line=dict(width=0.8), opacity=0.98, fillcolor=palette[i % len(palette)]
                ))
            fig_crop.update_layout(yaxis_title="Allocated area (share)",
                                   xaxis_title="Year", template="plotly_white",
                                   hovermode="x unified", legend_title_text="Crop",
                                   title=dict(text=municipio_label(code, mun_code_to_name),
                                              y=0.98, x=0.02, xanchor='left', yanchor='top'),
                                   yaxis=dict(range=[0,1], tickformat=".0%"))

    # ---------- 5) Water availability vs BAU baseline — y-range [0,100] ----------
    if is_donana(selected_mun):
        tw = wdf.copy().merge(w_hect, on="Municipio", how="left").dropna(subset=["Hectares","Water_rel"])
        if tw.empty:
            fig_w = go.Figure().update_layout(title="No water data", template="plotly_white",
                                              yaxis=dict(range=[0,100]))
        else:
            tw["w_val"] = tw["Water_rel"] * tw["Hectares"]
            wm = (tw.groupby(["Year","AM"], as_index=False)
                    .agg(val=("w_val","sum"), H=("Hectares","sum")))
            wm["Water_rel_w"] = wm["val"] / wm["H"]  # in %
            summary = wm.groupby("Year")["Water_rel_w"].agg(["mean","min","max"]).reset_index()
            fig_w = go.Figure()
            fig_w.add_trace(go.Scatter(x=summary["Year"], y=summary["mean"], mode="lines+markers",
                                       name="Water availability (rel. BAU=100)", line=dict(width=3)))
            fig_w.add_trace(go.Scatter(
                x=summary["Year"].tolist()+summary["Year"][::-1].tolist(),
                y=summary["max"].tolist()+summary["min"][::-1].tolist(),
                fill="toself", fillcolor="rgba(9,103,174,0.13)",
                line=dict(color="rgba(255,255,255,0)"), hoverinfo="skip",
                showlegend=True, name="Range (Min–Max)"
            ))
            fig_w.add_hline(y=100, line_dash="dash", line_color="#8e8e8e",
                            annotation_text="BAU baseline", annotation_position="top left")
            fig_w.update_layout(yaxis_title="Water availability [% of BAU baseline]",
                                xaxis_title="Year", template="plotly_white", hovermode="x unified",
                                yaxis=dict(range=[0,100]))
    else:
        code = int(selected_mun)
        tw = wdf[wdf["Municipio"]==code].copy().sort_values(["AM","Year"])
        if tw.empty:
            fig_w = go.Figure().update_layout(title="No water data", template="plotly_white",
                                              yaxis=dict(range=[0,100]))
        else:
            summary = tw.groupby("Year")["Water_rel"].agg(["mean","min","max"]).reset_index()
            fig_w = go.Figure()
            fig_w.add_trace(go.Scatter(x=summary["Year"], y=summary["mean"], mode="lines+markers",
                                       name="Water availability (rel. BAU=100)", line=dict(width=3)))
            fig_w.add_trace(go.Scatter(
                x=summary["Year"].tolist()+summary["Year"][::-1].tolist(),
                y=summary["max"].tolist()+summary["min"][::-1].tolist(),
                fill="toself", fillcolor="rgba(9,103,174,0.13)",
                line=dict(color="rgba(255,255,255,0)"), hoverinfo="skip",
                showlegend=True, name="Range (Min–Max)"
            ))
            fig_w.add_hline(y=100, line_dash="dash", line_color="#8e8e8e",
                            annotation_text="BAU baseline", annotation_position="top left")
            fig_w.update_layout(yaxis_title="Water availability [% of BAU baseline]",
                                xaxis_title="Year", template="plotly_white", hovermode="x unified",
                                title=dict(text=municipio_label(code, mun_code_to_name),
                                           y=0.98, x=0.02, xanchor='left', yanchor='top'),
                                yaxis=dict(range=[0,100]))

    # ---------- 6) Total water consumption (m³) ----------
    cons = compute_total_water_timeseries(selected_scenario, selected_mun, selected_rcp, selected_method)
    if cons is None or cons.empty:
        fig_c = go.Figure().update_layout(title="Total water consumption not available (check crop/WR columns).",
                                          template="plotly_white")
    else:
        fig_c = go.Figure()
        fig_c.add_trace(go.Scatter(x=cons["Year"], y=cons["mean"], mode="lines+markers",
                                   name="Total water (mean)", line=dict(width=3)))
        fig_c.add_trace(go.Scatter(
            x=cons["Year"].tolist()+cons["Year"][::-1].tolist(),
            y=cons["max"].tolist()+cons["min"][::-1].tolist(),
            fill="toself", fillcolor="rgba(9,103,174,0.13)",
            line=dict(color="rgba(255,255,255,0)"), hoverinfo="skip",
            showlegend=True, name="Range (Min–Max)"
        ))
        fig_c.update_layout(yaxis_title="m³", xaxis_title="Year",
                            template="plotly_white", hovermode="x unified")

    return desc_div, fig_gm, fig_crop, fig_w, fig_c

# ========= MAIN =========
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False)