tabs/tab5_sensitivity.py

import copy
import pandas as pd
import streamlit as st
import plotly.express as px
from config.config import CO2_UNIT, hidden_resource_names, method_order, resource_order, resource_unit_map, section_intros
from utils.data_utils import (get_resource_to_unit, get_resource_to_group, run_expanded_optimisation,
                              is_other_method, get_other_method_variants, make_variant_name,
                              default_constraint)
from utils.ui_helpers import register_chart_data, show_bar


def compute_baseline():
    """Run the baseline optimisation and build combined resource metadata.

    Returns:
        tuple[float, list[str], dict[str, str]]:
            baseline_removed  – total CO₂ removed at current inputs (0 if failed).
            all_resources     – standard resources + custom batch names.
            resource_to_unit  – {resource: unit} including custom batches.
    """
    baseline_success, baseline_result = run_expanded_optimisation(
        resource_caps=st.session_state.resource_caps,
        method_constraints=st.session_state.constraints,
        costs=st.session_state.costs,
        custom_resources=st.session_state.get("custom_resources", []),
        enabled_methods=st.session_state.get("enabled_methods", {}),
    )
    baseline_removed = baseline_result["total_removed"] if baseline_success else 0

    # Only offer resources that actually have a quantity set — sweeping a
    # resource currently at 0 always gives 0 regardless of the % slider
    # (any percentage of a 0 base cap is still 0), so it's never useful to
    # compare.
    custom_resource_names = [
        b["name"] for b in st.session_state.get("custom_resources", [])
        if float(b.get("amount", 0)) > 0
    ]
    all_resources = [
        r for r in resource_order
        if r not in hidden_resource_names and st.session_state.resource_caps.get(r, 0) > 0
    ] + custom_resource_names

    resource_to_unit = get_resource_to_unit()
    for b in st.session_state.get("custom_resources", []):
        resource_to_unit[b["name"]] = b.get("unit", "")

    return baseline_removed, all_resources, resource_to_unit


def panel_resource_sensitivity(baseline_removed, all_resources, resource_to_unit):
    """Render the resource-availability sensitivity sub-panel.

    Sweeps each selected resource from 0 % to 200 % of its current cap in
    10 % steps, re-solves the LP at each point, and plots total CO₂ removed.

    Args:
        baseline_removed (float): total CO₂ removed at current inputs.
        all_resources (list[str]): standard + custom resource names to select from.
        resource_to_unit (dict[str, str]): {resource: unit} for axis labels.

    Returns:
        None
    """
    st.markdown(
        '<p class="crra-sub-heading">Impact of resource availability</p>',
        unsafe_allow_html=True,
    )
    st.markdown(
        '<p class="caption">Each line shows how total CO₂ removal changes as one resource\'s availability varies from 0% to 200% of its current cap. '
        '<br>The range goes up to 200% so you can see both the impact of a reduction and the potential gains from doubling availability. '
        '<br>A steep curve means that resource strongly drives portfolio performance: small increases can yield large gains. '
        '<br>A flat line means the portfolio is limited by something else (another resource or a method cap).</p>',
        unsafe_allow_html=True,
    )
    selected_resources = st.multiselect(
        "Select one or more resources:", all_resources, key="crra-select-resource"
    )

    if not selected_resources:
        st.warning("Please select at least one resource to run the sensitivity analysis.", icon=":material/warning:")
        return

    percentages = list(range(0, 210, 10))
    custom_amount_by_name = {
        b["name"]: float(b.get("amount", 0))
        for b in st.session_state.get("custom_resources", [])
    }
    all_data = []

    for resource in selected_resources:
        unit = resource_to_unit.get(resource, "")
        base_cap = (
            st.session_state.resource_caps.get(resource)
            or custom_amount_by_name.get(resource, 0)
        )
        for pct in percentages:
            perturbed_caps = st.session_state.resource_caps.copy()
            actual_amount = base_cap * pct / 100
            perturbed_caps[resource] = actual_amount

            success, result = run_expanded_optimisation(
                resource_caps=perturbed_caps,
                method_constraints=st.session_state.constraints,
                costs=st.session_state.costs,
                custom_resources=st.session_state.get("custom_resources", []),
                enabled_methods=st.session_state.get("enabled_methods", {}),
            )
            all_data.append({
                "Input": resource,
                "% of cap": pct,
                "Actual amount": round(actual_amount, 4),
                "Unit": unit,
                f"CO₂ removed ({CO2_UNIT})": result["total_removed"] if success else 0,
            })

    df = pd.DataFrame(all_data)
    st.markdown("#### Sensitivity to resource availability")
    fig = px.line(
        df, x="% of cap", y=f"CO₂ removed ({CO2_UNIT})", color="Input", markers=True,
        hover_data={"Actual amount": ":.4f", "Unit": True},
    )
    fig.update_layout(
        xaxis_title="Resource availability (% of current cap)",
        yaxis_title=f"Total portfolio CO₂ removal ({CO2_UNIT})",
    )
    fig.add_hline(
        y=baseline_removed, line_dash="dash", line_color="gray",
        annotation_text=f"Baseline: {baseline_removed:,} {CO2_UNIT}",
    )
    st.plotly_chart(fig, width='stretch')
    with st.expander("📊 View / download chart data"):
        st.dataframe(df, width='stretch', hide_index=True)
    register_chart_data("Sensitivity to resource availability", df)


