tabs/tab4_optimisation.py

from collections import defaultdict
import datetime
import io
import json
import unicodedata
import zipfile
import openpyxl
import pandas as pd
import streamlit as st
import plotly.express as px
from config.config import CRRA_COLORS, PLOTLY_COLOR_SEQUENCE, PLOTLY_HEATMAP_SCALE, resource_order, resource_unit_map, resource_groups, secondary_resources, CO2_UNIT, method_order, section_intros
from optimisation.optimisation import run_optimisation
from utils.ui_helpers import apply_chart_style, register_chart_data, show_bar
from utils.data_utils import (is_valid_scenario, is_complete_scenario, load_inputs_from_excel,
                              neutralize_zero_coefficient_methods, strip_unit_suffix,
                              get_resource_to_group, get_resource_to_unit,
                              is_other_method, get_other_method_variants, make_variant_name,
                              default_constraint as _dc, to_json_safe, build_results_df)
from utils.state_utils import clear_widget_keys
from config.config import hidden_resource_names as _hidden

def panel_run_from_tabs():
    """Left sub-panel: run optimisation from manually filled tabs."""
    st.markdown("**Option 1 - From current tab values**", unsafe_allow_html=True)
    st.markdown('<p class="caption">Set resource availability and method constraints in the tabs above, then run.</p>', unsafe_allow_html=True)
    st.divider() 


def panel_load_from_excel():
    """Left sub-panel: load inputs from an Excel file, then run."""
    col_markdown, col_upload = st.columns([1, 1], vertical_alignment="center")

    with col_markdown:
        st.markdown("**Option 2 - From an Excel file**")

    st.markdown('<p class="caption">Upload an Excel file previously exported from the tool to reload all inputs at once.</p>', unsafe_allow_html=True)

    if "excel_loaded_file" in st.session_state:
        st.success(f"Loaded: **{st.session_state['excel_loaded_file']}** - All tabs updated.", icon="✅")
    for w in st.session_state.pop("import_warnings", []):
        st.warning(w, icon=":material/search:")

    uploaded = st.file_uploader("Upload inputs (.xlsx)", type=["xlsx"], key=f"inputs_loader_{st.session_state.excel_uploader_key}")

    if uploaded is None:
        return

    filename_stem = uploaded.name.rsplit(".", 1)[0]

    if st.button("↩️ Load inputs into all tabs", width='stretch'):
        st.session_state.setdefault("input_file_name_suggestion", filename_stem)
        try:
            resource_caps, constraints, costs, custom_resources, enabled_methods, import_warnings = load_inputs_from_excel(uploaded)
            st.session_state.resource_caps = resource_caps
            st.session_state.constraints = constraints
            st.session_state.costs = costs
            st.session_state.custom_resources = custom_resources
            st.session_state.enabled_methods = enabled_methods
            st.session_state.auto_disabled_methods = set()
            clear_widget_keys()
            for _m, _c in constraints.items():
                st.session_state[f"toggle_{_m}"] = bool(_c.get("active", True))
                if _c.get("active", True):
                    st.session_state[f"cap_type_{_m}"] = _c.get("cap_type", "percent")
                    if _c.get("cap_type", "percent") == "percent":
                        st.session_state[f"slider_{_m}"] = int(_c.get("cap_value", 100))
                    else:
                        st.session_state[f"manual_{_m}"] = float(_c.get("cap_value", 0.0))
            # Signal tab2 to force-refresh all widget keys on next visit.
            st.session_state["constraints_just_loaded"] = True
            st.session_state.pop("latest_result", None)
            st.session_state.pop("scenario_loaded_file", None)
            st.session_state.pop("scenario_name_suggestion", None)
            st.session_state.pop("resource_last_loaded_file", None)
            st.session_state["excel_loaded_file"] = uploaded.name
            st.session_state["input_file_name"] = filename_stem
            st.session_state["scenario_uploader_key"] += 1 # force reload in scenario uploader
            st.session_state["import_warnings"] = import_warnings
            st.rerun()
        except Exception as e:
            st.error(f"Error loading file: {e}")

def parse_zip_scenario(uploaded) -> dict | None:
    """Reconstruct a scenario dict from a ZIP containing the exported CSVs.

    Required: resource_caps.csv, method_constraints.csv, allocations.csv, method_costs.csv.
    Optional: custom_resources.csv, enabled_methods.csv.

    The detailed outputs (resource_usage, resource_used_total, resource_remaining) 
    are NOT trusted from allocations.csv alone — they are
    recomputed by re-running the optimiser on the reloaded inputs. This guarantees
    the reloaded scenario is always internally consistent (and lets results pages
    render without missing keys), instead of relying on whatever was exported.
    """

    try:
        with zipfile.ZipFile(io.BytesIO(uploaded.read())) as zf:
            names = zf.namelist()
            
            def read(fname):
                return pd.read_csv(io.BytesIO(zf.read(fname)), index_col=0)
 
            required = {"resource_caps.csv", "method_constraints.csv",
                        "allocations.csv", "method_costs.csv"}
            if not required.issubset(names):
                return None
 
            resource_caps_raw = read("resource_caps.csv")["Available Amount"].to_dict()
            resource_caps = {strip_unit_suffix(k): v for k, v in resource_caps_raw.items()}

            df_c = read("method_constraints.csv")

            # Compatibility with older files that used `max_share`
            if "cap_type (percent or absolute)" in df_c.columns:
                _VALID_CAP_TYPES = {"percent", "absolute"}
                constraints = {}
                for m, row in df_c.iterrows():
                    raw_ct = str(row["cap_type (percent or absolute)"])
                    raw_active = row["active"]
                    constraints[m] = {
                        "active": raw_active is True or str(raw_active).strip().lower() == "true",
                        # Store the raw value even if invalid so the UI can display
                        # a warning and let the user correct it without re-importing.
                        "cap_type": raw_ct if raw_ct in _VALID_CAP_TYPES else raw_ct,
                        "cap_value": float(row["cap_value"]),
                        "_cap_type_invalid": raw_ct not in _VALID_CAP_TYPES,
                    }
            else:
                constraints = {
                    m: {
                        "active": row["active"] is True or str(row["active"]).strip().lower() == "true",
                        "cap_type": "percent",
                        "cap_value": float(row["max_share"])
                    }
                    for m, row in df_c.iterrows()
                }

            df_costs = read("method_costs.csv")  # index=method, columns=resources

            costs = {m: row.dropna().to_dict() for m, row in df_costs.iterrows()}

            df = read("allocations.csv")
            col = f"{CO2_UNIT} removed"
            if col not in df.columns:
                old_col = "tCO₂ removed"
                if old_col in df.columns:
                    df[col] = df[old_col] / 1_000_000  # convert tCO2 → MtCO2
                else:
                    raise ValueError(f"Column '{col}' not found. Available: {list(df.columns)}")

            allocations = df[col].to_dict()

            total = sum(allocations.values())

            # Load optional custom resource batches
            custom_resources = []
            if "custom_resources.csv" in names:
                df_custom = pd.read_csv(io.BytesIO(zf.read("custom_resources.csv")))
                required_custom = {"name", "group", "amount", "unit", "method", "min", "median", "max"}
                if required_custom.issubset(df_custom.columns):
                    for bname in df_custom["name"].unique():
                        df_b = df_custom[df_custom["name"] == bname]
                        methods = {
                            row["method"]: {
                                "min": float(row["min"]),
                                "median": float(row["median"]),
                                "max": float(row["max"]),
                            }
                            for _, row in df_b.iterrows()
                        }
                        custom_resources.append({
                            "name": bname,
                            "group": str(df_b["group"].iloc[0]),
                            "amount": float(df_b["amount"].iloc[0]),
                            "unit": str(df_b["unit"].iloc[0]) if "unit" in df_b.columns else "?",
                            "methods": methods,
                        })

            # Load optional per-resource manual method allocation
            enabled_methods = {}
            if "enabled_methods.csv" in names:
                df_enabled = pd.read_csv(io.BytesIO(zf.read("enabled_methods.csv")))
                required_enabled = {"resource", "method", "enabled"}
                if required_enabled.issubset(df_enabled.columns):
                    for _, row in df_enabled.iterrows():
                        enabled_methods.setdefault(str(row["resource"]), {})[str(row["method"])] = bool(row["enabled"])

            # Use the loaded costs/constraints, not the old session state.
            # Old state may contain variant entries (e.g. "Method: Batch") from a
            # previous scenario; using them here causes false deactivation.
            # Guard: only report standard methods so variants never appear.
            _std = set(method_order)
            safe_constraints = neutralize_zero_coefficient_methods(costs, constraints)
            neutralized = [
                m for m in safe_constraints
                if safe_constraints[m].get("active") is False
                and constraints.get(m, {}).get("active") is True
                and m in _std
            ]
            
            # Recompute detailed outputs from the reloaded inputs so the scenario
            # is internally consistent (rather than trusting allocations.csv alone).
            success, result = run_optimisation(
                resource_caps=resource_caps,
                method_constraints=constraints,
                method_costs=costs,
                custom_resources=custom_resources,
                enabled_methods=enabled_methods,
                resource_to_group=get_resource_to_group(),
            )

            if neutralized:
                st.info(
                    f"{len(neutralized)} method(s) automatically deactivated (no active resource): "
                    + ", ".join(neutralized),
                    icon=":material/info:",
                )

            if success:
                allocations = result["method_usage"]
                total = result["total_removed"]
                resource_usage = result["resource_usage"]
                resource_used_total = result["resource_used_total"]
                resource_remaining = result["resource_remaining"]
                resource_actual_gain = result.get("resource_actual_gain", {})
            else:
                # Fall back to the raw exported allocations if re-optimisation
                # fails (e.g. incompatible inputs) — better than returning None.
                allocations = df[col].to_dict()
                total = sum(allocations.values())
                resource_usage = {}
                resource_used_total, resource_remaining = {}, {}
                resource_actual_gain = {}

            return {
                "summary": {"total_co2_removed": total, "total_cost": 0},
                "allocations": allocations,
                "resource_usage": resource_usage,
                "resource_used_total": resource_used_total,
                "resource_remaining": resource_remaining,
                "resource_actual_gain": resource_actual_gain,
                "resource_caps": resource_caps,
                "method_constraints": constraints,
                "method_costs": costs,
                "custom_resources": custom_resources,
                "enabled_methods": enabled_methods,
            }
        
    except Exception:
        return None

