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,
    hidden_resource_names,
)
from optimisation.optimisation import run_optimisation

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}"


def run_expanded_optimisation(resource_caps, method_constraints, costs, custom_resources, enabled_methods):
    """Apply the full Tab 4 preprocessing pipeline, then run the LP.

    This is the same pipeline as tab4_optimisation.run_and_store_result: apply
    enabled_methods overrides, inject custom-batch coefficients, neutralise
    all-zero methods, deactivate methods missing a required resource, and —
    crucially — expand any "Other"-tagged method into one independent method
    variant per assigned custom batch (each variant gets the batch as its own
    standard-style resource, with its own cap, instead of going through the
    z-variable pooling/substitution mechanism).

    Any caller that needs LP results consistent with what Tab 4 actually shows
    (e.g. Tab 5's sensitivity sweeps) should use this instead of calling
    run_optimisation() directly with raw session-state costs — otherwise a
    custom resource attached to an "Other" method would be silently computed
    through the wrong mechanism and never show up as its own entry.

    Args:
        resource_caps (dict): {resource: available amount} — may be a perturbed
            copy of st.session_state.resource_caps.
        method_constraints (dict): {method: {active, cap_type, cap_value}} — may
            be a perturbed copy of st.session_state.constraints.
        costs (dict): {method: {resource: coefficient}}, raw/unexpanded —
            normally st.session_state.costs.
        custom_resources (list): custom batch dicts.
        enabled_methods (dict): {resource: {method: bool}}.

    Returns:
        tuple[bool, dict]: same shape as optimisation.run_optimisation()'s
            return value — (True, result_dict) or (False, {"message": str}).
    """
    applied_costs = {m: dict(v) for m, v in costs.items()}
    for method, method_resources in applied_costs.items():
        for resource in list(method_resources):
            if not enabled_methods.get(resource, {}).get(method, True):
                applied_costs[method][resource] = 0.0

    # Inject custom batch coefficients so neutralize does not wrongly deactivate
    # a method that only has custom-batch access.
    for batch in custom_resources:
        batch_name = batch["name"]
        for method_name, coefficients in batch.get("methods", {}).items():
            if method_name in applied_costs and applied_costs[method_name].get(batch_name, 0.0) == 0.0:
                applied_costs[method_name][batch_name] = coefficients.get("median", 0.0)

    standard_constraints = {
        m: c for m, c in method_constraints.items()
        if m in applied_costs
    }
    safe_constraints = neutralize_zero_coefficient_methods(applied_costs, standard_constraints)

    # Deactivate any method that requires a resource not provided (standard cap
    # nor custom batch covers it).
    custom_batch_names_pre = {b["name"] for b in custom_resources}
    available_resources = set(resource_caps) | custom_batch_names_pre
    for m, m_costs in applied_costs.items():
        if not safe_constraints.get(m, {}).get("active", True):
            continue
        missing = [
            r for r, c in m_costs.items()
            if c > 0 and r not in available_resources
        ]
        if missing:
            safe_constraints.setdefault(m, default_constraint())
            safe_constraints[m] = {**safe_constraints[m], "active": False}

    # Expand "Other" methods: replace each with one variant per assigned batch.
    resource_caps_for_lp = dict(resource_caps)
    expanded_batch_names: set[str] = set()

    for m in list(applied_costs.keys()):
        if not is_other_method(m):
            continue
        batches_for_m = get_other_method_variants(m, custom_resources)
        if not batches_for_m:
            continue
        all_batch_names_for_m = {b["name"] for b in batches_for_m}

        for batch in batches_for_m:
            vname = make_variant_name(m, batch["name"])
            variant_costs = {
                r: coeff for r, coeff in applied_costs[m].items()
                if r not in hidden_resource_names and r not in all_batch_names_for_m
            }
            variant_costs[batch["name"]] = applied_costs[m].get(
                batch["name"],
                batch["methods"][m]["median"],
            )
            applied_costs[vname] = variant_costs
            safe_constraints[vname] = (
                method_constraints.get(vname) or default_constraint()
            )
            resource_caps_for_lp[batch["name"]] = float(batch["amount"])
            expanded_batch_names.add(batch["name"])
        del applied_costs[m]
        del safe_constraints[m]

    custom_resources_for_lp = [b for b in custom_resources if b["name"] not in expanded_batch_names]

    # Non-"Other" methods that share an expanded batch lose their z-variable
    # once the batch is removed from custom_resources_for_lp — inject it as a
    # standard resource coefficient so the constraint still applies to them.
    for batch in custom_resources:
        if batch["name"] not in expanded_batch_names:
            continue
        batch_group = batch.get("group", "")
        group_cap_is_zero = resource_caps_for_lp.get(batch_group, 0.0) == 0.0
        for method_name, coeff_dict in batch.get("methods", {}).items():
            if method_name not in applied_costs:
                continue
            applied_costs[method_name][batch["name"]] = coeff_dict.get("median", 0.0)
            if group_cap_is_zero and batch_group in applied_costs[method_name]:
                applied_costs[method_name][batch_group] = 0.0

    return run_optimisation(
        resource_caps=resource_caps_for_lp,
        method_constraints=safe_constraints,
        method_costs=applied_costs,
        custom_resources=custom_resources_for_lp,
        enabled_methods=enabled_methods,
        resource_to_group=get_resource_to_group(),
    )


# ---------------------------------------------------------------------------
# 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 resource_caps_to_excel_bytes(resource_caps: dict) -> bytes:
    """Serialise the current standard resource availabilities to an Excel file.

    Same "Resource" / "Unit" / "Available Amount" columns as the blank template
    (data/resource_template.xlsx), so the exported file can be re-uploaded as-is.

    Args:
        resource_caps (dict): {resource: available amount}.

    Returns:
        bytes: raw bytes of the .xlsx workbook.
    """
    resource_to_unit = get_resource_to_unit()
    rows = [
        {
            "Resource": r,
            "Unit (per year)": resource_to_unit.get(r, ""),
            "Available Amount": float(amount),
        }
        for r, amount in resource_caps.items()
        if r in resource_order
    ]
    excel_buffer = BytesIO()
    with pd.ExcelWriter(excel_buffer, engine="openpyxl") as writer:
        pd.DataFrame(rows).to_excel(writer, index=False, sheet_name="resource_caps")
    return excel_buffer.getvalue()


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']:,g} {CO2_UNIT}/yr"
            elif c.get("cap_type") == "percent":
                caps[m] = f"{int(c.get('cap_value', 100))}% of its 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)