from collections import defaultdict
import datetime
import io
import json
import re
import unicodedata
import zipfile
import openpyxl
import numpy as np
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
from optimization.optimization import run_optimization
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


def to_json_safe(obj):
    """Recursively convert numpy/pandas types to plain Python so json.dumps works."""
    if isinstance(obj, dict):
        return {k: to_json_safe(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [to_json_safe(v) for v in obj]
    if isinstance(obj, pd.Series):
        return obj.iloc[0] if len(obj) == 1 else obj.tolist()
    if isinstance(obj, pd.DataFrame):
        return obj.to_dict(orient="records")
    if isinstance(obj, np.integer):
        return int(obj)
    if isinstance(obj, np.floating):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    return obj


def panel_run_from_tabs():
    """Left sub-panel: run optimization 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**")

    with col_upload:
        with open("data/inputs_template.xlsx", "rb") as f:
            st.download_button(
                label="Excel Template",
                data=f,
                file_name="crra_input_template.xlsx",
                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                width='stretch',
                icon=":material/download:",
                key="download_button_input_template"
            )

    st.markdown('<p class="caption">Fill in the template with your values, then import and run.</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
            clear_widget_keys()
            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 optimizer 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:
                constraints = {
                    m: {
                        "active": bool(row["active"]),
                        "cap_type": str(row["cap_type (percent or absolute)"]),
                        "cap_value": float(row["cap_value"])
                    }
                    for m, row in df_c.iterrows()
                }
            else:
                constraints = {
                    m: {
                        "active": bool(row["active"]),
                        "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"])

            safe_constraints = neutralize_zero_coefficient_methods(
                st.session_state.costs, st.session_state.constraints
            )
            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
            ]
            
            # Recompute detailed outputs from the reloaded inputs so the scenario
            # is internally consistent (rather than trusting allocations.csv alone).
            success, result = run_optimization(
                resource_caps=resource_caps,
                method_constraints=constraints,
                method_costs=costs,
                custom_resources=custom_resources,
                enabled_methods=enabled_methods,
            )

            if neutralized:
                st.info(
                    f"ℹ️ {len(neutralized)} méthode(s) désactivée(s) automatiquement (aucune ressource active) : "
                    + ", ".join(neutralized)
                )

            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 - Optimization 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_optimization(
                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", {}),
            )
            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.")
    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 optimization", 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()
        st.session_state.costs = 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", {})
        clear_widget_keys()
        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.")
        st.rerun()

def clear_widget_keys():
    """Delete all Streamlit widget state keys for sliders, manual inputs, and toggles.

    Called after loading a new file or scenario so widgets re-initialise from
    fresh session-state values instead of showing stale cached values.

    Returns:
        None
    """
    prefixes = ("slider_", "manual_", "toggle_")
    for k in [k for k in st.session_state if k.startswith(prefixes)]:
        del st.session_state[k]


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_optimization, 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.")
        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)

    safe_constraints = neutralize_zero_coefficient_methods(
        applied_costs, st.session_state.constraints
    )
    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
    ]

    success, result = run_optimization(
        resource_caps=st.session_state.resource_caps,
        method_constraints=safe_constraints,
        method_costs=st.session_state.costs,
        custom_resources=custom_resources,
        enabled_methods=enabled_methods,
    )

    if not success:
        st.error("❌ Optimization failed")
        if "message" in result:
            st.error(result["message"])
        return
    
    if neutralized:
        st.info(
            f"ℹ️ {len(neutralized)} méthode(s) désactivée(s) automatiquement (aucune ressource active) : "
            + ", ".join(neutralized)
        )

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

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

    total_cost = sum(
        allocations[m] * sum(st.session_state.costs[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": st.session_state.resource_caps.copy(),
        "method_constraints": st.session_state.constraints,
        "method_costs": st.session_state.costs,
        "custom_resources": custom_resources,
        "enabled_methods": enabled_methods,
    }
    if custom_resources:
        st.info(f"ℹ️ {len(custom_resources)} new resource(s) included in optimisation.")
    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_optimization 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 build_results_df(allocations, method_constraints=None):
    """Build a DataFrame of per-method CO₂ allocations with share and cap columns.

    Args:
        allocations (dict): {method: CO₂ removed (float)}.
        method_constraints (dict | None): {method: {cap_type, cap_value}} for the Cap column.

    Returns:
        pd.DataFrame: indexed by method, columns CO₂ removed, Total share (%), Cap.
                      Sorted descending by CO₂ removed.
    """
    df = pd.DataFrame.from_dict(allocations, orient="index", columns=[f"{CO2_UNIT} removed"])

    df[f"{CO2_UNIT} removed"] = pd.to_numeric(df[f"{CO2_UNIT} removed"], errors="coerce")
    total = df[f"{CO2_UNIT} removed"].sum()
    if total > 0:
        df["Total share (%)"] = (df[f"{CO2_UNIT} removed"] / total * 100).round(2)

    if method_constraints:
        caps = {}
        for m in df.index:
            c = method_constraints.get(m, {})
            
            if c.get("cap_type") == "absolute":
                caps[m] = f"{c['cap_value']:,.2f} {CO2_UNIT}/yr"
            else:
                caps[m] = f"{int(c.get('cap_value', 100))}% of potential"
        df["Cap"] = pd.Series(caps)
    
    return df.sort_values(f"{CO2_UNIT} removed", ascending=False)


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.")
        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(used / cap * 100)
            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)

    fig = px.imshow(
        df_heat, color_continuous_scale=PLOTLY_HEATMAP_SCALE, text_auto=".1f",
        labels=dict(x="Resource", y="Method", color="% Used"),
    )
    matrix_exact_fmt = [[f"{v:.4f}" for v in row] for row in matrix_exact]
    fig.update_traces(
        customdata=matrix_exact_fmt,
        hovertemplate="<b>%{y}</b><br>%{x}<br>% of cap: %{z:.1f}%<br>Exact: %{customdata}<extra></extra>",
    )
    apply_chart_style(fig, height=max(400, 60 + 30 * len(df_heat.index)))
    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"):
        st.dataframe(df_export_exact, 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)
    resource_caps = latest["resource_caps"]
    used_total = latest["resource_used_total"]
    remaining = latest["resource_remaining"]

    resource_to_group = get_resource_to_group()
        
    rows = []
    for r in resource_caps:
        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": remaining.get(r, cap),
            "% used": round(used / cap * 100, 1) if cap > 0 else 0,
        })

    for batch in latest.get("custom_resources", []):
        r = batch["name"]
        cap = float(batch["amount"])
        used = used_total.get(r, 0.0)
        rows.append({
            "Resource": r, "Available amount": cap, "Quantity used": used,
            "Quantity remaining": remaining.get(r, cap),
            "% used": round(used / cap * 100, 1) 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 show_bottleneck_analysis(latest):
    """Render the resource sensitivity table (CDR gain per +1 unit of each resource).

    Args:
        latest (dict): result dict with keys resource_actual_gain, custom_resources.

    Returns:
        None
    """
    st.markdown('<p class="crra-sub-heading">Resource sensitivity</p>', unsafe_allow_html=True)
    st.markdown(
        '<p class="caption">'
        'For every resource: actual CDR gain if you add exactly +1 unit, recomputed with all other constraints in place.'
        '</p>',
        unsafe_allow_html=True,
    )

    resource_to_group = get_resource_to_group()
    custom_batch_lookup = {b["name"]: b for b in latest.get("custom_resources", [])}

    actual_gains = latest.get("resource_actual_gain", {})
    if actual_gains:
        sp_rows = []
        for r, gain in sorted(actual_gains.items(), key=lambda x: -x[1]):
            if r in custom_batch_lookup:
                unit = resource_unit_map.get(resource_to_group.get(custom_batch_lookup[r]["group"]), "?")
            else:
                unit = resource_unit_map.get(resource_to_group.get(r), "?")
            sp_rows.append({
                "Resource": r,
                "Unit": unit,
                f"CDR gain for +1 unit ({CO2_UNIT})": round(gain, 4),
            })
        df_sp = pd.DataFrame(sp_rows)
        with st.expander("📊 View / download resource sensitivity", expanded=True):
            st.dataframe(df_sp, width='stretch', hide_index=True)
        register_chart_data("Resource sensitivity (CDR gain per +1 unit)", df_sp)
    else:
        st.info("No resources to analyse.")

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:
        if st.button("Save", width='stretch', icon=":material/save:"):
            name = st.session_state.get("scenario_name", "").strip()
            if not name:
                st.warning("Please enter a name.")
            elif name in st.session_state.scenarios:
                st.warning("Name already exists.")
            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']:,.2f} {CO2_UNIT}/yr" if c.get("cap_type") == "absolute"
                else f"{int(c.get('cap_value', 100))}% of 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() 

    col_text, col_btn = st.columns([3, 1])
    with col_text:
        st.text_input(
            "File name",
            label_visibility="collapsed",
            key="input_file_name",
            placeholder="Enter a file name",
        )
    with col_btn:
        prepare = st.button("Prepare export", width='stretch', icon=":material/file_export:")

    if prepare:
        safe = st.session_state.get("input_file_name", "").strip() or suggested or "crra_inputs"

        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 optimizer 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.session_state["_input_export_buf"] = excel_buffer.getvalue()
        st.session_state["_input_export_name"] = safe

    if "_input_export_buf" in st.session_state:
        st.download_button(
            f"Download **{st.session_state['_input_export_name']}.xlsx**",
            data=st.session_state["_input_export_buf"],
            file_name=f"{st.session_state['_input_export_name']}.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.")
        return

    col_name, col_fmt = st.columns([2, 1])
    with col_name:
        base_name = st.text_input(
            "File name",
            value="charts_data",
            key="charts_export_name",
            help="The extension (.xlsx or .zip) is added automatically.",
        )
    with col_fmt:
        fmt = st.radio(
            "Format", ["Excel (.xlsx, one sheet per chart)", "ZIP (one CSV per chart)"],
            horizontal=True, key="chart_export_format",
        )

    safe_name = base_name.strip() or "charts_data"

    if fmt.startswith("Excel"):
        export_buffer = io.BytesIO()
        with pd.ExcelWriter(export_buffer, engine="openpyxl") as writer:
            for name, df in registry.items():
                sheet_name = name[:31]  # limite Excel
                df.to_excel(writer, sheet_name=sheet_name, index=False)
        st.download_button(
            f"Download {safe_name}.xlsx",
            data=export_buffer.getvalue(),
            file_name=f"{safe_name}.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            icon=":material/download:"
        )
    else:
        export_buffer = io.BytesIO()
        with zipfile.ZipFile(export_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))

        st.download_button(
            f"Download {safe_name}.zip",
            data=export_buffer.getvalue(),
            file_name=f"{safe_name}.zip",
            mime="application/zip",
            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'):
            st.rerun()

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

    Returns:
        None
    """
    st.divider()
    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="reset_scenario", type="primary", icon=":material/delete:"):
        confirm_clear_all()


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(
        """
        <div class="text_with_border">
            <p>Run the optimization tool, explore results and save scenarios for comparison or export.</p>
        </div>
        """,
        unsafe_allow_html=True
    )

    for key, default in [
        ("resource_caps", {}), ("constraints", {}),
        ("costs", {}), ("scenarios", {}),
        ("excel_uploader_key", 0),
        ("scenario_uploader_key", 0)
    ]:
        st.session_state.setdefault(key, default)
    if "enabled_methods" not in st.session_state:
        st.session_state.enabled_methods = {}

    st.markdown("""
        <p class="crra-sub-heading">Run optimization</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 Optimization", icon=":material/play_arrow:", key="btn-run-optimization", type="primary")
        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:
        run_and_store_result()

    if "latest_result" not in st.session_state:
        return

    latest = st.session_state.latest_result
    allocations = latest["allocations"]
    resource_caps = latest["resource_caps"]
    method_costs = latest["method_costs"]

    if not allocations or sum(allocations.values()) == 0:
        return

    show_metrics(latest)

    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:,.4f}" if pd.notna(x) else x
    )

    tab_results, tab_resources, tab_bottleneck, tab_export = st.tabs([
        "📊 Method allocation",
        "🌿 Resource usage",
        "🔍 Bottleneck analysis",
        "💾 Export & Save",
    ])

    with tab_results:
        st.markdown('<p class="crra-sub-heading">Optimization 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)
        show_donut(df_result)

    with tab_resources:
        st.markdown('<p class="crra-sub-heading">Resource usage heatmap (% of cap)</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)
        show_absolute_usage(allocations, method_costs,
                            custom_resources=latest.get("custom_resources", []),
                            resource_usage=latest.get("resource_usage"))
        show_resource_balance(latest)

    with tab_bottleneck:
        show_bottleneck_analysis(latest)

    with tab_export:
        save_scenario(latest)
        export_scenario()
        export_inputs_only()
        export_all_charts()
        reset_scenario()