def panel_load_from_json_or_csv():
    """Right column: reload a full scenario from JSON or ZIP."""
    st.markdown('<p class="caption">Accepts `.json` or `.zip` (CSV export)</p>', unsafe_allow_html=True)

    if "scenario_loaded_file" in st.session_state:
        st.success(
            f"Scenario **{st.session_state['scenario_loaded_file']}** loaded :\n\nAll tabs updated - optimisation results displayed.",
            icon="✅"
        )
    uploaded_json = st.file_uploader(
        "Upload scenario (.json or .zip)", type=["json", "zip"], key=f"scenario_loader_{st.session_state.scenario_uploader_key}"
    )

    if uploaded_json is None:
        return
 
    if uploaded_json.name.endswith(".json"):
        try:
            data = json.load(uploaded_json)
        except Exception as e:
            st.error(f"Cannot read JSON: {e}")
            return
        if not is_valid_scenario(data):
            st.error("Invalid scenario format.")
            return
        if not is_complete_scenario(data):
            success, result = run_optimisation(
                resource_caps=data["resource_caps"],
                method_constraints=data["method_constraints"],
                method_costs=data["method_costs"],
                custom_resources=data.get("custom_resources", []),
                enabled_methods=data.get("enabled_methods", {}),
                resource_to_group=get_resource_to_group(),
            )
            if success:
                data["resource_usage"] = result["resource_usage"]
                data["resource_used_total"] = result["resource_used_total"]
                data["resource_remaining"] = result["resource_remaining"]
                data["resource_actual_gain"] = result.get("resource_actual_gain", {})
                data.setdefault("enabled_methods", {})
            else:
                st.warning("Scenario loaded but detailed outputs could not be recomputed.", icon=":material/warning:")
    else:
        data = parse_zip_scenario(uploaded_json)
        if data is None:
            st.error(
                "Invalid ZIP. Expected: resource_caps.csv, method_constraints.csv, "
                "method_costs.csv, allocations.csv, "
            )
            return
    
    filename_stem = uploaded_json.name.rsplit(".", 1)[0]
 
    if st.button("↩️ Reload scenario and run optimisation", width='stretch'):
        st.session_state.setdefault("scenario_name_suggestion", filename_stem)
        st.session_state.pop("latest_result", None)
        st.session_state.pop("excel_loaded_file", None)
        st.session_state.pop("input_file_name_suggestion", None)
        st.session_state.pop("resource_last_loaded_file", None)
        st.session_state.resource_caps = data["resource_caps"].copy()
        st.session_state.constraints = data["method_constraints"].copy()
        # Use pre-expansion costs (has originals, no variants) so the next run
        # can expand "Other" methods correctly without skipping the resource check.
        # Fall back to method_costs for scenarios saved before this fix.
        st.session_state.costs = data.get("method_costs_base", data["method_costs"]).copy()
        # Restore custom resource batches if present in the scenario
        st.session_state.custom_resources = data.get("custom_resources", [])
        st.session_state.enabled_methods = data.get("enabled_methods", {})
        # Reset auto_disabled_methods so refresh_active_methods (app.py's shared
        # init) starts from a clean slate matching the loaded scenario (avoids
        # stale entries that would prevent or incorrectly trigger re-enable on
        # the next render).
        st.session_state.auto_disabled_methods = set()
        clear_widget_keys()
        for _m, _c in st.session_state.constraints.items():
            st.session_state[f"toggle_{_m}"] = bool(_c.get("active", True))
            if _c.get("active", True):
                st.session_state[f"cap_type_{_m}"] = _c.get("cap_type", "percent")
                if _c.get("cap_type", "percent") == "percent":
                    st.session_state[f"slider_{_m}"] = int(_c.get("cap_value", 100))
                else:
                    st.session_state[f"manual_{_m}"] = float(_c.get("cap_value", 0.0))
        st.session_state["constraints_just_loaded"] = True
        st.session_state["scenario_loaded_file"] = uploaded_json.name
        st.session_state["scenario_name"] = filename_stem
        st.session_state["latest_result"] = data
        st.session_state["excel_uploader_key"] += 1 # force reload in excel uploader
        if st.session_state.custom_resources:
            n = len(st.session_state.custom_resources)
            st.info(f"{n} new resource(s) restored from scenario.", icon=":material/info:")
        st.rerun()

