tabs/tab3_coefficients.py

import pandas as pd
import streamlit as st

from config.config import (method_groups, method_order, resource_groups,
                           resource_order, secondary_resources, resource_unit_map, CO2_UNIT,
                           hidden_resource_names)
from utils.ui_helpers import format_efficiency_summary, show_bar
from utils.data_utils import strip_unit_suffix, get_resource_to_group


def normalize_costs(costs_dict: dict) -> dict:
    """Strip unit suffixes from resource keys for legacy scenario compatibility."""
    return {strip_unit_suffix(r): v for r, v in costs_dict.items()}


def get_slider_params(conf: dict) -> tuple:
    """Derive slider step size and display format from a coefficient config dict.

    Format is driven by min (display precision); step is driven by range
    (usability) and floored at 1e-06 (Streamlit's hard lower limit).

    Args:
        conf (dict): slider config with keys "min", "median", "max".

    Returns:
        tuple[float, str]: (step, fmt) where step is the slider increment and
                           fmt is a printf-style format string for display.
    """
    min_coeff = conf["min"]
    coeff_range = conf["max"] - min_coeff

    # Format: enough decimals to show the minimum value without scientific notation
    if min_coeff < 1e-10:
        fmt = "%.11f"
    elif min_coeff < 1e-08:
        fmt = "%.9f"
    elif min_coeff < 1e-05:
        fmt = "%.6f"
    elif min_coeff < 0.01:
        fmt = "%.4f"
    elif min_coeff < 1:
        fmt = "%.3f"
    elif min_coeff < 100:
        fmt = "%.2f"
    else:
        fmt = "%.0f"

    # Step: ~1/1000 of the range, floored at 1e-06 (Streamlit minimum)
    if coeff_range < 0.001:
        step = 1e-06
    elif coeff_range < 0.1:
        step = 0.0001
    elif coeff_range < 1:
        step = 0.001
    elif coeff_range < 10:
        step = 0.01
    elif coeff_range < 100:
        step = 0.1
    else:
        step = 1.0

    return step, fmt


def is_method_enabled(method, resource, enabled_methods):
    """Check whether a method is currently enabled for a given resource.

    Args:
        method (str): CDR method name.
        resource (str): resource name.
        enabled_methods (dict): {resource: {method: bool}} from session state.

    Returns:
        bool: True if the method is enabled (defaults to True if not explicitly set).
    """
    return enabled_methods.get(resource, {}).get(method, True)


def is_custom_secondary(method_name, custom_r, custom_resource_to_sub_group):
    """Return True if the reference resource for custom_r is secondary for method_name.

    Args:
        method_name (str): CDR method name.
        custom_r (str): custom batch resource name.
        custom_resource_to_sub_group (dict): {custom_r: reference_resource}.

    Returns:
        bool
    """
    sub_group = custom_resource_to_sub_group.get(custom_r)
    return sub_group is not None and (method_name, sub_group) in secondary_resources


def make_sync_s_to_m(slider_key, manual_key):
    """on_change for slider: mirror its value into the number_input."""
    def sync():
        st.session_state[manual_key] = st.session_state[slider_key]
    return sync


def make_sync_m_to_s(slider_key, manual_key, min_val, max_val):
    """on_change for number_input: clamp to [min, max] and sync both widgets."""
    def sync():
        val = st.session_state[manual_key]
        if val is None:
            return
        clamped = max(min_val, min(max_val, val))
        st.session_state[manual_key] = clamped
        st.session_state[slider_key] = clamped
    return sync


