utils/data_utils.py

import io
import json
import re
import pandas as pd
import streamlit as st
import openpyxl
from openpyxl.styles import PatternFill, Font
from config.config import resource_unit_map, resource_order, method_order, resource_groups, method_groups

def get_resource_to_group() -> dict[str, str]:
    """Map each resource name to its resource group.

    Returns:
        dict[str, str]: resource name → group name (e.g. "Arable land" → "Land").
    """
    return {r: g for g, resources in resource_groups.items() for r in resources}

def get_resource_to_unit() -> dict[str, str]:
    """Map each resource name to its physical unit.

    Returns:
        dict[str, str]: resource name → unit (e.g. "Arable land" → "Mha").
    """
    return {r: resource_unit_map[g] for g, resources in resource_groups.items() for r in resources}

def get_method_group(method: str) -> str:
    """Return the group a CDR method belongs to.

    Args:
        method (str): CDR method name.

    Returns:
        str: group name, or "Other" if not found.
    """
    return next((g for g, methods in method_groups.items() if method in methods), "Other")

@st.cache_data
def load_slider_data():
    """Load and cache the cost slider configuration from data/cost_sliders.json.

    Returns:
        dict: nested dict {method: {resource: {min, median, max}}}.
    """
    with open("data/cost_sliders.json") as f:
        data = json.load(f)
        return data
    
def render_resource_input(name_clean, r, tooltip=None):
    """Render a number_input widget for a single resource availability cap.

    Writes the entered value into st.session_state.resource_caps[r].

    Args:
        name_clean (str): display label (resource name without unit suffix).
        r (str): canonical resource name used as the session state key.
        tooltip (str | None): optional help text shown on hover.

    Returns:
        None
    """
    widget_key = f"cap_{r}"

    if widget_key not in st.session_state:
        existing = st.session_state.resource_caps.get(r, 0.0)
        st.session_state[widget_key] = existing if existing else None

    cap = st.number_input(
        label=name_clean,
        min_value=0.0,
        format="%.2f",
        value=None,
        placeholder="0.00",
        key=widget_key,
        help=tooltip
    )

    st.session_state.resource_caps[r] = cap if cap is not None else 0.0

def is_valid_scenario(scenario):
    """Check that a scenario dict has the minimum required keys to be loaded.

    Args:
        scenario (dict): scenario data to validate.

    Returns:
        bool: True if all base keys are present, False otherwise.
    """
    base_ok = (
        isinstance(scenario, dict)
        and "allocations" in scenario
        and "resource_caps" in scenario
        and "method_costs" in scenario
        and "method_constraints" in scenario
        and "summary" in scenario
    )
    if not base_ok:
        return False
    return True

def is_complete_scenario(scenario):
    """Check that a scenario carries full optimisation outputs, not just inputs.

    A complete scenario can be displayed without re-running the optimisation.
    An incomplete scenario (e.g. exported before resource_actual_gain was added)
    triggers an automatic re-run on load.

    Args:
        scenario (dict): scenario data to validate.

    Returns:
        bool: True if the scenario includes all detailed output keys.
    """
    required_extra = {
        "resource_usage", "resource_used_total", "resource_remaining",
        "resource_actual_gain", "enabled_methods",
    }
    return is_valid_scenario(scenario) and required_extra.issubset(scenario.keys())