def run_and_store_result():
    """Run the LP optimisation from current session-state inputs and store the result.

    Reads resource_caps, constraints, costs, custom_resources, and enabled_methods
    from session state. Applies enabled_methods overrides to a temporary cost copy,
    neutralises zero-coefficient methods, runs run_optimisation, then writes the
    full result dict to st.session_state.latest_result.

    Shows a Streamlit warning/error if inputs are invalid or optimisation fails.

    Returns:
        None
    """
    custom_resources = st.session_state.get("custom_resources", [])
    enabled_methods = st.session_state.get("enabled_methods", {}) # for custom standard resources

    # Check that at least some resource is non-zero (standard OR custom)
    std_total = sum(st.session_state.resource_caps.values())

    custom_total = sum(float(b["amount"]) for b in custom_resources)
    if std_total <= 0 and custom_total <= 0:
        st.warning("⚠️ All resources are set to 0. Please set resource availability or load a scenario / input data.", icon=":material/warning:")
        return
    
    # Apply enabled_methods overrides to a temporary copy — session state is NOT
    # mutated, so re-enabling a resource in tab1 restores the original coefficient.
    # The LP receives zeroed coefficients only for this run.
    applied_costs = {m: dict(v) for m, v in st.session_state.costs.items()}
    for method, method_resources in applied_costs.items():
        for resource in list(method_resources):
            if not enabled_methods.get(resource, {}).get(method, True):
                applied_costs[method][resource] = 0.0

    # Inject custom batch coefficients into applied_costs so neutralize does
    # not wrongly deactivate a method that only has custom-batch access.
    # Custom pools are INDEPENDENT of the standard resource enable/disable toggle.
    for batch in custom_resources:
        batch_name = batch["name"]
        for method_name, coefficients in batch.get("methods", {}).items():
            if method_name in applied_costs and applied_costs[method_name].get(batch_name, 0.0) == 0.0:
                applied_costs[method_name][batch_name] = coefficients.get("median", 0.0)

    # Only neutralize standard methods (those present in applied_costs).
    # Variant names from previous runs may exist in st.session_state.constraints
    # but are not in applied_costs yet — passing them would cause false positives
    # (empty cost dict → all_zero → auto-deactivated before expansion runs).
    standard_constraints = {
        m: c for m, c in st.session_state.constraints.items()
        if m in applied_costs
    }
    safe_constraints = neutralize_zero_coefficient_methods(applied_costs, standard_constraints)

    # Deactivate any method that requires a resource not provided in tab1.
    # Hidden resources (e.g. "Other water") are never in resource_caps because
    # they are not shown in tab1.  Without this check, methods that consume such
    # resources go unconstrained in the LP and produce misleading results.
    # A resource is "missing" when it is absent from resource_caps (not set by
    # the user) AND has no custom batch covering it.
    custom_batch_names_pre = {b["name"] for b in custom_resources}
    available_resources = set(st.session_state.resource_caps) | custom_batch_names_pre
    for m, costs in applied_costs.items():
        if not safe_constraints.get(m, {}).get("active", True):
            continue
        missing = [
            r for r, c in costs.items()
            if c > 0 and r not in available_resources
        ]
        if missing:
            safe_constraints.setdefault(m, {"active": True, "cap_type": "percent", "cap_value": 100})
            safe_constraints[m] = {**safe_constraints[m], "active": False}

    # Expand "Other" methods: replace each with one variant per assigned batch.
    # Each variant uses the batch as its primary resource (standard LP constraint)
    # instead of the z-variable substitution mechanism.
    #
    # Snapshot taken before expansion: originals are present, variants absent.
    # Stored in the result so scenario reload can restore the original method costs
    # (post-expansion applied_costs deletes originals and adds variants, which would
    # cause the variants to bypass the missing-resources check on the next load).
    method_costs_base = {m: dict(v) for m, v in applied_costs.items()}
    resource_caps_for_lp = dict(st.session_state.resource_caps)
    custom_resources_for_lp = []
    expanded_batch_names: set[str] = set()

    for m in list(applied_costs.keys()):
        if not is_other_method(m):
            continue
        batches_for_m = get_other_method_variants(m, custom_resources)
        if not batches_for_m:
            continue
        # All batch names for this "Other" method — must be excluded from each
        # variant's base costs so variants don't inherit each other's batches.
        all_batch_names_for_m = {b["name"] for b in batches_for_m}

        for batch in batches_for_m:
            vname = make_variant_name(m, batch["name"])
            # Inherit only non-hidden, non-batch secondary resources from the
            # original method, then add THIS batch as the sole primary resource.
            variant_costs = {
                r: coeff for r, coeff in applied_costs[m].items()
                if r not in _hidden and r not in all_batch_names_for_m
            }
            variant_costs[batch["name"]] = applied_costs[m].get(
                batch["name"],
                batch["methods"][m]["median"],
            )
            applied_costs[vname] = variant_costs
            # Variant constraint: prefer an explicit user-set value, else default.
            safe_constraints[vname] = (
                st.session_state.constraints.get(vname) or _dc()
            )
            # Batch becomes a standard capped resource for the LP.
            resource_caps_for_lp[batch["name"]] = float(batch["amount"])
            expanded_batch_names.add(batch["name"])
        # Remove the original "Other" method — replaced by variants.
        del applied_costs[m]
        del safe_constraints[m]

    # Compute after expansion so "Other" methods (deleted above) don't appear.
    # Guard: only include standard methods so dynamically created variants (not in
    # method_order) never show up in this message even in edge cases.
    _standard_methods = set(method_order)
    neutralized = [
        m for m in safe_constraints
        if safe_constraints[m].get("active") is False
        and st.session_state.constraints.get(m, {}).get("active") is True
        and m in _standard_methods
    ]

    # LP receives only non-expanded batches.  Expanded batches are now standard
    # resources (in resource_caps_for_lp) so they get a type-A constraint — if
    # they stayed in custom_resources they'd land in custom_batch_names, be
    # excluded from standard_resources, and generate no constraint → unbounded LP.
    custom_resources_for_lp = [b for b in custom_resources if b["name"] not in expanded_batch_names]
    # Display functions keep the full list so group/unit info is available.
    custom_resources_for_display = list(custom_resources)

    # Non-"Other" methods that shared an expanded batch (e.g. "Mineral OAE" also
    # used "Limestone") lost their z-variable when the batch was removed from
    # custom_resources_for_lp.  Inject the batch as a standard resource coefficient
    # into their applied_costs so the type-A constraint still applies to them.
    #
    # Also zero out the standard group resource coefficient when the standard pool
    # is empty (cap = 0).  In the z-variable path the B-constraint allows
    # "y_m - z_{m,n} ≤ 0", letting the batch fully cover the method.  In the
    # standard-resource path that substitute mechanism is gone, so a method with
    # cost(group) > 0 and cap(group) = 0 would be incorrectly forced to y = 0
    # by the type-A constraint before it can use the batch at all.
    for batch in custom_resources:
        if batch["name"] not in expanded_batch_names:
            continue
        batch_group = batch.get("group", "")
        group_cap_is_zero = resource_caps_for_lp.get(batch_group, 0.0) == 0.0
        for method_name, coeff_dict in batch.get("methods", {}).items():
            if method_name not in applied_costs:
                continue  # "Other" method was already replaced by variants above
            applied_costs[method_name][batch["name"]] = coeff_dict.get("median", 0.0)
            # If the standard pool for this group is empty, zero its coefficient so
            # the type-A constraint doesn't force the method's output to zero.
            if group_cap_is_zero and batch_group in applied_costs[method_name]:
                applied_costs[method_name][batch_group] = 0.0

    success, result = run_optimisation(
        resource_caps=resource_caps_for_lp,
        method_constraints=safe_constraints,
        method_costs=applied_costs,
        custom_resources=custom_resources_for_lp,
        enabled_methods=enabled_methods,
        resource_to_group=get_resource_to_group(),
    )

    if not success:
        st.error("❌ optimisation failed")
        if "message" in result:
            st.error(result["message"])
        return
    
    if neutralized:
        st.info(
            f"{len(neutralized)} method(s) automatically deactivated (no active resource): "
            + ", ".join(neutralized),
            icon=":material/info:",
        )

    allocations = result["method_usage"]
    total_removed = result["total_removed"]

    if total_removed <= 0:
        st.warning(f"❌ optimisation removed 0 {CO2_UNIT}. Relax constraints or increase resources.", icon=":material/warning:")
        return

    total_cost = sum(
        allocations[m] * sum(applied_costs.get(m, st.session_state.costs.get(m, {})).values())
        for m in allocations
    )
    st.session_state.latest_result = {
        "summary": {"total_co2_removed": total_removed, "total_cost": total_cost},
        "allocations": allocations,
        "resource_usage": result["resource_usage"],
        "resource_used_total": result["resource_used_total"],
        "resource_remaining": result["resource_remaining"],
        "resource_actual_gain": result.get("resource_actual_gain", {}),
        "resource_caps": resource_caps_for_lp,
        "method_constraints": safe_constraints,
        "method_costs": applied_costs,
        "method_costs_base": method_costs_base,
        "custom_resources": custom_resources_for_display,
        "enabled_methods": enabled_methods,
    }
    if custom_resources:
        st.info(f"{len(custom_resources)} new resource(s) included in optimisation.", icon=":material/info:")

    # Warn when an active method has a primary resource whose coefficient is so
    # small (< 1e-7) that it is effectively unconstrained — the LP will run the
    # method freely regardless of how much of that resource is available.
    _COEFF_WARN = 1e-7
    near_zero_warnings = []
    for m, qty in allocations.items():
        if qty <= 0:
            continue
        for r, coeff in applied_costs.get(m, {}).items():
            if (m, r) in secondary_resources:
                continue
            if r not in st.session_state.resource_caps:
                continue
            if 0 < coeff < _COEFF_WARN:
                near_zero_warnings.append((m, r, coeff))
    if near_zero_warnings:
        lines = "\n".join(
            f"- **{m}** \n {r}: coefficient = {coeff:.2e} (effectively 0 — resource not constraining)"
            for m, r, coeff in near_zero_warnings
        )
        st.warning(
            f"The following primary resources have near-zero coefficients and do **not** "
            f"constrain the method: the method can run regardless of their availability:\n\n{lines}",
            icon=":material/warning:",
        )

    st.success(f"✅ {total_removed:,.4f} {CO2_UNIT} removed")


