tabs/tab2_constraints.py

import pandas as pd
import streamlit as st

from config.config import CO2_UNIT, method_groups, method_order, method_tooltips, secondary_resources, hidden_resource_names
from utils.ui_helpers import show_bar
from utils.data_utils import strip_unit_suffix, is_other_method, get_other_method_variants, make_variant_name



def default_constraint():
    """Return the default constraint dict for a CDR method (active, 100% cap).

    Returns:
        dict: {"active": True, "cap_type": "percent", "cap_value": 100}.
    """
    return {"active": True, "cap_type": "percent", "cap_value": int(100)}


def _make_slider_cb(m):
    """on_change for the percent-cap slider: write the new value into constraints."""
    def _cb():
        st.session_state.constraints[m]["cap_value"] = st.session_state[f"slider_{m}"]
    return _cb


def _make_manual_cb(m):
    """on_change for the absolute-cap number_input: write the new value into constraints."""
    def _cb():
        v = st.session_state.get(f"manual_{m}")
        st.session_state.constraints[m]["cap_value"] = float(v) if v is not None else 0.0
    return _cb


def get_hidden_resources_for_method(m, method_to_primary_resources, cost_sliders_norm):
    """Return (required, optional) lists of hidden resource names for method m.

    Returns (None, None) if the method has at least one non-hidden primary resource
    or a custom batch already assigned.
    """
    primary = method_to_primary_resources.get(m, [])
    if not primary or not all(r in hidden_resource_names for r in primary):
        return None, None
    custom_resources = st.session_state.get("custom_resources", [])
    if any(m in b.get("methods", {}) for b in custom_resources):
        return None, None
    required = [r for r in primary if r in hidden_resource_names]
    optional = [
        r for r in cost_sliders_norm.get(m, {})
        if r in hidden_resource_names and (m, r) in secondary_resources
    ]
    return required, optional


def get_uncovered_hidden_resources(m, cost_sliders_norm):
    """Return (required, optional) hidden resources for method m that leave it unconstrained.

    A hidden resource is "uncovered" only when the user has provided no resource
    at all from the same group — neither a standard cap (> 0 in tab1) nor a
    custom batch assigned to this method.  If at least one standard resource in
    the group has a positive cap, the method is already constrained by it and the
    missing hidden resource does not inflate results.

    Returns:
        tuple[list[str], list[str]]: (required, optional) uncovered hidden resource names.
            required — must be added for the method to run (not in secondary_resources).
            optional — in secondary_resources, shown as informational only.
            Returns ([], []) when nothing is uncovered.
    """
    custom_resources = st.session_state.get("custom_resources", [])
    covered = {b["name"] for b in custom_resources if m in b.get("methods", {})}

    required, optional = [], []
    for r, coeff_dict in cost_sliders_norm.get(m, {}).items():
        if r not in hidden_resource_names:
            continue
        if coeff_dict.get("median", 0.0) <= 0:
            continue
        if r in covered:
            continue
        if (m, r) in secondary_resources:
            optional.append(r)
        else:
            required.append(r)
    return required, optional