def panel_method_sensitivity(baseline_removed):
    """Render the method-cap sensitivity sub-panel.

    Sweeps each selected method's cap from 0 % to 100 % of its current value
    in 10 % steps, re-solves the LP at each point, and plots total CO₂ removed.
    Methods whose current cap is already 0 are excluded from the selector
    entirely (sweeping "X% of 0" is always 0, never useful to compare).

    Args:
        baseline_removed (float): total CO₂ removed at current inputs.

    Returns:
        None
    """
    st.markdown(
        '<p class="crra-sub-heading">Impact of max allocation per method</p>',
        unsafe_allow_html=True,
    )
    st.markdown(
        '<p class="caption">Each line shows how total CO₂ removal changes as one method\'s maximum deployment cap varies from 0% to 100% of its current setting: '
        'whether that cap was defined as a percentage of its potential or as an absolute value in the Method constraints tab. '
        '<br>The scale stops at 100% because the cap represents a ceiling you set: going beyond it would mean overriding your own constraint. '
        '<br>A steep curve means the cap is a binding constraint: relaxing it would allow more removal. '
        '<br>A flat line means the method is already limited by resource availability, not by its cap.</p>',
        unsafe_allow_html=True,
    )

    # "Other"-tagged methods with custom batches attached are expanded into
    # independent variants (e.g. "Enhanced weathering - Other mineral: Ganite")
    # by run_expanded_optimisation — offer those variants as selectable
    # entries too, alongside the standard method names, so their own cap can
    # be swept just like any other method. An "Other"-tagged method with NO
    # batch attached has no variant AND can never produce any result on its
    # own (its only resource is the hidden "Other X" pool, which is never
    # directly settable) — skip it entirely rather than offer a sweep that
    # would always show a flat 0 line.
    #
    # Same reasoning as the resource panel: a method whose current cap is
    # already 0 would sweep to "X% of 0", always 0 regardless of the slider —
    # never useful to compare — so it's excluded from the list up front
    # instead of being selectable and only warned about afterwards. Same for
    # a method that is simply inactive (missing resources, manually toggled
    # off in Tab 2...): sweeping its cap_value doesn't touch "active", so the
    # perturbed run stays forced to 0 regardless — just as pointless to offer.
    def _current_cap(method_name):
        constraint = st.session_state.constraints.get(method_name) or default_constraint()
        cap_type = constraint.get("cap_type", "percent")
        return constraint.get("cap_value", 100 if cap_type == "percent" else 0.0)

    def _is_sweepable(method_name):
        constraint = st.session_state.constraints.get(method_name) or default_constraint()
        return constraint.get("active", True) and _current_cap(method_name) > 0

    custom_resources = st.session_state.get("custom_resources", [])
    method_options = []
    for m in method_order:
        variants = get_other_method_variants(m, custom_resources)
        if variants:
            method_options.extend(
                make_variant_name(m, b["name"]) for b in variants
                if _is_sweepable(make_variant_name(m, b["name"]))
            )
        elif not is_other_method(m) and _is_sweepable(m):
            method_options.append(m)

    selected_methods = st.multiselect(
        "Select one or more methods:", method_options, key="crra-select-method"
    )

    if not selected_methods:
        st.warning("Please select at least one method to run the sensitivity analysis.", icon=":material/warning:")
        return

    percentages = list(range(0, 110, 10))
    all_data = []

    for method in selected_methods:
        constraint = st.session_state.constraints.get(method) or default_constraint()
        cap_type = constraint.get("cap_type", "percent")
        current_cap = _current_cap(method)

        for pct in percentages:
            perturbed_constraints = copy.deepcopy(st.session_state.constraints)
            perturbed_constraints.setdefault(method, default_constraint())
            if cap_type == "percent":
                actual_cap = min(100, current_cap * pct / 100)
                perturbed_constraints[method]["cap_value"] = actual_cap
                cap_label = f"{actual_cap:.0f}% of its potential"
            else:
                actual_cap = current_cap * pct / 100
                perturbed_constraints[method]["cap_value"] = actual_cap
                cap_label = f"{actual_cap:.2f} {CO2_UNIT}/yr"

            success, result = run_expanded_optimisation(
                resource_caps=st.session_state.resource_caps,
                method_constraints=perturbed_constraints,
                costs=st.session_state.costs,
                custom_resources=custom_resources,
                enabled_methods=st.session_state.get("enabled_methods", {}),
            )
            all_data.append({
                "Method": method,
                "Cap type": cap_type,
                "% of current cap": pct,
                "Actual cap": cap_label,
                f"CO₂ removed ({CO2_UNIT})": result["total_removed"] if success else 0,
            })

    if not all_data:
        return

    df = pd.DataFrame(all_data)
    fig = px.line(
        df, x="% of current cap", y=f"CO₂ removed ({CO2_UNIT})",
        color="Method", markers=True,
        hover_data={"Actual cap": True, "Cap type": True},
    )
    fig.add_hline(
        y=baseline_removed, line_dash="dash", line_color="lightgray",
        annotation_text=f"Baseline: {baseline_removed:,} {CO2_UNIT}",
    )
    fig.update_layout(
        xaxis_title="% of current cap setting",
        yaxis_title=f"Total portfolio CO₂ removal ({CO2_UNIT})",
    )
    st.plotly_chart(fig, width='stretch')
    with st.expander("📊 View / download chart data"):
        st.dataframe(df, width='stretch', hide_index=True)
    register_chart_data("Sensitivity to method cap", df)