def show_metrics(latest):
    """Render top-level KPI metrics and an optional custom-resource summary expander.

    Args:
        latest (dict): result dict from run_optimisation with keys summary, allocations,
                       custom_resources.

    Returns:
        None
    """
    custom_resources = latest.get("custom_resources", [])
    cols = st.columns(3) if custom_resources else st.columns(2)
    cols[0].metric("Total CO₂ Removed", f"{latest['summary']['total_co2_removed']:,.2f} {CO2_UNIT}")
    cols[1].metric("Methods Used", len(latest["allocations"]))

    resource_to_group = get_resource_to_group()
    # Show which custom batches were included and how much of each was available
    if custom_resources:
        with st.expander(f"✅ {len(custom_resources)} new resource(s) included in this optimisation", expanded=False):
            rows = []
            for batch in custom_resources:
                group = resource_to_group.get(batch["group"])
                unit = resource_unit_map.get(group, '?')
                rows.append({
                    "Batch name": batch["name"],
                    "Group": batch["group"],
                    "Amount available": f"{float(batch['amount']):,.2f} {unit}",
                    "Methods linked": ", ".join(batch["methods"].keys())
                })
            st.dataframe(pd.DataFrame(rows), width='stretch', hide_index=True)

def show_donut(df_result):
    """Render a donut chart of CO₂ removal by method and register its data.

    Args:
        df_result (pd.DataFrame): output of build_results_df, indexed by method.

    Returns:
        None
    """
    df_plot = df_result.reset_index().rename(columns={"index": "Method"})
    fig = px.pie(df_plot, names="Method", values=f"{CO2_UNIT} removed", hole=0.4,
                 color_discrete_sequence=PLOTLY_COLOR_SEQUENCE)
    fig.update_traces(textinfo="label+percent", textposition="outside",
                   marker=dict(line=dict(color=CRRA_COLORS["white"], width=1)))
    apply_chart_style(fig)
    st.plotly_chart(fig, width='stretch')
    register_chart_data("Method allocation (donut)", df_plot)

def show_heatmap(df_result, resource_caps, method_costs, custom_resources=None, resource_usage=None):
    """Render a resource-usage heatmap (% of cap per method × resource) and register data.

    Args:
        df_result (pd.DataFrame): output of build_results_df, indexed by method.
        resource_caps (dict): {resource: available amount}.
        method_costs (dict): {method: {resource: coefficient}}.
        custom_resources (list | None): list of custom batch dicts.
        resource_usage (dict | None): {resource: {method: qty}} from LP output;
                                      falls back to cost × allocation if None.

    Returns:
        None
    """
    # Normalize standard method_costs keys (strip legacy unit suffixes)
    normalized_costs = {
        m: {strip_unit_suffix(r): v for r, v in resources.items()}
        for m, resources in method_costs.items()
    }

    # Build complete caps — resource_caps may not contain custom batches
    full_caps = dict(resource_caps)
    for batch in (custom_resources or []):
        full_caps.setdefault(batch["name"], float(batch["amount"]))

    standard_available = [r for r in resource_order if full_caps.get(r, 0) > 0]
    extra_available = [b["name"] for b in (custom_resources or []) if float(b["amount"]) > 0]
    available = standard_available + extra_available

    if not available:
        st.info("No resources with non-zero availability to display.", icon=":material/info:")
        return

    matrix_pct, matrix_exact, y_labels = [], [], []
    for method in df_result.index:
        row_pct, row_exact = [], []
        for r in available:
            cap = full_caps.get(r, 0)
            if cap <= 0:
                row_pct.append(0)
                row_exact.append(0)
                continue
            if resource_usage is not None:
                used = resource_usage.get(r, {}).get(method, 0.0)
            else:
                used = normalized_costs.get(method, {}).get(r, 0.0) * df_result.at[method, f"{CO2_UNIT} removed"]
            row_pct.append(min(used / cap * 100, 100.0))
            row_exact.append(used)
        matrix_pct.append(row_pct)
        matrix_exact.append(row_exact)
        y_labels.append(method)

    df_heat = pd.DataFrame(matrix_pct, index=y_labels, columns=available)
    df_exact = pd.DataFrame(matrix_exact, index=y_labels, columns=available)

    # A genuine but tiny usage (e.g. 0.03%) rounds to "0.0" at 1 decimal,
    # visually indistinguishable from a true zero. Below 0.05 show the exact
    # value instead (enough decimals to reveal it, trailing zeros stripped)
    # so a real (if small) contribution never looks like "not used at all".
    def _fmt_pct(v):
        if v <= 0:
            return "0"
        if v < 0.1:
            s = f"{v:.4f}".rstrip("0").rstrip(".")
            return s if s else "0"
        return f"{v:.1f}"

    text_matrix = [[_fmt_pct(v) for v in row] for row in matrix_pct]

    fig = px.imshow(
        df_heat, color_continuous_scale=PLOTLY_HEATMAP_SCALE,
        labels=dict(x="Resource", y="Method", color="% Used"),
    )
    matrix_exact_fmt = [[f"{v:.10g}" for v in row] for row in matrix_exact]
    pct_and_exact = [
        [[text_matrix[i][j], matrix_exact_fmt[i][j]] for j in range(len(available))]
        for i in range(len(y_labels))
    ]
    fig.update_traces(
        text=text_matrix, texttemplate="%{text}",
        customdata=pct_and_exact,
        hovertemplate="<b>%{y}</b><br>%{x}<br>% of cap: %{customdata[0]}%<br>Exact: %{customdata[1]}<extra></extra>",
    )
    n_rows = len(df_heat.index)
    n_cols = len(df_heat.columns)
    height = max(500, 80 + 50 * n_rows)
    # For wide heatmaps, force a minimum width so each column stays readable
    if n_cols > 8:
        fig.update_layout(width=max(900, 90 * n_cols))
    apply_chart_style(fig, height=height)
    st.plotly_chart(fig, width='stretch')

    resource_to_group = get_resource_to_group()
    unit_map = {r: resource_unit_map.get(resource_to_group.get(r), "?") for r in available}
    for batch in (custom_resources or []):
        unit_map[batch["name"]] = resource_unit_map.get(
            resource_to_group.get(batch["group"]), batch.get("unit", "?")
        )
    col_rename = {r: f"{r} ({unit_map.get(r, '?')})" for r in available}

    df_export_pct = df_heat.reset_index().rename(columns={"index": "Method"})
    df_export_exact = (
        df_exact.reset_index()
        .rename(columns={"index": "Method"})
        .rename(columns=col_rename)
    )

    with st.expander("📊 View / download chart data"):
        df_pct_display = df_export_pct.copy()
        for col in df_pct_display.columns:
            if col != "Method":
                df_pct_display[col] = df_pct_display[col].map(lambda x: f"{_fmt_pct(x)}%")
        st.dataframe(df_pct_display, width='stretch', hide_index=True)

    register_chart_data("Resource usage heatmap (% of cap)", df_export_pct)
    register_chart_data("Resource usage exact values", df_export_exact)

