Blame for utils/data_utils.py

import json
import re
from io import BytesIO
import numpy as np
import pandas as pd
import streamlit as st
from config.config import (
    resource_unit_map, resource_order, method_order,
    resource_groups, method_groups, secondary_resources, CO2_UNIT,
)

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 default_constraint() -> dict:
    """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 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}"


# ---------------------------------------------------------------------------
# Resource / Excel helpers (moved from tab1)

def standard_resource_names() -> set:
    """Return the flat set of all standard resource names defined in resource_groups."""
    return {r for rlist in resource_groups.values() for r in rlist}


def to_excel_bytes(batches: list) -> bytes:
    """Serialise all custom resource batches to an Excel file (one row per batch × method).

    Args:
        batches (list): list of custom batch dicts with keys name, group, amount, unit, methods.

    Returns:
        bytes: raw bytes of the .xlsx workbook.
    """
    rows = []
    for batch in batches:
        for m, coefs in batch["methods"].items():
            rows.append({
                "name":   batch["name"],
                "group":  batch["group"],
                "amount": float(batch["amount"]),
                "unit":   batch["unit"],
                "method": m,
                "min":    float(coefs["min"]),
                "median": float(coefs["median"]),
                "max":    float(coefs["max"]),
            })
    excel_buffer = BytesIO()
    with pd.ExcelWriter(excel_buffer, engine="openpyxl") as writer:
        pd.DataFrame(rows).to_excel(writer, index=False, sheet_name="custom_resources")
    return excel_buffer.getvalue()


def parse_custom_excel(file) -> list:
    """Parse a custom resource Excel file into a list of batch dicts.

    Expects columns: name, group, amount, unit, method, min, median, max.

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

    Returns:
        list[dict]: one dict per unique batch name with keys name, group, amount, unit, methods.

    Raises:
        ValueError: if required columns are missing.
    """
    df = pd.read_excel(file, sheet_name=0)
    required = {"name", "group", "amount", "unit", "method", "min", "median", "max"}
    if not required.issubset(df.columns):
        raise ValueError(f"Missing columns: {required - set(df.columns)}")
    batches = []
    for batch_name in df["name"].unique():
        df_b = df[df["name"] == batch_name]
        methods = {
            row["method"]: {
                "min":    float(row["min"]),
                "median": float(row["median"]),
                "max":    float(row["max"]),
            }
            for _, row in df_b.iterrows()
        }
        group_name = str(df_b["group"].iloc[0])
        resource_to_group = get_resource_to_group()
        batches.append({
            "name":    str(batch_name),
            "group":   group_name,
            "amount":  float(df_b["amount"].iloc[0]),
            "unit":    resource_unit_map.get(resource_to_group.get(group_name), "?"),
            "methods": methods,
        })
    return batches


# ---------------------------------------------------------------------------
# Coefficient / slider helpers (moved from tab3)

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.

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

    Returns:
        tuple[float, str]: (step, fmt).
    """
    min_coeff = conf["min"]
    coeff_range = conf["max"] - min_coeff

    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"

    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: str, resource: str, enabled_methods: dict) -> bool:
    """Check whether a method is 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 enabled (defaults to True when not explicitly set).
    """
    return enabled_methods.get(resource, {}).get(method, True)


def is_custom_secondary(method_name: str, custom_r: str, custom_resource_to_sub_group: dict) -> bool:
    """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


# ---------------------------------------------------------------------------
# JSON / results helpers (moved from tab4)

def to_json_safe(obj):
    """Recursively convert numpy/pandas types to plain Python so json.dumps works."""
    if isinstance(obj, dict):
        return {k: to_json_safe(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [to_json_safe(v) for v in obj]
    if isinstance(obj, pd.Series):
        return obj.iloc[0] if len(obj) == 1 else obj.tolist()
    if isinstance(obj, pd.DataFrame):
        return obj.to_dict(orient="records")
    if isinstance(obj, np.integer):
        return int(obj)
    if isinstance(obj, np.floating):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    return obj


def build_results_df(allocations: dict, method_constraints: dict | None = None) -> pd.DataFrame:
    """Build a DataFrame of per-method CO₂ allocations with share and cap columns.

    Args:
        allocations (dict): {method: CO₂ removed (float)}.
        method_constraints (dict | None): {method: {cap_type, cap_value}} for the Cap column.

    Returns:
        pd.DataFrame: indexed by method, sorted descending by CO₂ removed.
    """
    df = pd.DataFrame.from_dict(allocations, orient="index", columns=[f"{CO2_UNIT} removed"])
    df[f"{CO2_UNIT} removed"] = pd.to_numeric(df[f"{CO2_UNIT} removed"], errors="coerce")
    total = df[f"{CO2_UNIT} removed"].sum()
    if total > 0:
        df["Total share (%)"] = (df[f"{CO2_UNIT} removed"] / total * 100).round(2)

    if method_constraints:
        caps = {}
        for m in df.index:
            c = method_constraints.get(m, {})
            if c.get("cap_type") == "absolute":
                caps[m] = f"{c['cap_value']:,.2f} {CO2_UNIT}/yr"
            elif c.get("cap_type") == "percent":
                caps[m] = f"{int(c.get('cap_value', 100))}% of potential"
            else:
                caps[m] = f"⚠️ Invalid type ({c.get('cap_type')!r}) — fix in Tab 2"
        df["Cap"] = pd.Series(caps)

    return df.sort_values(f"{CO2_UNIT} removed", ascending=False)