def _render_resource_status(m, cost_sliders_norm):
    """Render an expander showing all resources of method m with their availability status.

    Standard resources: ✅ if cap > 0, ❌ if = 0.
    Hidden resources:   ✅ if covered by a custom batch, ❌ if not.
    Secondary resources are shown in a separate section.
    """
    resource_caps   = st.session_state.get("resource_caps", {})
    custom_resources = st.session_state.get("custom_resources", [])
    covered = {b["name"] for b in custom_resources if m in b.get("methods", {})}

    all_res = cost_sliders_norm.get(m, {})
    required_rows = []
    optional_rows = []

    for r, coeff_dict in all_res.items():
        if coeff_dict.get("median", 0.0) <= 0:
            continue
        if r in hidden_resource_names:
            available = r in covered
            source = "Add as new resource in Resource inputs"
        else:
            available = resource_caps.get(r, 0.0) > 0
            source = "Set quantity in Resource inputs"

        icon = "✅" if available else "❌"
        row = (icon, r, "" if available else source)

        if (m, r) in secondary_resources:
            optional_rows.append(row)
        else:
            required_rows.append(row)

    with st.expander("Why is this method disabled?", expanded=False):
        if required_rows:
            st.markdown("**Required resources**")
            for _, r, hint in required_rows:
                color = "#55C95B" if not hint else "#FF5F4D"
                label = f'<span style="color:{color};font-weight:600"><strong>{r}</strong></span>'
                if hint:
                    label += f' <span style="color:#888;font-style:italic;font-weight:400">: {hint}</span>'
                st.markdown(label, unsafe_allow_html=True)
        if optional_rows:
            st.markdown("**Optional resources**")
            for _, r, hint in optional_rows:
                color = "#55C95B" if not hint else "#FF5F4D"
                label = f'<span style="color:{color};font-weight:600"><strong>{r}</strong></span>'
                if hint:
                    label += f' <span style="color:#888;font-style:italic;font-weight:400">: {hint}</span>'
                st.markdown(label, unsafe_allow_html=True)