def show_absolute_usage(allocations, method_costs, custom_resources=None, resource_usage=None):
    """Render grouped bar charts of absolute resource consumption by method and unit.

    Args:
        allocations (dict): {method: CO₂ removed (float)}.
        method_costs (dict): {method: {resource: coefficient}}.
        custom_resources (list | None): list of custom batch dicts.
        resource_usage (dict | None): {resource: {method: qty}} from LP output;
                                      falls back to cost × allocation if None.

    Returns:
        None
    """
    # Normalize method_costs keys: strip legacy unit suffixes
    normalized_costs = {
        m: {strip_unit_suffix(r): v for r, v in resources.items()}
        for m, resources in method_costs.items()
    }
    # Build resource_units from standard groups
    resource_units = defaultdict(list)
    resource_to_group = get_resource_to_group()
    for cat, resources in resource_groups.items():
        resource_units[resource_unit_map[cat]].extend(resources)

    # Inject custom resource batches using the correct unit from their reference group
    custom_batch_names = set()
    for batch in (custom_resources or []):
        custom_batch_names.add(batch["name"])
        group_name = resource_to_group.get(batch["group"])
        correct_unit = resource_unit_map.get(group_name, "?")
        if batch["name"] not in resource_units[correct_unit]:
            resource_units[correct_unit].append(batch["name"])

    local_resource_units = dict(resource_units)

    rows = []
    for m, co2 in allocations.items():
        for unit, resources in local_resource_units.items():
            for r in resources:
                if r in custom_batch_names:
                    # Use actual LP usage — coefficient × total_y_m overcounts when
                    # standard and custom pools are split across z variables
                    used_qty = (resource_usage or {}).get(r, {}).get(m, 0.0)
                    if used_qty > 1e-9:
                        rows.append({"Method": m, "Resource": r, "Used": used_qty, "Unit": unit})
                else:
                    if resource_usage is not None:
                        used_qty = resource_usage.get(r, {}).get(m, 0.0)
                        if used_qty > 1e-9:
                            rows.append({"Method": m, "Resource": r, "Used": used_qty, "Unit": unit})
                    else:
                        cost = normalized_costs.get(m, {}).get(r)
                        if cost is not None and cost * co2 > 0:
                            rows.append({"Method": m, "Resource": r, "Used": cost * co2, "Unit": unit})
    if not rows:
        return

    df = pd.DataFrame(rows)
    for unit in df["Unit"].unique():
        df_unit = df[df["Unit"] == unit]
        fig = px.bar(
            df_unit, x="Resource", y="Used", color="Method", barmode="group",
            color_discrete_sequence=PLOTLY_COLOR_SEQUENCE,
        )
        apply_chart_style(fig)
        fig.update_layout(yaxis_title=f"Used ({unit})", xaxis_tickangle=-45)
        st.plotly_chart(fig, width='stretch')
    with st.expander("📊 View / download chart data"):
        st.dataframe(df, width='stretch', hide_index=True)
    register_chart_data("Absolute resource usage by method", df)

def show_resource_balance(latest):
    """Render a stacked bar chart comparing quantity used vs remaining for each resource.

    Args:
        latest (dict): result dict with keys resource_caps, resource_used_total,
                       resource_remaining, custom_resources.

    Returns:
        None
    """
    st.markdown('<p class="crra-sub-heading">Resource balance: used vs available</p>', unsafe_allow_html=True)
    st.markdown(
        '<p class="caption">Green bars show quantities consumed by the optimised portfolio; grey bars show the unused remainder. '
        '<br>Resources with almost no grey bar are fully utilised and may be limiting further CDR. '
        '<br>Resources with a large grey bar have spare capacity: either no method needs them or their demand is already saturated elsewhere.</p>',
        unsafe_allow_html=True,
    )
    resource_caps = latest["resource_caps"]
    used_total = latest["resource_used_total"]

    resource_to_group = get_resource_to_group()
    # Custom batch names: their amounts come from resource_caps (expanded batches
    # are also in resource_caps) but we display them via the custom_resources loop
    # to retain group/unit info — skip them in the standard loop to avoid duplicates.
    custom_names = {b["name"] for b in latest.get("custom_resources", [])}

    rows = []
    for r in resource_caps:
        if r in custom_names:
            continue  # handled in the custom_resources loop below
        group = resource_to_group.get(r)
        unit = resource_unit_map.get(group)
        cap = resource_caps[r]
        if cap <= 0:
            continue
        used = used_total.get(r, 0.0)
        rows.append({
            "Resource": f"{r} ({unit})",
            "Available amount": cap,
            "Quantity used": used,
            "Quantity remaining": max(0.0, cap - used),
            "% used": (used / cap * 100) if cap > 0 else 0,
        })

    for batch in latest.get("custom_resources", []):
        r = batch["name"]
        # Use resource_caps for the cap (includes expanded batch amounts).
        cap = resource_caps.get(r, float(batch.get("amount", 0)))
        if cap <= 0:
            continue
        used = used_total.get(r, 0.0)
        resource_to_group_local = get_resource_to_group()
        group_name = resource_to_group_local.get(batch.get("group", ""))
        unit = resource_unit_map.get(group_name, "?")
        rows.append({
            "Resource": f"{r} ({unit})",
            "Available amount": cap,
            "Quantity used": used,
            "Quantity remaining": max(0.0, cap - used),
            "% used": (used / cap * 100) if cap > 0 else 0,
        })

    df = pd.DataFrame(rows).sort_values("% used", ascending=False)
    fig = px.bar(
        df, x="Resource", y=["Quantity used", "Quantity remaining"], barmode="stack",
        color_discrete_map={"Quantity used": CRRA_COLORS["green"], "Quantity remaining": CRRA_COLORS["off_white"]},
    )
    apply_chart_style(fig)
    fig.update_layout(yaxis_title="Quantity", xaxis_tickangle=-45)
    st.plotly_chart(fig, width='stretch')
    with st.expander("📊 View / download chart data"):
        st.dataframe(df, width='stretch', hide_index=True)
    register_chart_data("Resource balance (used vs remaining)", df)