def build_custom_slider_index(cost_sliders_norm, custom_resources):
    """Build the per-method slider configs and group mappings for custom resource batches.

    Also initialises session-state cost entries for custom resources (primary at
    median, secondary at 0) and removes stale entries for methods no longer linked
    to a given batch.

    Args:
        cost_sliders_norm (dict): {method: {resource: {min, median, max}}} with unit suffixes stripped.
        custom_resources (list): list of custom batch dicts from session state.

    Returns:
        tuple[dict, set, dict, dict]:
            custom_sliders_by_method     – {method: {resource: {min, median, max}}}.
            custom_resource_names        – set of all custom batch names.
            custom_resource_to_sub_group – {custom_r: reference_resource}.
            custom_resource_to_group     – {custom_r: physical_group}.
    """
    custom_sliders_by_method = {}
    custom_resource_names = set()
    for batch in custom_resources:
        r_name = batch["name"]
        custom_resource_names.add(r_name)
        for m, coefs in batch["methods"].items():
            custom_sliders_by_method.setdefault(m, {})[r_name] = coefs
        ref = batch["group"]
        for m, resources in cost_sliders_norm.items():
            if (m, ref) in secondary_resources and r_name not in custom_sliders_by_method.get(m, {}):
                ref_coefs = resources.get(ref)
                if ref_coefs:
                    custom_sliders_by_method.setdefault(m, {})[r_name] = ref_coefs

    custom_resource_to_sub_group = {batch["name"]: batch["group"] for batch in custom_resources}
    resource_to_group = get_resource_to_group()
    custom_resource_to_group = {
        r: resource_to_group[s_g]
        for r, s_g in custom_resource_to_sub_group.items()
        if s_g in resource_to_group
    }

    # Remove stale cost entries for methods no longer linked to a batch
    valid_methods_for_custom = {}
    for m, resources in custom_sliders_by_method.items():
        for r in resources:
            valid_methods_for_custom.setdefault(r, set()).add(m)
    for r_name, valid in valid_methods_for_custom.items():
        for m in list(st.session_state.get("costs", {}).keys()):
            if r_name in st.session_state.costs.get(m, {}) and m not in valid:
                del st.session_state.costs[m][r_name]

    # Pre-initialise cost entries for custom resources
    for m, resources in custom_sliders_by_method.items():
        st.session_state.costs.setdefault(m, {})
        for r, conf in resources.items():
            if r not in st.session_state.costs[m]:
                st.session_state.costs[m][r] = (
                    0.0 if is_custom_secondary(m, r, custom_resource_to_sub_group) else conf["median"]
                )

    return custom_sliders_by_method, custom_resource_names, custom_resource_to_sub_group, custom_resource_to_group


def render_standard_resource_row(method, r, conf, display_label, enabled_methods, has_custom_batch=False):
    """Render the slider row for a single standard resource coefficient.

    Handles the secondary-resource toggle, fixed-value display, and the
    slider + manual number_input pair. Writes the result to
    st.session_state.costs[method][r].

    Args:
        method (str): CDR method name.
        r (str): resource name.
        conf (dict): {min, median, max} for this (method, resource) pair.
        display_label (str): human-readable resource label with unit.
        enabled_methods (dict): {resource: {method: bool}} from session state.
        has_custom_batch (bool): True if a custom batch derived from r exists.

    Returns:
        None
    """
    if not is_method_enabled(method, r, enabled_methods):
        st.info(f"INFO: Method disabled for **{r}** (see Resource Inputs tab)")
        return

    is_secondary = (method, r) in secondary_resources
    step, fmt = get_slider_params(conf)
    toggle_key = f"toggle_{method}_{r}".replace(" ", "_")
    slider_key = f"slider_{method}_{r}".replace(" ", "_")
    manual_key = f"manual_{method}_{r}".replace(" ", "_")
    current_val = normalize_costs(st.session_state.costs.get(method, {})).get(r, 0.0)
    if is_secondary:
        if toggle_key not in st.session_state:
            st.session_state[toggle_key] = current_val != 0.0
        if not st.checkbox(f"Use {display_label}", key=toggle_key):
            st.session_state.costs[method][r] = 0.0
            st.markdown("**DISABLED**")
            return
        elif current_val == 0.0:
            st.session_state.costs[method][r] = conf["median"]
            current_val = conf["median"]

    if conf["min"] == conf["max"]:
        st.markdown(
            f'<p class="crra-resource-name">{display_label} = <strong>{conf["median"]}</strong> (fixed)</p>',
            unsafe_allow_html=True,
        )
        st.session_state.costs[method][r] = conf["median"]
        if st.session_state.get("resource_caps", {}).get(r, 0) <= 0:
            if has_custom_batch:
                st.warning(
                    f"No available quantity set for **{r}** in Resource Inputs tab.",
                    icon=":material/warning:",
                )
            else:
                st.warning(
                    f"No available quantity set for **{r}** in Resource Inputs tab: "
                    f"**{method}** won't be used in the optimization.",
                    icon=":material/warning:",
                )
        return

    if slider_key not in st.session_state:
        st.session_state[slider_key] = current_val
    if manual_key not in st.session_state:
        st.session_state[manual_key] = st.session_state[slider_key]
    if st.session_state.get("resource_caps", {}).get(r, 0) <= 0:
        if has_custom_batch:
            st.warning(
                f"No available quantity set for **{r}** in Resource Inputs tab.",
                icon=":material/warning:",
            )
        else:
            st.warning(
                f"No available quantity set for **{r}** in Resource Inputs tab: "
                f"**{method}** won't be used in the optimization.",
                icon=":material/warning:",
            )

    col1, col2 = st.columns([3, 1])
    with col1:
        slider_val = st.slider(
            label=display_label,
            min_value=conf["min"], max_value=conf["max"], step=step, format=fmt,
            key=slider_key,
            on_change=make_sync_s_to_m(slider_key, manual_key)
        )
    with col2:
        st.number_input(
            label="label", label_visibility="collapsed",
            key=manual_key, help="Manual override (optional, any value)",
            format=fmt, value=None, placeholder=str(round(slider_val, 6)),
            on_change=make_sync_m_to_s(slider_key, manual_key, conf["min"], conf["max"]),
        )

    manual_value = st.session_state.get(manual_key)
    st.session_state.costs[method][r] = slider_val if manual_value is None else manual_value