def load_inputs_from_excel(file):
    """Parse an Excel workbook exported from the tool into session-state-ready dicts.

    Expects sheets: resource_caps, constraints, coefficients.
    Optional sheets: custom_resources, enabled_methods.

    Args:
        file: file-like object pointing to an .xlsx workbook.

    Returns:
        tuple: (resource_caps, constraints, costs, custom_resources, enabled_methods, warnings)
               where warnings is a list[str] of non-fatal import issues.
    """
    xl = pd.read_excel(file, sheet_name=None)

    resource_caps = {}
    if "resource_caps" in xl:
        for _, row in xl["resource_caps"].iterrows():
            resource_caps[row["Resource"]] = float(row["Available Amount"])

    constraints = {}
    if "constraints" in xl:
        for _, row in xl["constraints"].iterrows():
                constraints[row["method"]] = {
                    "active": bool(row["active"]),
                    "cap_type": str(row["cap_type (percent or absolute)"]),
                    "cap_value": float(row["cap_value"])
                }
    costs = {}
    if "coefficients" in xl:
        for _, row in xl["coefficients"].iterrows():
            m, r, v = row["method"], row["resource"], float(row["value"])
            costs.setdefault(m, {})[r] = v

    custom_resources = []
    warnings = []
    valid_resources = set(resource_order)
    valid_methods = set(method_order)
    if "custom_resources" in xl:
        df_custom = xl["custom_resources"]
        required_custom = {"name", "group", "amount", "unit", "method", "min", "median", "max"}
        if required_custom.issubset(df_custom.columns):
            for bname in df_custom["name"].dropna().unique():
                df_b = df_custom[df_custom["name"] == bname]
                group = str(df_b["group"].iloc[0])
                if group not in valid_resources:
                    warnings.append(f"New resource \"{bname}\" - group \"{group}\" not found: check spelling.")
                    continue
                methods = {}
                for _, row in df_b.iterrows():
                    m = str(row["method"])
                    if m not in valid_methods:
                        warnings.append(f"New resource \"{bname}\" - method \"{m}\" not found: check spelling.")
                        continue
                    methods[m] = {
                        "min": float(row["min"]),
                        "median": float(row["median"]),
                        "max": float(row["max"]),
                    }
                custom_resources.append({
                    "name": str(bname),
                    "group": group,
                    "amount": float(df_b["amount"].iloc[0]),
                    "unit": str(df_b["unit"].iloc[0]),
                    "methods": methods,
                })

    # Per-resource manual method allocation (which CDR methods are allowed to
    # use a given resource). Optional sheet for backward compatibility with
    # older templates that predate this feature.
    enabled_methods = {}
    if "enabled_methods" in xl:
        df_enabled = xl["enabled_methods"]
        required_enabled = {"resource", "method", "enabled"}
        if required_enabled.issubset(df_enabled.columns):
            for _, row in df_enabled.iterrows():
                r, m = row["resource"], row["method"]
                enabled_methods.setdefault(r, {})[m] = bool(row["enabled"])

    return resource_caps, constraints, costs, custom_resources, enabled_methods, warnings

def strip_unit_suffix(resource_name: str) -> str:
    """Remove trailing unit suffix from a resource name.

    E.g. "Other land (Mha/MtCO₂)" → "Other land".

    Args:
        resource_name (str): raw resource name, possibly with a parenthetical unit.

    Returns:
        str: resource name without the unit suffix.
    """
    return re.sub(r"\s*\(.*?\)\s*$", "", resource_name).strip()

def neutralize_zero_coefficient_methods(method_costs: dict, method_constraints: dict) -> dict:
    """Return a copy of method_constraints with methods that have all-zero costs deactivated.

    Prevents the LP from including methods with no usable resource coefficients.

    Args:
        method_costs (dict): {method: {resource: coefficient}} cost matrix.
        method_constraints (dict): {method: {active, cap_type, cap_value}}.

    Returns:
        dict: modified copy of method_constraints with zero-cost methods deactivated.
    """
    safe_constraints = {m: c.copy() for m, c in method_constraints.items()}

    for method, constraint in safe_constraints.items():
        if not constraint.get("active", True):
            continue
        coeffs = method_costs.get(method, {})
        all_zero = all(v == 0 for v in coeffs.values()) if coeffs else True
        if all_zero:
            constraint["active"] = False

    return safe_constraints


# ---------------------------------------------------------------------------
# "Other" method expansion helpers
# ---------------------------------------------------------------------------

def is_other_method(method_name: str) -> bool:
    """True when the method name contains 'other' (case-insensitive).

    These methods accept any resource from a generic 'Other X' pool and are
    expanded into one variant per custom batch assigned by the user.
    """
    return "other" in method_name.lower()


def get_other_method_variants(method_name: str, custom_resources: list) -> list:
    """Return the custom batch dicts assigned to an 'Other' method.

    Returns an empty list when the method is not an 'Other' method or has no
    custom batches assigned to it.
    """
    if not is_other_method(method_name):
        return []
    return [b for b in custom_resources if method_name in b.get("methods", {})]


def make_variant_name(method_name: str, batch_name: str) -> str:
    """Build the display/LP name for a method variant: '<Method>: <Batch>'."""
    return f"{method_name}: {batch_name}"