def save_scenario(latest):
    """Render the save-scenario widget and persist the current result to session state.

    Args:
        latest (dict): full result dict to store under the user-chosen scenario name.

    Returns:
        None
    """
    st.markdown('<p class="crra-sub-heading">Save scenario</p>', unsafe_allow_html=True)
    st.markdown('<p class="caption">Save the current scenario with inputs and outputs for comparison and/or export.</p>', unsafe_allow_html=True)

    suggested = st.session_state.get("scenario_loaded_file", "").rsplit(".", 1)[0]
    
    # Initialize only if not already set
    if "scenario_name" not in st.session_state:
        st.session_state["scenario_name"] = suggested

    col_name, col_btn = st.columns([3, 1])
    with col_name:
        st.text_input(
            "Scenario name",
            key="scenario_name",
            label_visibility="collapsed",
            placeholder="Enter a scenario name",
        )
    with col_btn:
        save_clicked = st.button("Save", width='stretch', icon=":material/save:")

    if save_clicked:
        name = st.session_state.get("scenario_name", "").strip()
        if not name:
            st.warning("Please enter a name.", icon=":material/warning:")
        elif name in st.session_state.scenarios:
            st.warning("Name already exists.", icon=":material/warning:")
        else:
            st.session_state.scenarios[name] = latest
            st.success(f"Scenario '{name}' saved for comparison.")

def export_scenario():
    """Render JSON and ZIP export download buttons for the currently saved scenario.

    Reads the scenario from st.session_state.scenarios[scenario_name] and
    produces two download buttons: a JSON file and a ZIP of CSVs (allocations,
    resource caps, constraints, costs, balance, usage detail, sensitivity, summary).

    Returns:
        None
    """
    name = st.session_state.get("scenario_name", "").strip()
    if not name or name not in st.session_state.scenarios:
        return
    
    resource_to_group = get_resource_to_group() 

    scenario = st.session_state.scenarios[name]
    safe = unicodedata.normalize("NFKD", name).encode("ascii", "ignore").decode("ascii")
    col_json, col_csv = st.columns(2)
    with col_json:
        st.download_button(
            "Export JSON", data=json.dumps(to_json_safe(scenario), indent=2),
            file_name=f"{safe}.json", mime="application/json",
            width='stretch',
            icon=":material/download:"
        )

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zf:
        zf.writestr("resource_caps.csv",
                    pd.DataFrame.from_dict(scenario["resource_caps"], orient="index",
                                           columns=["Available Amount"]).to_csv())
        zf.writestr(
            "method_constraints.csv",
            pd.DataFrame.from_dict(
                scenario["method_constraints"],
                orient="index",
                columns=["active", "cap_type", "cap_value"]
            ).rename(columns={
                "cap_type": "cap_type (percent or absolute)"
            }).to_csv()
        )

        zf.writestr("method_costs.csv",                      
                    pd.DataFrame.from_dict(scenario["method_costs"], orient="index").to_csv())
        
        alloc_df = pd.DataFrame.from_dict(scenario["allocations"], orient="index",
                                          columns=[f"{CO2_UNIT} removed"])
        alloc_df[f"{CO2_UNIT} removed"] = pd.to_numeric(alloc_df[f"{CO2_UNIT} removed"], errors="coerce")
        alloc_total = alloc_df[f"{CO2_UNIT} removed"].sum()
        if alloc_total > 0:
            alloc_df["Total share (%)"] = (alloc_df[f"{CO2_UNIT} removed"] / alloc_total * 100).round(2)
        method_constraints_export = scenario.get("method_constraints", {})
        alloc_df["Cap constraint"] = pd.Series({
            m: (f"{c['cap_value']:,g} {CO2_UNIT}/yr" if c.get("cap_type") == "absolute"
                else f"{int(c.get('cap_value', 100))}% of its potential")
            for m, c in method_constraints_export.items()
        })
        alloc_df = alloc_df.sort_values(f"{CO2_UNIT} removed", ascending=False)
        zf.writestr("allocations.csv", alloc_df.to_csv())
        
        custom_batches = scenario.get("custom_resources", [])
        if custom_batches:
            custom_rows = []
            for batch in custom_batches:
                for method, coefs in batch["methods"].items():
                    custom_rows.append({
                        "name": batch["name"],
                        "group": batch["group"],
                        "amount": batch["amount"],
                        "unit": resource_unit_map.get(resource_to_group.get(batch["group"])),
                        "method": method,
                        "min": coefs["min"],
                        "median": coefs["median"],
                        "max": coefs["max"],
                    })
            zf.writestr("custom_resources.csv", pd.DataFrame(custom_rows).to_csv(index=False))

        enabled_methods = scenario.get("enabled_methods", {})
        if enabled_methods:
            enabled_rows = [
                {"resource": r, "method": m, "enabled": ok}
                for r, methods in enabled_methods.items()
                for m, ok in methods.items()
            ]
            zf.writestr("enabled_methods.csv", pd.DataFrame(enabled_rows).to_csv(index=False))

        resource_used_total = scenario.get("resource_used_total", {})
        resource_remaining = scenario.get("resource_remaining", {})
        all_resource_caps = dict(scenario.get("resource_caps", {}))
        custom_batch_by_name = {b["name"]: b for b in scenario.get("custom_resources", [])}
        for batch_name, batch in custom_batch_by_name.items():
            all_resource_caps.setdefault(batch_name, float(batch["amount"]))
        if resource_used_total or resource_remaining:
            balance_rows = []
            for r, cap in all_resource_caps.items():
                if cap <= 0:
                    continue
                if r in custom_batch_by_name:
                    unit = resource_unit_map.get(resource_to_group.get(custom_batch_by_name[r]["group"]), "?")
                else:
                    unit = resource_unit_map.get(resource_to_group.get(r), "?")
                used = resource_used_total.get(r, 0.0)
                balance_rows.append({
                    "Resource": r,
                    "Unit": unit,
                    "Available amount": cap,
                    "Quantity used": used,
                    "Quantity remaining": resource_remaining.get(r, cap),
                    "% used": round(used / cap * 100, 1) if cap > 0 else 0,
                })
            zf.writestr("resource_balance.csv", pd.DataFrame(balance_rows).to_csv(index=False))

        resource_usage = scenario.get("resource_usage", {})
        if resource_usage:
            usage_rows = []
            for r, methods in resource_usage.items():
                if r in custom_batch_by_name:
                    unit = resource_unit_map.get(resource_to_group.get(custom_batch_by_name[r]["group"]), "?")
                else:
                    unit = resource_unit_map.get(resource_to_group.get(r), "?")
                for m, qty in methods.items():
                    usage_rows.append({"resource": r, "unit": unit, "method": m, "used": qty})
            zf.writestr("resource_usage_detail.csv", pd.DataFrame(usage_rows).to_csv(index=False))

        resource_actual_gains = scenario.get("resource_actual_gain", {})
        if resource_actual_gains:
            sens_rows = []
            for r, gain in sorted(resource_actual_gains.items(), key=lambda x: -x[1]):
                if r in custom_batch_by_name:
                    unit = resource_unit_map.get(resource_to_group.get(custom_batch_by_name[r]["group"]), "?")
                else:
                    unit = resource_unit_map.get(resource_to_group.get(r), "?")
                sens_rows.append({
                    "resource": r,
                    "unit": unit,
                    f"CDR_gain_for_+1_unit_{CO2_UNIT}": round(gain, 4),
                })
            zf.writestr("resource_sensitivity.csv", pd.DataFrame(sens_rows).to_csv(index=False))

        summary_metrics = scenario.get("summary", {})
        summary_rows = [
            {"metric": "Scenario name",                    "value": name,                                        "unit": ""},
            {"metric": "Export timestamp",                 "value": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "unit": ""},
            {"metric": f"Total {CO2_UNIT} removed",        "value": summary_metrics.get("total_co2_removed", ""),       "unit": CO2_UNIT},
            {"metric": "Total cost (resource units sum)",  "value": summary_metrics.get("total_cost", ""),              "unit": "mixed"},
            {"metric": "Methods used",                     "value": len(scenario.get("allocations", {})),        "unit": "count"},
        ]
        zf.writestr("summary.csv", pd.DataFrame(summary_rows).to_csv(index=False))

    with col_csv:
        st.download_button(
            "Export CSV (zip)", data=zip_buffer.getvalue(),
            file_name=f"{safe}.zip", mime="application/zip",
            width='stretch',
            icon=":material/download:"
        )