def render_custom_resource_row(method, r, conf, enabled_methods, custom_resource_to_sub_group, custom_resources):
    """Render the slider row for a single custom resource batch coefficient.

    Args:
        method (str): CDR method name.
        r (str): custom resource batch name.
        conf (dict): {min, median, max} for this (method, resource) pair.
        enabled_methods (dict): {resource: {method: bool}} from session state.
        custom_resource_to_sub_group (dict): {custom_r: reference_resource}.
        custom_resources (list): list of custom batch dicts from session state.

    Returns:
        None
    """
    if not is_method_enabled(method, r, enabled_methods):
        st.info("INFO: Method disabled for this resource (see Resource Inputs tab)")
        return

    is_sec = is_custom_secondary(method, r, custom_resource_to_sub_group)
    step, fmt = get_slider_params(conf)
    toggle_key = f"toggle_{method}_{r}".replace(" ", "_")
    slider_key = f"slider_{method}_{r}".replace(" ", "_")
    manual_key = f"manual_{method}_{r}".replace(" ", "_")
    current_val = st.session_state.costs.get(method, {}).get(r, 0.0 if is_sec else conf["median"])

    if is_sec:
        if toggle_key not in st.session_state:
            st.session_state[toggle_key] = current_val != 0.0
        if not st.checkbox(f"Use {r}", key=toggle_key):
            st.session_state.costs.setdefault(method, {})[r] = 0.0
            st.markdown("**DISABLED**")
            return
        elif current_val == 0.0:
            st.session_state.costs.setdefault(method, {})[r] = conf["median"]
            current_val = conf["median"]

    if conf["min"] == conf["max"]:
        st.markdown(
            f'<p class="crra-resource-name">{r} = <strong>{conf["median"]}</strong> (fixed)</p>',
            unsafe_allow_html=True,
        )
        st.session_state.costs.setdefault(method, {})[r] = conf["median"]
        return

    if slider_key not in st.session_state:
        st.session_state[slider_key] = current_val
    if manual_key not in st.session_state:
        st.session_state[manual_key] = st.session_state[slider_key]

    batch_amount = next((float(b.get("amount", 0)) for b in custom_resources if b["name"] == r), 0.0)
    if batch_amount <= 0:
        if is_sec:
            st.info(
                f"No available quantity set for **{r}** in Resource Inputs.",
                icon=":material/info:",
            )
        else:
            st.info(
                f"No available quantity set for **{r}** in Resource Inputs — "
                "this method won't be used in the optimization.",
                icon=":material/info:",
            )

    col1, col2 = st.columns([3, 1])
    with col1:
        slider_val = st.slider(
            label=r,
            min_value=conf["min"], max_value=conf["max"], step=step, format=fmt,
            key=slider_key,
            on_change=make_sync_s_to_m(slider_key, manual_key),
        )
    with col2:
        st.number_input(
            label="label", label_visibility="collapsed",
            key=manual_key, help="Manual override (optional, any value)",
            format=fmt, value=None, placeholder=str(round(slider_val, 6)),
            on_change=make_sync_m_to_s(slider_key, manual_key, conf["min"], conf["max"]),
        )

    st.session_state.costs.setdefault(method, {})
    manual_value = st.session_state.get(manual_key)
    st.session_state.costs[method][r] = slider_val if manual_value is None else manual_value