def render_method(m, cols, i, method_to_primary_resources, cost_sliders_norm):
    """Render the constraint widget group (toggle + cap slider/input) for one method.

    Writes the result into st.session_state.constraints[m].

    Args:
        m (str): CDR method name.
        cols (list): list of Streamlit column objects.
        i (int): index used to pick the column (i % 3).

    Returns:
        None
    """
    with cols[i % 3]:
        tooltip = method_tooltips.get(m, None)

        if m not in st.session_state.constraints:
            st.session_state.constraints[m] = default_constraint()

        constraint = st.session_state.constraints[m]

        toggle_key = f"toggle_{m}"
        cap_type_key = f"cap_type_{m}"

        # Check for missing hidden resources. This runs before the regular
        # st.toggle so we never touch toggle_{m} at all for blocked methods —
        # the pre-render sync already set that key and Streamlit would raise
        # "cannot be modified after widget is instantiated" if we touched it again.
        # Instead we force the constraint to inactive in the data layer and show
        # a static label + info box without any widget that uses toggle_{m}.
        uncovered_required, _ = get_uncovered_hidden_resources(m, cost_sliders_norm)
        if uncovered_required:
            st.session_state.constraints[m]["active"] = False
            st.session_state.pop(toggle_key, None)
            st.toggle(m, value=False, key=toggle_key, disabled=True, help=tooltip)
            _render_resource_status(m, cost_sliders_norm)
            return

        was_active = constraint["active"]
        st.session_state.setdefault(toggle_key, was_active)

        is_active = st.toggle(m, key=toggle_key, help=tooltip)

        st.session_state.constraints[m]["active"] = is_active

        if is_active and not was_active:
            st.session_state.constraints[m]["cap_type"] = "percent"
            st.session_state.constraints[m]["cap_value"] = 100
            st.session_state[cap_type_key] = "percent"

        if not is_active:
            st.session_state.pop(cap_type_key, None)
            return

        # If cap_type is invalid (e.g. "nan" from a corrupted import), show a
        # warning and reset to "percent" so the radio and slider render correctly.
        # The flag is cleared once the user interacts with the radio.
        if constraint.get("_cap_type_invalid") or constraint.get("cap_type") not in ("percent", "absolute"):
            st.warning(
                f"Invalid **cap_type** value (`{constraint.get('cap_type')}`) loaded from file. "
                "Select a constraint type below to fix it."
            )
            constraint["cap_type"] = "percent"
            constraint["cap_value"] = 100
            constraint.pop("_cap_type_invalid", None)
            st.session_state.pop(cap_type_key, None)
            st.session_state.pop(f"slider_{m}", None)
            st.session_state.pop(f"manual_{m}", None)

        # cap_type radio keeps a key so the user's choice persists across
        # reruns; slider and number_input are keyless so they always read
        # directly from the constraint — no stale session-state value possible.
        st.session_state.setdefault(cap_type_key, constraint["cap_type"])

        required, optional = get_hidden_resources_for_method(m, method_to_primary_resources, cost_sliders_norm)
        if required is not None:
            req_list = "\n".join(f"- **{r}** *(required)*" for r in required)
            opt_list = "\n".join(f"- **{r}** *(optional)*" for r in optional) if optional else ""
            body = f"To use this method, add a new resource in the **Resource inputs** tab:\n\n{req_list}"
            if opt_list:
                body += f"\n{opt_list}"
            st.info(body, icon=":material/info:")
            return

        # Capture the cap_type that was active in the previous render so we can
        # detect a mode switch (absolute → percent) below.
        old_cap_type = constraint.get("cap_type", "percent")

        cap_type = st.radio(
            "Cap value type",
            ["percent", "absolute"],
            horizontal=True,
            key=cap_type_key,
            format_func=lambda x: "% of total potential"
            if x == "percent"
            else f"Absolute ({CO2_UNIT}/yr)",
            label_visibility="collapsed",
        )

        st.session_state.constraints[m]["cap_type"] = cap_type

        # When the user switches from absolute → percent the stored cap_value is
        # an absolute number (MtCO₂/yr) which is meaningless as a percentage.
        # Reset both the constraint and the slider key to the default 100 % so
        # neither value nor widget is "stuck" at whatever number was left over.
        if cap_type == "percent" and old_cap_type == "absolute":
            constraint["cap_value"] = 100
            st.session_state[f"slider_{m}"] = 100
            st.session_state.pop(f"manual_{m}", None)

        if cap_type == "absolute" and old_cap_type == "percent":
            constraint["cap_value"] = 0.0
            st.session_state[f"manual_{m}"] = 0.0
            st.session_state.pop(f"slider_{m}", None)

        if cap_type == "percent":
            # Seed the slider key if absent or if it's spuriously 0 while the
            # constraint still records a non-zero intent.  constraint["cap_value"]
            # is only written via on_change callbacks (user-initiated), so it
            # faithfully reflects the user's last explicit choice even across tab
            # navigations — making the discrepancy check reliable.
            _cur = st.session_state.get(f"slider_{m}")
            if _cur is None or (_cur == 0 and int(constraint["cap_value"]) > 0):
                st.session_state[f"slider_{m}"] = int(constraint["cap_value"])

            value = st.slider(
                "Max share (%)", 0, 100, step=5,
                key=f"slider_{m}",
                on_change=_make_slider_cb(m),
            )

            if value == 0:
                st.warning("Cap is 0% but method is active.", icon=":material/warning:")

        else:

            if f"manual_{m}" not in st.session_state:
                st.session_state[f"manual_{m}"] = float(constraint["cap_value"]) if constraint["cap_value"] else 0.0

            value = st.number_input(
                f"Max cap ({CO2_UNIT}/yr)",
                min_value=0.0,
                format="%.2f",
                key=f"manual_{m}",
                placeholder="0.00",
                on_change=_make_manual_cb(m),
            )

            if not value:
                st.warning(f"Cap is 0 {CO2_UNIT}/yr but method is active.", icon=":material/warning:")

def reset_all_constraints():
    """Reset all method constraints to their default values and rerun the app.

    Returns:
        None
    """
    for m in method_order:
        st.session_state.constraints[m] = default_constraint()
        for key in [f"toggle_{m}", f"cap_type_{m}", f"slider_{m}", f"manual_{m}"]:
            st.session_state.pop(key, None)
    st.success("All method constraints reset to default.")
    st.rerun()


def _iter_display_methods(custom_resources):
    """Yield (display_name, original_method_name) for all methods.

    'Other' methods with custom batches are expanded into one entry per batch.
    All other methods yield (m, m).
    """
    for m in method_order:
        variants = get_other_method_variants(m, custom_resources)
        if variants:
            for batch in variants:
                yield make_variant_name(m, batch["name"]), m
        else:
            yield m, m