def show_bottleneck_analysis():
    """Render the resource sensitivity table (CDR gain per +1 unit of each resource).

    Reads the latest optimisation result from session state. Shows nothing if no
    result is available yet.

    Returns:
        None
    """
    latest = st.session_state.get("latest_result")
    if not latest:
        st.info("Run the optimisation first (Portfolio generation tab) to see the bottleneck analysis.", icon=":material/info:")
        return

    st.markdown('<p class="crra-sub-heading">Resource sensitivity</p>', unsafe_allow_html=True)
    st.markdown(
        '<p class="caption">'
        'For each resource: the additional CO₂ that could be removed by adding exactly +1 unit, '
        'recomputed with all other constraints held fixed. '
        '<br>Resources at the top of the list are the most binding bottlenecks in your current portfolio: '
        'prioritise increasing their availability to unlock the largest gains.'
        '</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)
        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.", icon=":material/info:")


def render_tab():
    """Render the Sensitivity Analysis tab.

    Computes a baseline optimisation from session state, then lets the user
    sweep one or more resources (or method caps) over a range and plots how
    total CO₂ removed changes. resource_caps, constraints, and costs are
    guaranteed present by app.py's shared init, which runs before any tab.

    Returns:
        None
    """
    st.markdown('<h1 class="crra-heading">Sensitivity analysis</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("tab5_sensitivity", "")}</p></div>',
        unsafe_allow_html=True,
    )

    baseline_removed, all_resources, resource_to_unit = compute_baseline()

    tab1, tab2, tab3 = st.tabs(["🔁 Resource sensitivity", "🔁 Method sensitivity", "🔍 Bottleneck analysis"])

    with tab1:
        panel_resource_sensitivity(baseline_removed, all_resources, resource_to_unit)

    with tab2:
        panel_method_sensitivity(baseline_removed)

    with tab3:
        show_bottleneck_analysis()

    if st.button("Scenario comparison →", key="btn-nav-next-sensitivity", type="primary", help="Compare multiple saved scenarios side by side."):
        st.session_state["_navigate_to"] = "Scenario comparison"
        st.rerun()