def panel_method_editor(
    method, cost_sliders_norm, custom_sliders_by_method,
    resource_display, enabled_methods, custom_resources,
    custom_resource_to_group, custom_resource_to_sub_group,
):
    """Render the coefficient-editing expander for one CDR method.

    Shows one slider row per standard resource, then one per custom resource batch,
    grouped by resource group. Includes a per-method reset button.

    Args:
        method (str): CDR method name.
        cost_sliders_norm (dict): {method: {resource: {min, median, max}}}.
        custom_sliders_by_method (dict): {method: {resource: {min, median, max}}}.
        resource_display (dict): {resource: label_with_unit}.
        enabled_methods (dict): {resource: {method: bool}} from session state.
        custom_resources (list): list of custom batch dicts from session state.
        custom_resource_to_group (dict): {custom_r: physical_group}.
        custom_resource_to_sub_group (dict): {custom_r: reference_resource}.

    Returns:
        None
    """
    with st.expander(f"Manage resources for **{method}**", expanded=True):
        method_sliders = cost_sliders_norm.get(method, {})
        custom_for_method = custom_sliders_by_method.get(method, {})

        primary_resources = [r for r in method_sliders if (method, r) not in secondary_resources]
        if primary_resources and all(r in hidden_resource_names for r in primary_resources) and not custom_for_method:
            required = [r for r in primary_resources if r in hidden_resource_names]
            optional = [r for r in method_sliders if r in hidden_resource_names and (method, r) in secondary_resources]
            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

        for group_name, rlist in resource_groups.items():
            used = [r for r in rlist if r in method_sliders and r not in hidden_resource_names]
            custom_in_group = [
                r for r in custom_for_method
                if custom_resource_to_group.get(r) == group_name
            ]
            hidden_secondary_in_group = [
                r for r in rlist
                if r in method_sliders
                and r in hidden_resource_names
                and (method, r) in secondary_resources
                and not any(custom_resource_to_sub_group.get(cr) == r for cr in custom_for_method)
            ]
            if not used and not custom_in_group and not hidden_secondary_in_group:
                continue

            st.markdown(f'<p class="crra-group-name">{group_name}</p>', unsafe_allow_html=True)

            resources_with_batches = set(custom_resource_to_sub_group.values())
            for r in used:
                conf = method_sliders[r]
                render_standard_resource_row(
                    method, r, conf, resource_display.get(r, r), enabled_methods,
                    has_custom_batch=(r in resources_with_batches),
                )

            for r in custom_in_group:
                render_custom_resource_row(
                    method, r, custom_for_method[r],
                    enabled_methods, custom_resource_to_sub_group, custom_resources
                )

            for r in hidden_secondary_in_group:
                st.info(
                    f"**{r}** can be used as an optional secondary resource for this method. "
                    f"Add a custom resource of type **{r}** in the **Resource inputs** tab to activate it.",
                    icon=":material/add_circle:",
                )

        if st.button("Reset selected method", key="btn-reset-method", icon=":material/reset_settings:"):
            for r, conf in cost_sliders_norm.get(method, {}).items():
                is_secondary = (method, r) in secondary_resources
                st.session_state.costs[method][r] = 0.0 if is_secondary else conf["median"]
                for prefix in ["slider", "manual", "toggle"]:
                    st.session_state.pop(f"{prefix}_{method}_{r}".replace(" ", "_"), None)
            for r, conf in custom_sliders_by_method.get(method, {}).items():
                st.session_state.costs.setdefault(method, {})[r] = conf["median"]
                for prefix in ["slider", "manual"]:
                    st.session_state.pop(f"{prefix}_{method}_{r}".replace(" ", "_"), None)
            st.success(f"{method} reset to default values.")
            st.rerun()


def panel_summary_table(
    cost_sliders_norm, custom_sliders_by_method,
    enabled_methods, custom_resource_names, custom_resource_to_sub_group,
):
    """Render the full coefficient summary table with cells highlighted when modified.

    Args:
        cost_sliders_norm (dict): {method: {resource: {min, median, max}}}.
        custom_sliders_by_method (dict): {method: {resource: {min, median, max}}}.
        enabled_methods (dict): {resource: {method: bool}} from session state.
        custom_resource_names (set[str]): names of all custom resource batches.
        custom_resource_to_sub_group (dict): {custom_r: reference_resource}.

    Returns:
        None
    """
    st.markdown(
        '<p class="crra-sub-heading">Summary of resource usage coefficients</p>',
        unsafe_allow_html=True,
    )
    resource_display = {
        r: f"{r} ({resource_unit_map[cat]}/{CO2_UNIT})"
        for cat, resources in resource_groups.items()
        for r in resources
    }
    all_resource_columns = [r for r in resource_order if r not in hidden_resource_names] + sorted(custom_resource_names)

    df_default = pd.DataFrame(index=method_order, columns=all_resource_columns)
    df_current = pd.DataFrame(index=method_order, columns=all_resource_columns)

    for m in method_order:
        current_costs = normalize_costs(st.session_state.costs.get(m, {}))
        sliders = dict(cost_sliders_norm.get(m, {}))
        sliders.update(custom_sliders_by_method.get(m, {}))

        for r in all_resource_columns:
            is_sec = (m, r) in secondary_resources or is_custom_secondary(m, r, custom_resource_to_sub_group)
            disabled = not is_method_enabled(m, r, enabled_methods)
            current_val = 0.0 if disabled else current_costs.get(r, 0.0)
            default_val = 0.0 if disabled else sliders.get(r, {}).get("median", 0.0)
            if is_sec and current_val == 0.0:
                default_val = 0.0
            df_default.at[m, r] = default_val
            df_current.at[m, r] = current_val

    df_default = df_default.astype(float).fillna(0.0)
    df_current = df_current.astype(float).fillna(0.0)

    resource_display_full = dict(resource_display)
    resource_display_full.update({r: r for r in custom_resource_names})
    df_default = df_default.rename(columns=resource_display_full)
    df_current = df_current.rename(columns=resource_display_full)

    styled = format_efficiency_summary(df_current, df_default)
    st.dataframe(
        data=styled,
        height=(len(df_current) * 36 + 36),
        width='stretch',
        column_config={"_index": st.column_config.TextColumn("Method", width="large")},
        row_height=36
    )