def export_inputs_only():
    """Render the inputs-only Excel export widget (no optimisation results).

    Writes an .xlsx with sheets: resource_caps, constraints, coefficients,
    custom_resources, enabled_methods. Disabled (method, resource) pairs are
    exported with coefficient 0 to stay consistent with what the LP sees.

    Returns:
        None
    """
    resource_to_unit = get_resource_to_unit()
    st.markdown('<p class="crra-sub-heading">Export inputs</p>', unsafe_allow_html=True)
    st.markdown('<p class="caption">Export only the resource capabilities, method constraints and coefficients without results.</p>', unsafe_allow_html=True)

    suggested = st.session_state.get("excel_loaded_file", "").rsplit(".", 1)[0]
    # Update when a new file is loaded
    if st.session_state.get("_input_file_name_source") != suggested:
        st.session_state["input_file_name"] = suggested
        st.session_state["_input_file_name_source"] = suggested

    resource_to_group = get_resource_to_group() 

    enabled_methods = st.session_state.get("enabled_methods", {})

    def is_allowed(method, resource):
        return enabled_methods.get(resource, {}).get(method, True)

    wb = openpyxl.Workbook()
    ws1 = wb.active
    ws1.title = "resource_caps"
    ws1.append(["Resource", "Unit", "Available Amount"])
    for r, v in st.session_state.resource_caps.items():
        ws1.append([r, resource_to_unit.get(r, ""), v])

    ws2 = wb.create_sheet("constraints")
    ws2.append(["method", "active", "cap_type (percent or absolute)", "cap_value"])
    for m, c in st.session_state.constraints.items():
        ws2.append([m, c.get("active", True), c.get("cap_type", "percent"), c.get("cap_value", 100)])

    # Coefficients — a disabled (method, resource) pair is always exported
    # as 0, regardless of whether the user ever opened it in tab3. This
    # keeps the export consistent with what the optimiser actually uses.
    ws3 = wb.create_sheet("coefficients")
    ws3.append(["method", "resource", "value", "secondary"])
    for m, resources in st.session_state.costs.items():
        for r, v in resources.items():
            value = 0.0 if not is_allowed(m, r) else v
            ws3.append([m, r, value, (m, r) in secondary_resources])

    ws4 = wb.create_sheet("custom_resources")
    ws4.append(["name", "group", "amount", "unit", "method", "min", "median", "max"])
    for batch in st.session_state.get("custom_resources", []):
        unit = resource_unit_map.get(resource_to_group.get(batch["group"]))
        for method, c in batch["methods"].items():
            ws4.append([
                batch["name"],
                batch["group"],
                batch["amount"],
                unit,
                method,
                c["min"],
                c["median"],
                c["max"],
            ])

    # Per-resource manual method allocation (objective 2) — which CDR
    # methods are allowed to use each resource batch.
    ws5 = wb.create_sheet("enabled_methods")
    ws5.append(["resource", "method", "enabled"])
    for r, methods in enabled_methods.items():
        for m, ok in methods.items():
            ws5.append([r, m, bool(ok)])

    excel_buffer = io.BytesIO()
    wb.save(excel_buffer)

    st.text_input(
        "File name",
        label_visibility="collapsed",
        key="input_file_name",
        placeholder="Enter a file name",
    )
    safe = st.session_state.get("input_file_name", "").strip() or suggested or "crra_inputs"
    st.download_button(
        "Export excel",
        data=excel_buffer.getvalue(),
        file_name=f"{safe}.xlsx",
        mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        width='stretch',
        icon=":material/download:",
    )

def export_all_charts():
    """Render a download widget for all chart data registered in this session.

    Reads chart_data_registry from session state and offers the user a choice of
    Excel (one sheet per chart) or ZIP (one CSV per chart).

    Returns:
        None
    """
    st.markdown('<p class="crra-sub-heading">Export all chart data</p>', unsafe_allow_html=True)
    st.markdown(
        '<p class="caption">Download a single file containing the underlying data '
        'of every chart shown in this session.</p>',
        unsafe_allow_html=True
    )

    registry = st.session_state.get("chart_data_registry", {})
    if not registry:
        st.info("No charts have been generated yet in this session.", icon=":material/info:")
        return

    base_name = st.text_input(
        "File name",
        value=None,
        key="charts_export_name",
        label_visibility="collapsed",
        placeholder="File name",
        help="The extension (.xlsx or .zip) is added automatically.",
    )

    safe_name = (base_name or "").strip() or "charts_data"

    xlsx_buffer = io.BytesIO()
    with pd.ExcelWriter(xlsx_buffer, engine="openpyxl") as writer:
        for name, df in registry.items():
            df.to_excel(writer, sheet_name=name[:31], index=False)

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zf:
        for name, df in registry.items():
            safe_chart_name = name.replace(" ", "_").replace("/", "-").replace("(", "").replace(")", "")
            zf.writestr(f"{safe_chart_name}.csv", df.to_csv(index=False))

    col_xlsx, col_zip = st.columns(2)
    with col_xlsx:
        st.download_button(
            "Export Excel",
            data=xlsx_buffer.getvalue(),
            file_name=f"{safe_name}.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            width='stretch',
            icon=":material/download:",
        )
    with col_zip:
        st.download_button(
            "Export csv (ZIP)",
            data=zip_buffer.getvalue(),
            file_name=f"{safe_name}.zip",
            mime="application/zip",
            width='stretch',
            icon=":material/download:",
        )