def build_summary_df():
    """Build a summary DataFrame of currently active methods and their caps.

    Returns:
        pd.DataFrame | None: index = method name, column "Cap value" as a formatted
                             string. Returns None if no method is active.
    """
    custom_resources = st.session_state.get("custom_resources", [])
    rows = {}
    for display_m, _ in _iter_display_methods(custom_resources):
        c = st.session_state.constraints.get(display_m, default_constraint())
        if not c["active"]:
            continue
        if c["cap_type"] == "percent":
            rows[display_m] = {"Cap value": f"{int(c['cap_value'])}%"}
        elif c["cap_type"] == "absolute":
            rows[display_m] = {"Cap value": f"{c['cap_value']:,.2f} {CO2_UNIT}/yr"}
        else:
            rows[display_m] = {"Cap value": f"⚠️ Invalid type ({c['cap_type']!r}) — fix in Tab 2"}
    return pd.DataFrame.from_dict(rows, orient="index") if rows else None



def render_tab(cost_sliders):
    """Render the Method Constraints tab.

    Args:
        cost_sliders (dict): full cost slider config {method: {resource: {min, median, max}}}.

    Returns:
        None
    """
    cost_sliders_norm = {
        m: {strip_unit_suffix(name): conf for name, conf in resources.items()}
        for m, resources in cost_sliders.items()
    }
    method_to_primary_resources = {
        m: [r for r in resources if (m, r) not in secondary_resources]
        for m, resources in cost_sliders_norm.items()
    }

    # Pre-render sync: drive alzzl widget keys from the constraint dict before
    # any widget is drawn.
    #
    # Two modes:
    #   force mode  – after a file/scenario load (constraints_just_loaded flag)
    #                 or for freshly activated methods (just_activated_methods).
    #                 All keys are set unconditionally from the constraint.
    #   normal mode – setdefault only; never overrides a value the user already
    #                 set by interacting with the widget this session.
    just_loaded = st.session_state.pop("constraints_just_loaded", False)
    just_activated = st.session_state.pop("just_activated_methods", set())

    for m, c in st.session_state.constraints.items():
        # Methods with uncovered hidden resources are always forced inactive —
        # skip the normal active-branch so the toggle key is never set to True.
        uncov_req, _ = get_uncovered_hidden_resources(m, cost_sliders_norm)
        if uncov_req:
            st.session_state.constraints[m]["active"] = False
            st.session_state.pop(f"toggle_{m}", None)
            st.session_state.pop(f"cap_type_{m}", None)
            st.session_state.pop(f"slider_{m}", None)
            st.session_state.pop(f"manual_{m}", None)
            continue

        force = just_loaded or (m in just_activated)
        if c["active"]:
            cap_type = c.get("cap_type", "percent")
            if force:
                st.session_state[f"toggle_{m}"] = True
                st.session_state[f"cap_type_{m}"] = cap_type
                if cap_type == "percent":
                    st.session_state[f"slider_{m}"] = int(c.get("cap_value", 100))
                    st.session_state.pop(f"manual_{m}", None)
                else:
                    st.session_state[f"manual_{m}"] = float(c.get("cap_value", 0.0))
                    st.session_state.pop(f"slider_{m}", None)
            else:
                st.session_state.setdefault(f"toggle_{m}", True)
                st.session_state.setdefault(f"cap_type_{m}", cap_type)
                if cap_type == "percent":
                    st.session_state.setdefault(f"slider_{m}", int(c.get("cap_value", 100)))
                    st.session_state.pop(f"manual_{m}", None)
                else:
                    st.session_state.setdefault(f"manual_{m}", float(c.get("cap_value", 0.0)))
                    st.session_state.pop(f"slider_{m}", None)
        else:
            if force:
                st.session_state[f"toggle_{m}"] = False
            else:
                st.session_state.setdefault(f"toggle_{m}", False)
            st.session_state.pop(f"cap_type_{m}", None)
            st.session_state.pop(f"slider_{m}", None)
            st.session_state.pop(f"manual_{m}", None)

    st.markdown('<h1 class="crra-heading">Method constraints</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>
            This section lets you <strong>control which CDR methods are included in or excluded from the optimization</strong> according to the specific context of the studied country.<br>
            For example, you can switch off marine CDR methods for those country with no access to the seas. Some countries may also have legal or tacit bans on specific CDR methods.<br><br>
            For the included methods, you can also set a <strong>maximum deployment cap</strong> per method, by <strong>% of its total potential</strong> or by an <strong>absolute value</strong> in {CO2_UNIT}/yr. Use this parameter in situations where the resources are largely available in the country but other parameters will hinder the deployment, for example low maturity CDR methods, slow deployment of CO2 transport infrastructure, documented low adoption rates...
            </p>
        </div>
        <br>
        """.format(CO2_UNIT=CO2_UNIT),
        unsafe_allow_html=True,
    )

    custom_resources = st.session_state.get("custom_resources", [])

    # Build the full list of display names (variants replace "Other" methods with batches)
    all_display = list(_iter_display_methods(custom_resources))

    any_active = any(
        st.session_state.constraints.get(display_m, default_constraint()).get("active", True)
        for display_m, _ in all_display
    )
    # Activatable: variants are always activatable; standard methods only if they
    # have at least one visible (non-hidden) primary resource AND no uncovered
    # hidden resource that would block them.
    activatable = []
    for display_m, orig_m in all_display:
        if display_m != orig_m:  # it's a variant → batch provides the resource
            activatable.append(display_m)
        elif get_hidden_resources_for_method(orig_m, method_to_primary_resources, cost_sliders_norm)[0] is None:
            uncov_req, _ = get_uncovered_hidden_resources(orig_m, cost_sliders_norm)
            if not uncov_req:
                activatable.append(orig_m)

    if any_active:
        if st.button("Deactivate all", icon=":material/remove_done:", key="btn-deactivate-all"):
            for display_m, _ in all_display:
                st.session_state.constraints.setdefault(display_m, default_constraint())["active"] = False
                st.session_state[f"toggle_{display_m}"] = False
            st.rerun()
    else:
        if st.button("Activate all", icon=":material/done_all:", key="btn-activate-all"):
            for m in activatable:
                st.session_state.constraints.setdefault(m, default_constraint())["active"] = True
                st.session_state[f"toggle_{m}"] = True
            st.rerun()

    with st.expander("Manage methods", expanded=True, icon=":material/settings:"):
        for group, method_list in method_groups.items():
            st.markdown(f"""
                <p class="crra-group-name">{group}</p>""",
                unsafe_allow_html=True,
            )
            cols = st.columns(3)
            col_idx = 0
            for m in method_list:
                variants = get_other_method_variants(m, custom_resources)
                if variants:
                    for batch in variants:
                        vname = make_variant_name(m, batch["name"])
                        render_method(vname, cols, col_idx, method_to_primary_resources, cost_sliders_norm)
                        col_idx += 1
                else:
                    render_method(m, cols, col_idx, method_to_primary_resources, cost_sliders_norm)
                    col_idx += 1

    st.divider()
    if st.button("RESET ALL", key="btn-reset", icon=":material/reset_settings:"):
        reset_all_constraints()

    st.markdown("""
        <p class="crra-sub-heading">Summary of Active methods and caps</p>""",
        unsafe_allow_html=True,
    )
    df = build_summary_df()
    if df is not None:
        df_display = df.reset_index().rename(columns={"index": "CDR Method"})
        df_display = df_display[["CDR Method", "Cap value"]]
        num_rows = len(df_display)
        st.dataframe(
            df_display,
            width='stretch',
            hide_index=True,
            height=(num_rows * 36 + 36),
            column_config={
                "CDR Method": st.column_config.TextColumn("CDR Method", width="large"),
                "Cap value": st.column_config.TextColumn("Cap value"),
            },
            row_height=36
        )
    else:
        st.warning("No active methods selected.")