def render_tab(cost_sliders):
    """Render the Resource Efficiency Coefficients tab.

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

    Returns:
        None
    """
    st.markdown('<h1 class="crra-heading">Resource efficiency coefficients</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 adjust the resource usage coefficients for each CDR method "
        "to best fit your country conditions.<br>"
        "<li>Primary resources are always active and initialized with default values.</li>"
        "<li>Secondary resources are initialized at zero and toggled off by default.</li>"
        "<li>Toggling on a secondary resource restores its default value.</li>"
        "</p></div>",
        unsafe_allow_html=True,
    )
    with st.container(key="warning_resources"):
        st.info(
            "Red cells in the table highlight coefficients modified from their **default median values**",
            icon=":material/info:",
        )

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

    cost_sliders_norm = {m: normalize_costs(resources) for m, resources in cost_sliders.items()}
    resource_display = {
        r: f"{r} ({resource_unit_map[cat]}/{CO2_UNIT})"
        for cat, resources in resource_groups.items()
        for r in resources
    }

    (
        custom_sliders_by_method,
        custom_resource_names,
        custom_resource_to_sub_group,
        custom_resource_to_group,
    ) = build_custom_slider_index(cost_sliders_norm, custom_resources)

    constraints = st.session_state.get("constraints", {})
    active_methods = {m for m, c in constraints.items() if c.get("active", True)}
    method_choices = [
        f"{g} | {m}"
        for g, methods in method_groups.items()
        for m in methods
        if not constraints or m in active_methods
    ]
    if not method_choices:
        st.warning("No active methods. Enable at least one method in the **Method Constraints** tab.")
        return
    method_lookup  = {label: label.split(" | ", 1)[1] for label in method_choices}
    selected_label = st.selectbox("Select a method:", method_choices, key="crra-select")

    if st.button("RESET ALL", key="btn-reset", icon=":material/reset_settings:"):
        for m, sliders in cost_sliders_norm.items():
            for r, conf in sliders.items():
                st.session_state.costs[m][r] = 0.0 if (m, r) in secondary_resources else conf["median"]
                for prefix in ["slider", "manual", "toggle"]:
                    st.session_state.pop(f"{prefix}_{m}_{r}".replace(" ", "_"), None)
        for m, sliders in custom_sliders_by_method.items():
            for r, conf in sliders.items():
                st.session_state.costs.setdefault(m, {})[r] = (
                    0.0 if is_custom_secondary(m, r, custom_resource_to_sub_group) else conf["median"]
                )
                for prefix in ["slider", "manual", "toggle"]:
                    st.session_state.pop(f"{prefix}_{m}_{r}".replace(" ", "_"), None)
        st.success("All methods reset to default values.")
        st.rerun()

    panel_method_editor(
        method=method_lookup[selected_label],
        cost_sliders_norm=cost_sliders_norm,
        custom_sliders_by_method=custom_sliders_by_method,
        resource_display=resource_display,
        enabled_methods=enabled_methods,
        custom_resources=custom_resources,
        custom_resource_to_group=custom_resource_to_group,
        custom_resource_to_sub_group=custom_resource_to_sub_group,
    )

    panel_summary_table(
        cost_sliders_norm=cost_sliders_norm,
        custom_sliders_by_method=custom_sliders_by_method,
        enabled_methods=enabled_methods,
        custom_resource_names=custom_resource_names,
        custom_resource_to_sub_group=custom_resource_to_sub_group,
    )