def clear_all_data():
    """Wipe all scenario data and widget keys from session state, then rerun.

    Removes resource caps, constraints, costs, scenario list, file-load markers,
    and all slider/toggle widget keys. Increments uploader keys to force re-render.

    Returns:
        None
    """
    for key in ["resource_caps", "constraints", "costs", "scenarios",
                "latest_result", "scenario_loaded_file", "scenario_name",
                "excel_loaded_file", "input_file_name", "resource_last_loaded_file",
                "resource_uploader_key", "custom_resources"]:
        st.session_state.pop(key, None)

    for uploader_key in ["resource_uploader_key", "excel_uploader_key", "scenario_uploader_key"]:
        st.session_state[uploader_key] = st.session_state.get(uploader_key, 0) + 1

    clear_widget_keys()
    st.success("All data cleared.")
    st.rerun()

@st.dialog("Confirm deletion")
def confirm_clear_all():
    """Confirmation dialog before wiping all session data.

    Shows a destructive-action warning and two buttons: confirm (calls clear_all_data)
    or cancel (closes the dialog).

    Returns:
        None
    """
    st.error("This action will delete all data from the current scenario(s).", icon=":material/warning:")
    col1, col2 = st.columns([2, 1])
    with col1:
        if st.button("Yes, delete everything", type="primary", width='stretch', icon=":material/delete:"):
            clear_all_data()
            st.rerun()
    with col2:
        if st.button("Cancel", width='stretch', type="secondary", key="btn-cancel"):
            st.rerun()

def reset_scenario(tab_name):
    """Render the reset section with a button that opens the confirm_clear_all dialog.

    Returns:
        None
    """
    st.markdown('<p class="crra-sub-heading">Reset</p>', unsafe_allow_html=True)
    st.markdown('<p class="caption">Clear all current inputs, results and saved scenarios from the session.</p>', unsafe_allow_html=True)

    if st.button("Clear all data", key=f"reset_scenario_{tab_name}", type="primary", icon=":material/delete:"):
        confirm_clear_all()


@st.fragment
def render_export_tab(latest):
    """Render the full export/save section as a fragment so Save reruns only this area."""
    save_scenario(latest)
    export_scenario()
    export_inputs_only()
    export_all_charts()
    reset_scenario("export")


def render_tab():
    """Render the Portfolio Generation tab.

    Orchestrates the full optimisation workflow: input panels, run button,
    result metrics, method-allocation table and donut chart, resource-usage
    heatmap and bar charts, bottleneck analysis, and export/save controls.

    Returns:
        None
    """
    st.markdown('<h1 class="crra-heading">Portfolio generation</h1>', unsafe_allow_html=True)
    show_bar()

    st.markdown(
        """
        <p class="objective_title">OBJECTIVE</p>
        """,
        unsafe_allow_html=True,
    )

    st.markdown(
        f"""
        <div class="text_with_border">
            <p>{section_intros.get("tab4_optimisation", "")}</p>
        </div>
        """,
        unsafe_allow_html=True
    )

    # resource_caps / constraints / costs / scenarios / enabled_methods are
    # already guaranteed by app.py's shared init (runs before any tab renders)
    # — only these two upload-widget counters are tab4-specific.
    for key, default in [
        ("excel_uploader_key", 0),
        ("scenario_uploader_key", 0)
    ]:
        st.session_state.setdefault(key, default)

    tab_run, tab_results, tab_resources, tab_export, tab_reset = st.tabs([
        "▶️ Run",
        "📊 Method allocation",
        "🌿 Resource usage",
        "💾 Export & Save",
        "🔄 Reset"
    ])

    with tab_run:
        st.markdown("""
            <p class="crra-sub-heading">Run optimisation</p>""",
            unsafe_allow_html=True,
        )
        col_inputs, col_scenario = st.columns(2, gap="large")

        with col_inputs:
            st.markdown("""
                <p class="crra-sub-sub-heading">From inputs</p>""",
                unsafe_allow_html=True,
            )
            run_clicked = st.button("Run optimisation", icon=":material/play_arrow:", key="btn-run-optimisation", type="primary")
            run_msg_area = st.container()
            panel_run_from_tabs()
            panel_load_from_excel()

        with col_scenario:
            st.markdown("""
                <p class="crra-sub-sub-heading">From a saved scenario</p>""",
                unsafe_allow_html=True,
            )
            panel_load_from_json_or_csv()

        if run_clicked:
            with run_msg_area:
                run_and_store_result()

    no_results = '<p class="caption">Run the optimisation first to see results here.</p>'

    latest = st.session_state.get("latest_result")
    if not latest:
        for tab in [tab_results, tab_resources, tab_export]:
            with tab:
                st.markdown(no_results, unsafe_allow_html=True)
        with tab_reset:
            reset_scenario("reset")
        return

    allocations = latest["allocations"]
    resource_caps = latest["resource_caps"]
    method_costs = latest["method_costs"]

    if not allocations or sum(allocations.values()) == 0:
        for tab in [tab_results, tab_resources, tab_export]:
            with tab:
                st.markdown(no_results, unsafe_allow_html=True)
        with tab_reset:
            reset_scenario("reset")
        return

    df_result = build_results_df(
        allocations,
        method_constraints=latest.get("method_constraints")
    )
    df_display = df_result.copy()
    df_display[f"{CO2_UNIT} removed"] = df_display[f"{CO2_UNIT} removed"].map(
        lambda x: f"{x:,g}" if pd.notna(x) else x
    )

    with tab_results:
        show_metrics(latest)
        st.markdown('<p class="crra-sub-heading">Optimisation Summary</p>', unsafe_allow_html=True)
        st.dataframe(df_display, width='stretch')
        st.markdown('<p class="crra-sub-heading">Method allocation breakdown</p>', unsafe_allow_html=True)
        st.markdown(
            '<p class="caption">Share of total CO₂ removal contributed by each active method. '
            'Larger slices indicate methods that dominate the portfolio.</p>',
            unsafe_allow_html=True,
        )
        show_donut(df_result)

    with tab_resources:
        st.markdown('<p class="crra-sub-heading">Resource usage heatmap (% of cap)</p>', unsafe_allow_html=True)
        st.markdown(
            '<p class="caption">Each cell shows what percentage of a resource\'s available cap is consumed by a given method.'
            '<br>Dark cells indicate near-full use: those resource–method pairs are the most constrained.</p>',
            unsafe_allow_html=True,
        )
        show_heatmap(df_result, resource_caps, method_costs,
                     custom_resources=latest.get("custom_resources", []),
                     resource_usage=latest.get("resource_usage"))
        st.markdown('<p class="crra-sub-heading">Absolute resource usage by unit</p>', unsafe_allow_html=True)
        st.markdown(
            '<p class="caption">Absolute quantities consumed per resource and per method, in their native units. '
            'Use this view to compare the scale of demand across resources with different units (e.g. ha, m³/yr, t/yr).</p>',
            unsafe_allow_html=True,
        )
        show_absolute_usage(allocations, method_costs,
                            custom_resources=latest.get("custom_resources", []),
                            resource_usage=latest.get("resource_usage"))
        show_resource_balance(latest)

    with tab_export:
        render_export_tab(latest)

    with tab_reset:
        reset_scenario("reset")

    if st.button("Sensitivity analysis →", key="btn-nav-next-optimisation", type="primary", help="Explore how resource or method changes affect your portfolio."):
        st.session_state["_navigate_to"] = "Sensitivity analysis"
        st.rerun()