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, method_order, resource_order
from optimization.optimization import run_optimization
from utils.data_utils import get_resource_to_unit
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_optimization(
        resource_caps=st.session_state.resource_caps,
        method_constraints=st.session_state.constraints,
        method_costs=st.session_state.costs,
        custom_resources=st.session_state.get("custom_resources", []),
    )
    baseline_removed = baseline_result["total_removed"] if baseline_success else 0

    custom_resource_names = [b["name"] for b in st.session_state.get("custom_resources", [])]
    all_resources = resource_order + 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,
    )
    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.")
        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_optimization(
                resource_caps=perturbed_caps,
                method_constraints=st.session_state.constraints,
                method_costs=st.session_state.costs,
                custom_resources=st.session_state.get("custom_resources", []),
            )
            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.
    Skips methods whose current cap is 0.

    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,
    )
    selected_methods = st.multiselect(
        "Select one or more methods:", method_order, key="crra-select-method"
    )

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

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

    for method in selected_methods:
        cap_type = st.session_state.constraints[method].get("cap_type", "percent")
        current_cap = st.session_state.constraints[method].get(
            "cap_value", 100 if cap_type == "percent" else 0.0
        )

        if current_cap <= 0:
            skipped.append(method)
            continue

        for pct in percentages:
            perturbed_constraints = copy.deepcopy(st.session_state.constraints)
            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 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_optimization(
                resource_caps=st.session_state.resource_caps,
                method_constraints=perturbed_constraints,
                method_costs=st.session_state.costs,
                custom_resources=st.session_state.get("custom_resources", []),
            )
            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,
            })

    for m in skipped:
        st.warning(f"Method **{m}** has a cap of 0 — set a value first to include it.", icon="⚠️")

    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 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. Requires resource_caps, constraints, and costs
    to be present in session state.

    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(
        '<div class="text_with_border">'
        "<p>Explore how varying <strong>one or multiple resources' availability</strong> "
        "or <strong>methods' maximum allocation</strong> affects total CO₂ removed.</p>"
        "</div>",
        unsafe_allow_html=True,
    )

    if (
        "resource_caps" not in st.session_state
        or "constraints" not in st.session_state
        or "costs" not in st.session_state
    ):
        st.warning("Please complete the setup in previous tabs before running sensitivity analysis.")
        return

    baseline_removed, all_resources, resource_to_unit = compute_baseline()

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

    with tab1:
        panel_resource_sensitivity(baseline_removed, all_resources, resource_to_unit)

    with tab2:
        panel_method_sensitivity(baseline_removed)