optimisation/optimisation.py

import numpy as np
from scipy.optimize import linprog


def apply_enabled_methods(method_costs, enabled_methods):
    """Return a copy of method_costs with coefficients zeroed for disabled (method, resource) pairs.

    Args:
        method_costs (dict): {method: {resource: coefficient}}.
        enabled_methods (dict): {resource: {method: bool}} — False entries disable the coefficient.

    Returns:
        dict: deep copy of method_costs with disabled coefficients set to 0.0.
    """
    result = {m: dict(v) for m, v in method_costs.items()}
    enabled_methods = enabled_methods or {}
    for method, resources in result.items():
        for resource in list(resources):
            if not enabled_methods.get(resource, {}).get(method, True):
                result[method][resource] = 0.0
    return result


def build_custom_batch_structures(method_costs, resource_caps, custom_resources):
    """Build the custom-batch index structures and inject fallback coefficients.

    Mutates method_costs (adds fallback median coefficients) and resource_caps
    (adds batch caps that are not yet present). Both dicts are already copies
    when this function is called from run_optimisation.

    Args:
        method_costs (dict): {method: {resource: coefficient}} — mutated in place.
        resource_caps (dict): {resource: cap} — mutated in place.
        custom_resources (list | None): list of custom batch dicts with keys
                                        name, group, amount, methods.

    Returns:
        tuple[dict, dict, set]:
            custom_batch_specs       – {batch_name: {group, amount, primary_methods}}.
            group_to_custom_batches  – {reference_resource: [batch_name, ...]}.
            custom_batch_names       – set of all batch names.
    """
    custom_batch_specs = {}
    for batch in (custom_resources or []):
        custom_batch_specs[batch["name"]] = {
            "group":           batch["group"],
            "amount":          float(batch["amount"]),
            "primary_methods": set(batch.get("methods", {}).keys()),
        }

    custom_batch_names = set(custom_batch_specs)
    group_to_custom_batches = {}
    for batch_name, specs in custom_batch_specs.items():
        group_to_custom_batches.setdefault(specs["group"], []).append(batch_name)
        if batch_name not in resource_caps:
            resource_caps[batch_name] = specs["amount"]

    # Fallback: use batch median coefficient if method has none
    for batch in (custom_resources or []):
        batch_name = batch["name"]
        for method, coefs in batch["methods"].items():
            if method_costs.get(method, {}).get(batch_name, 0.0) == 0.0:
                method_costs.setdefault(method, {})[batch_name] = coefs["median"]

    return custom_batch_specs, group_to_custom_batches, custom_batch_names


def build_lp_indices(active_methods, num_methods, custom_batch_specs, resource_caps, method_costs, pool_of_entity=None):
    """Build the LP variable index for custom-batch draw variables z_{m,n}.

    When the user adds a custom resource batch (e.g. "degraded land patch A"),
    the solver needs to track separately how much of each CDR method's output
    comes from that specific batch vs. the standard resource pool.
    These "draw variables" z_{m,n} represent:
        z_{m,n} = amount of CO₂ method m produces using custom batch n.
    They sit alongside the main y_m variables in the LP column vector.

    Args:
        active_methods (list[str]): ordered list of active CDR method names.
        num_methods (int): len(active_methods) — index offset for z variables.
        custom_batch_specs (dict): {batch_name: {group, amount, primary_methods}}.
        resource_caps (dict): {resource: cap} — used to skip zero-cap batches.
        method_costs (dict): {method: {resource: coefficient}}.
        pool_of_entity (dict | None): {(method, entity_name): pool_id} — batches
            already covered by a family-resource pool (see build_family_pools)
            are excluded here since they get a w_{m,e} variable instead.

    Returns:
        tuple[list, dict, int]:
            custom_usage_pairs   – [(method, batch_name)] with non-zero coefficient.
            custom_pair_to_index – {(method, batch_name): lp_column_index}.
            num_lp_vars          – total number of LP variables.
    """
    pool_of_entity = pool_of_entity or {}
    custom_usage_pairs = []
    for batch_name, specs in custom_batch_specs.items():
        primary_methods = specs["primary_methods"]
        if resource_caps.get(batch_name, 0.0) <= 0:
            continue
        for m in active_methods:
            if m not in primary_methods:
                continue
            if (m, batch_name) in pool_of_entity:
                continue
            if method_costs.get(m, {}).get(batch_name, 0.0) > 0:
                custom_usage_pairs.append((m, batch_name))

    custom_pair_to_index = {pair: num_methods + i for i, pair in enumerate(custom_usage_pairs)}
    num_lp_vars = num_methods + len(custom_usage_pairs)
    return custom_usage_pairs, custom_pair_to_index, num_lp_vars


def build_family_pools(active_methods, method_costs, resource_to_group, group_to_custom_batches, custom_batch_specs):
    """Group a method's own resources into OR-substitutable pools.

    Two (or more) resources are pooled for a given method when they belong to
    the same higher-level resource family (e.g. "Arable land" and "Other land"
    are both "Land") and both currently have a non-zero coefficient for that
    method — i.e. both are active (primary, or a secondary resource the user
    switched on). Pooling means the method can satisfy its need from either
    one (or a mix) instead of requiring all of them at once (AND), the same
    way a custom batch already complements its reference standard resource.

    Any custom batch attached to a pooled standard resource joins the same
    pool, so the batch and its reference resource stay complementary exactly
    as they already do outside a family pool.

    Args:
        active_methods (list[str]): ordered active CDR method names.
        method_costs (dict): {method: {resource: coefficient}}.
        resource_to_group (dict | None): {resource: family group name}. When
            None or empty, no method has any pool (fully backward compatible).
        group_to_custom_batches (dict): {reference_resource: [batch_name, ...]}.
        custom_batch_specs (dict): {batch_name: {group, amount, primary_methods}}.

    Returns:
        tuple[dict, dict]:
            pool_of_entity – {(method, entity_name): pool_id} for every entity
                (standard resource or custom batch) that belongs to a real
                pool (2+ standard resources sharing a family for that method).
            pools          – {pool_id: [entity_name, ...]}.
    """
    resource_to_group = resource_to_group or {}
    pool_of_entity = {}
    pools = {}
    for m in active_methods:
        by_group = {}
        for r, c in method_costs.get(m, {}).items():
            if c <= 0:
                continue
            g = resource_to_group.get(r)
            if g is None:
                continue
            by_group.setdefault(g, []).append(r)
        for g, members in by_group.items():
            if len(members) < 2:
                continue
            pool_id = (m, g)
            entities = list(members)
            for r in members:
                for batch_name in group_to_custom_batches.get(r, []):
                    if m in custom_batch_specs.get(batch_name, {}).get("primary_methods", set()):
                        entities.append(batch_name)
            pools[pool_id] = entities
            for e in entities:
                pool_of_entity[(m, e)] = pool_id
    return pool_of_entity, pools


def build_pool_indices(offset, pool_of_entity, resource_caps):
    """Build the LP variable index for family-pool draw variables w_{m,e}.

    Analogous to build_lp_indices's z_{m,n}, but for resources pooled because
    they share a resource family (see build_family_pools), not because a
    custom batch substitutes a single standard resource.

    Args:
        offset (int): first free LP column index (after y_m and any z_{m,n}).
        pool_of_entity (dict): {(method, entity_name): pool_id}.
        resource_caps (dict): {resource: cap} — used to skip zero-cap entities.

    Returns:
        tuple[list, dict, int]:
            pool_pairs         – [(method, entity_name), ...].
            pool_pair_to_index – {(method, entity_name): lp_column_index}.
            num_lp_vars        – total LP variable count including these columns.
    """
    pool_pairs = sorted(
        pair for pair in pool_of_entity if resource_caps.get(pair[1], 0.0) > 0
    )
    pool_pair_to_index = {pair: offset + i for i, pair in enumerate(pool_pairs)}
    num_lp_vars = offset + len(pool_pairs)
    return pool_pairs, pool_pair_to_index, num_lp_vars


def build_constraint_matrix(
    active_methods, method_to_index, num_lp_vars, num_methods,
    resource_caps, method_costs, method_constraints,
    custom_batch_specs, group_to_custom_batches,
    resources_no_batches, resources_with_batches,
    custom_usage_pairs, custom_pair_to_index,
    pool_of_entity=None, pools=None, pool_pair_to_index=None,
):
    """Assemble the LP inequality constraint matrix (A_ub · x ≤ b_ub).

    Builds six groups of constraints:
      A – standard resources with no custom substitutes.
      B – standard resources that have custom-batch substitutes.
      C – custom batch capacity limits.
      D1/D2 – coupling constraints between y_m and z_{m,n}.
      F – family-pool entity capacity limits, and their own D1/D2 coupling to y_m.
      E – per-method output caps (percent or absolute).

    Args:
        active_methods (list[str]): ordered active CDR method names.
        method_to_index (dict): {method: lp_column_index}.
        num_lp_vars (int): total LP variable count.
        num_methods (int): number of y_m variables.
        resource_caps (dict): {resource: cap}.
        method_costs (dict): {method: {resource: coefficient}}.
        method_constraints (dict): {method: {cap_type, cap_value}}.
        custom_batch_specs (dict): {batch_name: {group, amount, primary_methods}}.
        group_to_custom_batches (dict): {reference_resource: [batch_name]}.
        resources_no_batches (list[str]): standard resources without substitutes.
        resources_with_batches (list[str]): standard resources that have substitutes.
        custom_usage_pairs (list): [(method, batch_name)].
        custom_pair_to_index (dict): {(method, batch_name): lp_column_index}.
        pool_of_entity (dict | None): {(method, entity_name): pool_id} — see
            build_family_pools. Entities here are excluded from A/B/C's direct
            y_m contribution; they get their own w_{m,e} variable instead.
        pools (dict | None): {pool_id: [entity_name, ...]}.
        pool_pair_to_index (dict | None): {(method, entity_name): lp_column_index}.

    Returns:
        tuple[np.ndarray, np.ndarray, list]:
            A_ub               – constraint coefficient matrix.
            b_ub               – constraint right-hand-side vector.
            constraint_labels  – list of (type_str, identifier) per row.
    """
    pool_of_entity = pool_of_entity or {}
    pools = pools or {}
    pool_pair_to_index = pool_pair_to_index or {}

    A_rows, b_rows, constraint_labels = [], [], []
    resource_rows: dict[str, np.ndarray] = {}

    def get_row(name):
        if name not in resource_rows:
            resource_rows[name] = np.zeros(num_lp_vars)
        return resource_rows[name]

    # ── CONSTRAINT GROUPS ─────────────────────────────────────────────────────
    # Each group adds to one row per resource/method in the matrix A_ub.
    # The solver will find the y_m values that maximise total CO₂ removal while
    # keeping every A_ub · x ≤ b_ub inequality satisfied.
    #
    # Plain-English summary of each group:
    #   A – "Don't use more of a standard resource than is available."
    #   B – "Same, but subtract the part already covered by custom batches."
    #   C – "Each custom batch can only supply up to its declared quantity."
    #   D – "Book-keeping: custom-batch draw variables can't exceed the method's
    #        total output, and if a method has no standard resource it must draw
    #        entirely from custom batches."
    #   F – "A method with two resources from the same family (e.g. two kinds
    #        of land) can draw from either, not both at once — each still can't
    #        exceed its own declared quantity."
    #   E – "Respect each method's user-defined cap (% or absolute MtCO₂/yr)."
    # ─────────────────────────────────────────────────────────────────────────

    # (A) Standard resources — no custom substitutes.
    # Rule: total resource consumed across all methods ≤ available amount.
    # A method whose use of r is covered by a family pool (F) is skipped here —
    # its contribution comes from its w_{m,r} variable instead.
    for r in resources_no_batches:
        row = get_row(r)
        for j, m in enumerate(active_methods):
            if (m, r) in pool_of_entity:
                continue
            row[j] += method_costs.get(m, {}).get(r, 0.0)

    # (B) Standard resources that have custom-batch substitutes.
    # Rule: only the portion of the resource NOT already covered by custom batches
    # counts against the standard pool cap. Formula: Σ cost(m,r)·(y_m − Σ z_{m,n}) ≤ cap(r)
    for r in resources_with_batches:
        row = get_row(r)
        for j, m in enumerate(active_methods):
            if (m, r) in pool_of_entity:
                continue
            c_s_m = method_costs.get(m, {}).get(r, 0.0)   # coefficient for standard resource r
            row[j] += c_s_m
            for batch_name in group_to_custom_batches.get(r, []):
                if (m, batch_name) in custom_pair_to_index:
                    row[custom_pair_to_index[(m, batch_name)]] -= c_s_m   # subtract batch draw

    # (C) Custom batch capacity.
    # Rule: total resource drawn from this specific batch ≤ batch quantity Q_n.
    for batch_name, specs in custom_batch_specs.items():
        if resource_caps.get(batch_name, 0.0) <= 0:
            continue
        row = get_row(batch_name)
        for j, m in enumerate(active_methods):
            if (m, batch_name) not in custom_pair_to_index:
                continue
            c_n_m = method_costs.get(m, {}).get(batch_name, 0.0)   # coefficient for batch n
            if c_n_m > 0:
                row[custom_pair_to_index[(m, batch_name)]] += c_n_m

    # (F) Family-pool entities — each pooled entity (standard resource or batch)
    # is drawn from via its own w_{m,e} variable instead of y_m directly.
    # Rule: total resource drawn from entity e ≤ its own declared quantity —
    # exactly like A/B/C, just keyed by the pool draw variable.
    for (m, e), col in pool_pair_to_index.items():
        row = get_row(e)
        row[col] += method_costs.get(m, {}).get(e, 0.0)

    for name, row in resource_rows.items():
        if np.any(row != 0):
            A_rows.append(row)
            b_rows.append(resource_caps.get(name, 0.0))
            constraint_labels.append(("resource", name))

    # (D) Coupling constraints — link z_{m,n} draw variables to y_m output.
    # D1: Σ z_{m,n} ≤ y_m  — can't draw more from batches than the method produces.
    # D2: y_m ≤ Σ z_{m,n}  — only added when the method has no standard resource (c_s=0),
    #     which forces the method to source 100 % of its output from custom batches.
    coupling_done = set()
    for (m, batch_name) in custom_usage_pairs:
        ref = custom_batch_specs[batch_name]["group"]
        key = (m, ref)
        if key in coupling_done:
            continue
        coupling_done.add(key)
        j     = method_to_index[m]
        c_s_m = method_costs.get(m, {}).get(ref, 0.0)
        batch_cols = [
            custom_pair_to_index[(m, sibling)]
            for sibling in group_to_custom_batches.get(ref, [])
            if (m, sibling) in custom_pair_to_index
        ]
        row = np.zeros(num_lp_vars)
        row[j] = -1.0
        for col in batch_cols:
            row[col] = 1.0
        A_rows.append(row)
        b_rows.append(0.0)
        constraint_labels.append(("coupling_D1", (m, ref)))

        if c_s_m == 0:
            row2 = np.zeros(num_lp_vars)
            row2[j] = 1.0
            for col in batch_cols:
                row2[col] = -1.0
            A_rows.append(row2)
            b_rows.append(0.0)
            constraint_labels.append(("coupling_D2", (m, ref)))

    # (F, coupling) Every unit a method produces must come from exactly one (or
    # a mix) of its pool's entities: Σ w_{m,e} == y_m, expressed as two ≤ rows.
    # Unlike D1/D2 above, both directions always apply — a family pool has no
    # single "reference" resource to silently absorb the remainder.
    for pool_id, entities in pools.items():
        m = pool_id[0]
        j = method_to_index.get(m)
        if j is None:
            continue
        cols = [pool_pair_to_index[(m, e)] for e in entities if (m, e) in pool_pair_to_index]
        if not cols:
            continue
        row1 = np.zeros(num_lp_vars)
        row1[j] = -1.0
        for col in cols:
            row1[col] = 1.0
        A_rows.append(row1)
        b_rows.append(0.0)
        constraint_labels.append(("pool_D1", pool_id))

        row2 = np.zeros(num_lp_vars)
        row2[j] = 1.0
        for col in cols:
            row2[col] = -1.0
        A_rows.append(row2)
        b_rows.append(0.0)
        constraint_labels.append(("pool_D2", pool_id))

    # (E) Method output caps — respect the user's per-method limits from Tab 2.
    # Absolute cap: y_m ≤ cap_value (in MtCO₂/yr).
    # Percent cap:  y_m ≤ cap_ratio · Σ y_all  (fraction of total portfolio output).
    for j, m in enumerate(active_methods):
        constraint = method_constraints.get(m, {})
        cap_type   = constraint.get("cap_type",  "percent")
        cap_value  = constraint.get("cap_value", 100)
        if cap_type == "absolute":
            row    = np.zeros(num_lp_vars)
            row[j] = 1.0
            A_rows.append(row)
            b_rows.append(cap_value)
            constraint_labels.append(("method_cap", m))
        elif cap_type == "percent":
            cap_ratio = cap_value / 100
            row = np.zeros(num_lp_vars)
            row[:num_methods] -= cap_ratio
            row[j]            += 1.0
            A_rows.append(row)
            b_rows.append(0.0)
            constraint_labels.append(("method_cap", m))
        else:
            raise ValueError(
                f"Unknown cap_type '{cap_type}' for method '{m}'. "
                "Expected 'percent' or 'absolute'."
            )

    A_ub = np.array(A_rows) if A_rows else np.empty((0, num_lp_vars))
    b_ub = np.array(b_rows)
    return A_ub, b_ub, constraint_labels


def extract_resource_usage(
    active_methods, num_methods, method_outputs, lp_solution, method_costs,
    resource_caps, custom_batch_specs, group_to_custom_batches,
    custom_pair_to_index, standard_resources,
    resources_no_batches, resources_with_batches, custom_batch_names,
    pool_of_entity=None, pool_pair_to_index=None,
):
    """Compute per-method resource consumption from the LP solution.

    Args:
        active_methods (list[str]): ordered active CDR method names.
        num_methods (int): number of y_m variables.
        method_outputs (np.ndarray): y_m values from LP solution (length num_methods).
        lp_solution (np.ndarray): full LP solution vector x.
        method_costs (dict): {method: {resource: coefficient}}.
        resource_caps (dict): {resource: cap} — used for resource_remaining.
        custom_batch_specs (dict): {batch_name: {group, amount, primary_methods}}.
        group_to_custom_batches (dict): {reference_resource: [batch_name]}.
        custom_pair_to_index (dict): {(method, batch_name): lp_column_index}.
        standard_resources (list[str]): all non-batch resource names.
        resources_no_batches (list[str]): standard resources without substitutes.
        resources_with_batches (list[str]): standard resources with substitutes.
        custom_batch_names (set[str]): set of all custom batch names.
        pool_of_entity (dict | None): {(method, entity_name): pool_id} — see
            build_family_pools. Entities here report usage via their own
            w_{m,e} variable instead of the formulas below.
        pool_pair_to_index (dict | None): {(method, entity_name): lp_column_index}.

    Returns:
        tuple[dict, dict, dict]:
            resource_usage       – {resource: {method: qty_used}}.
            resource_used_total  – {resource: total_qty_used}.
            resource_remaining   – {resource: qty_remaining}.
    """
    pool_of_entity = pool_of_entity or {}
    pool_pair_to_index = pool_pair_to_index or {}

    all_resource_names = standard_resources + list(custom_batch_names)
    resource_usage = {r: {} for r in all_resource_names}

    # LP solvers (HiGHS) return floating-point residuals on the order of 1e-14 to
    # 1e-12, so 1e-14 safely filters out numerical noise while preserving
    # legitimately tiny usage values (e.g. Non-arable land coefficient = 1e-11).
    _USAGE_EPS = 1e-14
    for r in resources_no_batches:
        for j, m in enumerate(active_methods):
            if (m, r) in pool_of_entity:
                continue
            used = method_costs.get(m, {}).get(r, 0.0) * method_outputs[j]
            if used > _USAGE_EPS:
                resource_usage[r][m] = float(used)

    for r in resources_with_batches:
        for j, m in enumerate(active_methods):
            if (m, r) in pool_of_entity:
                continue
            c_s_m = method_costs.get(m, {}).get(r, 0.0)
            z_sum = sum(
                lp_solution[custom_pair_to_index[(m, batch_name)]]
                for batch_name in group_to_custom_batches.get(r, [])
                if (m, batch_name) in custom_pair_to_index
            )
            used = c_s_m * (method_outputs[j] - z_sum)
            if used > _USAGE_EPS:
                resource_usage[r][m] = float(used)

    for batch_name, specs in custom_batch_specs.items():
        primary_methods = specs["primary_methods"]
        for j, m in enumerate(active_methods):
            if m not in primary_methods:
                continue
            if (m, batch_name) in pool_of_entity:
                continue
            c_n_m = method_costs.get(m, {}).get(batch_name, 0.0)
            if c_n_m <= 0:
                continue
            used = (
                c_n_m * lp_solution[custom_pair_to_index[(m, batch_name)]]
                if (m, batch_name) in custom_pair_to_index
                else 0.0
            )
            if used > _USAGE_EPS:
                resource_usage[batch_name][m] = float(used)

    # Family-pool entities (standard resources or batches pooled across a
    # method's own resource family) — usage comes directly from w_{m,e}.
    for (m, e), col in pool_pair_to_index.items():
        used = method_costs.get(m, {}).get(e, 0.0) * lp_solution[col]
        if used > _USAGE_EPS:
            resource_usage.setdefault(e, {})[m] = float(used)

    _EPS = 1e-14
    resource_used_total = {
        r: round(sum(v.values()), 10) for r, v in resource_usage.items()
    }
    all_caps = {
        **resource_caps,
        **{batch_name: specs["amount"] for batch_name, specs in custom_batch_specs.items()},
    }
    resource_remaining = {
        r: max(0.0, v) if abs(v) >= _EPS else 0.0
        for r, v in {
            r: all_caps.get(r, 0.0) - resource_used_total.get(r, 0.0)
            for r in all_caps
        }.items()
    }
    return resource_usage, resource_used_total, resource_remaining


def resolve_pooling_tiebreak(
    num_methods, num_lp_vars, method_outputs, custom_usage_pairs,
    custom_pair_to_index, custom_batch_specs, method_costs, A_ub, b_ub, bounds,
    pool_pair_to_index=None,
):
    """Re-solve the LP with each method's output pinned, preferring the more
    efficient resource within each standard/custom pooling pair or family pool.

    The main LP (run_optimisation) maximises total CO2 removed only — nothing
    in its objective distinguishes a standard resource from a custom batch
    that substitutes for it (or from a same-family alternative resource), so
    when several could satisfy the same demand the solver picks an arbitrary
    split. This second pass fixes that: it pins every active method's output
    y_m to its pass-1 value (so total_removed and method_usage never change)
    and minimises Σ (c_n - c_s) · z_{m,n} across all (method, batch) pooling
    pairs, plus Σ c_e · w_{m,e} across all family-pool entities — i.e. it
    prefers fully using whichever resource has the lower coefficient (more
    MtCO2 removed per unit of resource) before drawing on the less efficient
    one. For a two-entity pool this reduces exactly to the batch-pooling
    formula (pinning w_other = y_m - w_e turns Σ c_e·w_e into a (c_e - c_other)
    term up to a constant), so it generalises the original tie-break rather
    than replacing it.

    Args:
        num_methods (int): number of y_m variables.
        num_lp_vars (int): total LP variable count (pass-1 sizing).
        method_outputs (np.ndarray): y_m values from the pass-1 solution.
        custom_usage_pairs (list): [(method, batch_name)] with non-zero coefficient.
        custom_pair_to_index (dict): {(method, batch_name): lp_column_index}.
        custom_batch_specs (dict): {batch_name: {group, amount, primary_methods}}.
        method_costs (dict): {method: {resource: coefficient}}.
        A_ub, b_ub: pass-1 constraint matrix/vector (reused unchanged — pass-1's
            own solution already satisfies them, so this second LP is always feasible).
        bounds: pass-1 variable bounds (reused unchanged).
        pool_pair_to_index (dict | None): {(method, entity_name): lp_column_index}
            for family-pool draw variables — see build_family_pools.

    Returns:
        np.ndarray | None: the tie-broken solution vector, or None if there are
            no pooling pairs to break ties on, or if the second solve fails
            (in which case the caller should fall back to the pass-1 solution).
    """
    pool_pair_to_index = pool_pair_to_index or {}
    if not custom_usage_pairs and not pool_pair_to_index:
        return None

    A_eq = np.zeros((num_methods, num_lp_vars))
    for j in range(num_methods):
        A_eq[j, j] = 1.0
    b_eq = np.array(method_outputs, dtype=float)

    objective2 = np.zeros(num_lp_vars)
    for (m, batch_name) in custom_usage_pairs:
        ref = custom_batch_specs[batch_name]["group"]
        c_s = method_costs.get(m, {}).get(ref, 0.0)
        c_n = method_costs.get(m, {}).get(batch_name, 0.0)
        objective2[custom_pair_to_index[(m, batch_name)]] = c_n - c_s

    for (m, e), col in pool_pair_to_index.items():
        objective2[col] = method_costs.get(m, {}).get(e, 0.0)

    try:
        result = linprog(
            objective2, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq,
            bounds=bounds, method="highs",
        )
    except Exception:
        return None
    return result.x if result.success else None


def run_perturbation(constraint_labels, A_ub, b_ub, objective_coefficients, bounds, num_methods, total_removed):
    """Compute marginal CDR gain per unit for each binding resource constraint.

    Re-solves the LP with each resource constraint right-hand side increased by 1
    to measure the actual CO₂ gain that one extra unit would unlock.

    Args:
        constraint_labels (list): list of (type_str, identifier) per constraint row.
        A_ub (np.ndarray): LP constraint coefficient matrix.
        b_ub (np.ndarray): LP constraint right-hand-side vector.
        objective_coefficients (np.ndarray): LP objective vector.
        bounds (list): variable bounds for linprog.
        num_methods (int): number of y_m variables (for summing total removed).
        total_removed (float): baseline total CO₂ removed.

    Returns:
        dict: {resource_or_batch_name: CDR_gain_for_plus_1_unit (float)}.
    """
    resource_actual_gain = {}
    for i, (ctype, cid) in enumerate(constraint_labels):
        if ctype not in ("resource", "custom_batch"):
            continue
        b_perturbed = b_ub.copy()
        b_perturbed[i] += 1.0
        try:
            perturbed = linprog(
                objective_coefficients, A_ub=A_ub, b_ub=b_perturbed, bounds=bounds, method="highs"
            )
            if perturbed.success:
                new_total = float(perturbed.x[:num_methods].sum())
                resource_actual_gain[cid] = max(new_total - total_removed, 0.0)
        except Exception:
            pass
    return resource_actual_gain


def compute_standalone_potential(method, resource_caps, method_costs, custom_resources, enabled_methods, resource_to_group=None):
    """Compute a method's own maximum achievable output, in isolation.

    "Isolation" means solving as if this were the only active method — no
    competition for shared resources from any other method, and no cap of
    its own applied. This is what "% of its total potential" (the method's
    OWN ceiling) should be a percentage of, as opposed to a self-referential
    "% of the combined output of every active method", which depends on
    what every other method happens to be doing.

    Args:
        method (str): the method to isolate.
        resource_caps (dict): {resource: available amount}.
        method_costs (dict): {method: {resource: coefficient}} — already
            reflecting any enabled_methods zeroing.
        custom_resources (list | None): custom batch dicts (safe to pass the
            full list — batches tied to other methods simply produce no
            usage pairs once every other method is excluded).
        enabled_methods (dict | None): {resource: {method: bool}}.
        resource_to_group (dict | None): {resource: family group name} — see
            build_family_pools. Must match what the outer run_optimisation
            call uses, so the method's standalone potential already reflects
            OR-availability across its own same-family resources.

    Returns:
        float: the method's standalone potential in MtCO₂ (0.0 if it can't
            produce anything on its own).
    """
    _UNCAPPED = 1e15
    success, result = run_optimisation(
        resource_caps=resource_caps,
        method_constraints={method: {"active": True, "cap_type": "absolute", "cap_value": _UNCAPPED}},
        method_costs={method: dict(method_costs.get(method, {}))},
        custom_resources=custom_resources,
        enabled_methods=enabled_methods,
        resource_to_group=resource_to_group,
    )
    return result["total_removed"] if success else 0.0


def run_optimisation(
    resource_caps,
    method_constraints,
    method_costs,
    custom_resources=None,
    enabled_methods=None,
    resource_to_group=None,
):
    """Maximize total CO₂ removal subject to resource caps and method constraints.

    Solves a linear programme (HiGHS solver via scipy.optimize.linprog).
    Also runs a perturbation analysis: for each binding resource constraint,
    re-solves with cap + 1 to compute the marginal CDR gain per unit.

    LP variables (three types, packed into one vector x):
      y_m     = MtCO₂ removed by CDR method m — the main "answer" (indices 0..num_methods-1).
      z_{m,n} = MtCO₂ that method m draws specifically from custom batch n — a helper
                variable introduced when the user adds a custom resource batch; it lets the
                solver track which batch supplies which method without double-counting.
      w_{m,e} = MtCO₂ that method m draws specifically from same-family entity e (a
                standard resource or custom batch inside an OR-substitutable resource
                pool — see build_family_pools); same bookkeeping role as z_{m,n}, for
                resources that compete with each other rather than with the standard/
                custom-batch split.
                (z and w columns together occupy indices num_methods..num_lp_vars-1.)

    Args:
        resource_caps (dict): {resource: available amount (float)}.
        method_constraints (dict): {method: {active (bool), cap_type (str), cap_value (float)}}.
        method_costs (dict): {method: {resource: coefficient (float)}}.
        custom_resources (list | None): list of custom batch dicts with keys
                                         name, group, amount, methods.
        enabled_methods (dict | None): {resource: {method: bool}} — False entries
                                        zero-out the corresponding coefficient before solving.
        resource_to_group (dict | None): {resource: family group name} (e.g.
            "Arable land" → "Land"). When a method has 2+ resources from the
            same family with a non-zero coefficient, they're treated as
            OR-substitutable (draw from either, or a mix) instead of both
            being required at once — see build_family_pools. None (default)
            disables this and preserves the original AND-of-all-resources
            behaviour.

    Returns:
        tuple[bool, dict]:
            (True, result) on success, where result contains:
                total_removed (float), method_usage (dict), resource_usage (dict),
                resource_used_total (dict), resource_remaining (dict),
                resource_actual_gain (dict).
            (False, {"message": str}) on failure.
    """
    if resource_caps is None:
        return False, {"message": "resource_caps cannot be None."}

    # Sanitize inputs that may bypass the UI's own safeguards (widgets floor
    # at 0, but a scenario loaded from Excel/JSON/ZIP, or a hand-edited
    # data/inputs.csv, is never passed through those widgets). A negative
    # resource cap is caught cleanly downstream as an infeasible LP, but a
    # negative coefficient is not: with no other binding constraint on that
    # resource it lets the solver report an unbounded-looking "success" with
    # an astronomically large total instead of an error. Clamping everything
    # to >= 0 here — the single funnel every input path (UI, Excel, ZIP,
    # JSON, or CSV) passes through before reaching the solver — closes both.
    resource_caps  = {r: max(0.0, float(v)) for r, v in resource_caps.items()}
    method_costs   = {
        m: {r: max(0.0, float(v)) for r, v in resources.items()}
        for m, resources in method_costs.items()
    }
    if custom_resources:
        custom_resources = [
            {
                **batch,
                "amount": max(0.0, float(batch.get("amount", 0.0))),
                "methods": {
                    m: {k: max(0.0, float(v)) for k, v in coefs.items()}
                    for m, coefs in batch.get("methods", {}).items()
                },
            }
            for batch in custom_resources
        ]

    method_costs  = apply_enabled_methods(method_costs, enabled_methods)

    # Clamp negative absolute caps to 0 to avoid infeasible LP.
    for m, c in method_constraints.items():
        if c.get("cap_type") == "absolute" and c.get("cap_value", 0) < 0:
            method_constraints = {
                **method_constraints,
                m: {**c, "cap_value": 0.0},
            }

    # Convert "percent" caps into an equivalent absolute cap computed from the
    # method's OWN standalone potential (its max output if it were the only
    # active method, unconstrained by any cap) — matches the documented
    # meaning of "% of its total potential", instead of a self-referential
    # "% of the combined output of every active method". The old
    # self-referential form made one method's cap depend on what every other
    # active method happened to produce, and — since it's a ratio against a
    # value that includes the capped method itself — an out-of-range percent
    # (e.g. a negative value from a corrupted import) could force the ENTIRE
    # portfolio to zero, not just the misconfigured method. Clamp cap_value to
    # [0, 100] first (the UI slider already restricts to this range, but a
    # loaded scenario file might not).
    resolved_constraints = {}
    for m, c in method_constraints.items():
        if c.get("cap_type") == "percent" and c.get("active", True):
            cap_value = max(0.0, min(100.0, c.get("cap_value", 100)))
            potential = compute_standalone_potential(
                m, resource_caps, method_costs, custom_resources, enabled_methods,
                resource_to_group=resource_to_group,
            )
            resolved_constraints[m] = {**c, "cap_type": "absolute", "cap_value": potential * cap_value / 100.0}
        else:
            resolved_constraints[m] = c
    method_constraints = resolved_constraints

    custom_batch_specs, group_to_custom_batches, custom_batch_names = build_custom_batch_structures(
        method_costs, resource_caps, custom_resources
    )

    # Deactivate methods whose effective coefficients are all zero — they would
    # otherwise create unconstrained LP variables and cause an unbounded problem.
    # Runs AFTER build_custom_batch_structures so that batch coefficients injected
    # into method_costs are visible here. Without this ordering, a method that relies
    # solely on a custom batch (standard pool blocked via enabled_methods) would be
    # wrongly deactivated because its standard coefficients all appear zero.
    # Iterate over method_costs (not just method_constraints) to catch methods
    # that are active by default (not present in method_constraints at all).
    safe = {m: c.copy() for m, c in method_constraints.items()}
    for m in method_costs:
        c = safe.get(m, {"active": True, "cap_type": "percent", "cap_value": 100})
        if not c.get("active", True):
            continue
        coeffs = method_costs.get(m, {})
        if all(v == 0 for v in coeffs.values()) if coeffs else True:
            safe[m] = {**c, "active": False}
    method_constraints = safe

    active_methods  = [m for m in method_costs if method_constraints.get(m, {}).get("active", True)]
    num_methods     = len(active_methods)
    method_to_index = {m: j for j, m in enumerate(active_methods)}
    if not num_methods:
        return False, {"message": "No active methods."}

    # If there are no resource caps at all (not even custom batches), the LP has no
    # upper bound on y_m — only percent-type method caps apply, which are relative and
    # don't fix the scale. Return zero rather than letting HiGHS report "unbounded".
    all_caps = {**resource_caps, **{b["name"]: float(b.get("amount", 0)) for b in (custom_resources or [])}}
    if not all_caps:
        empty_result = {
            "total_removed": 0.0,
            "method_usage": {},
            "resource_usage": {},
            "resource_used_total": {},
            "resource_remaining": {},
            "resource_actual_gain": {},
        }
        return True, empty_result

    pool_of_entity, pools = build_family_pools(
        active_methods, method_costs, resource_to_group, group_to_custom_batches, custom_batch_specs
    )

    custom_usage_pairs, custom_pair_to_index, num_lp_vars = build_lp_indices(
        active_methods, num_methods, custom_batch_specs, resource_caps, method_costs, pool_of_entity
    )
    pool_pairs, pool_pair_to_index, num_lp_vars = build_pool_indices(
        num_lp_vars, pool_of_entity, resource_caps
    )

    standard_resources     = [r for r in resource_caps if r not in custom_batch_names]
    resources_with_batches = [r for r in standard_resources if r in group_to_custom_batches]
    resources_no_batches   = [r for r in standard_resources if r not in group_to_custom_batches]

    try:
        A_ub, b_ub, constraint_labels = build_constraint_matrix(
            active_methods, method_to_index, num_lp_vars, num_methods,
            resource_caps, method_costs, method_constraints,
            custom_batch_specs, group_to_custom_batches,
            resources_no_batches, resources_with_batches,
            custom_usage_pairs, custom_pair_to_index,
            pool_of_entity, pools, pool_pair_to_index,
        )
    except ValueError as e:
        return False, {"message": str(e)}

    # linprog minimises by convention, so we negate the objective to maximise Σ y_m.
    objective_coefficients = np.zeros(num_lp_vars)
    objective_coefficients[:num_methods] = -1.0   # only y_m variables enter the objective
    bounds = [(0, None)] * num_lp_vars             # all variables must be ≥ 0

    try:
        result = linprog(objective_coefficients, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method="highs")
        if not result.success:
            return False, {"message": result.message}

        lp_solution    = result.x
        method_outputs = lp_solution[:num_methods]
        total_removed  = float(method_outputs.sum())
        method_usage   = {
            active_methods[j]: float(method_outputs[j])
            for j in range(num_methods) if method_outputs[j] > 1e-9
        }

        tiebreak_solution = resolve_pooling_tiebreak(
            num_methods, num_lp_vars, method_outputs, custom_usage_pairs,
            custom_pair_to_index, custom_batch_specs, method_costs, A_ub, b_ub, bounds,
            pool_pair_to_index,
        )
        lp_solution_for_usage = tiebreak_solution if tiebreak_solution is not None else lp_solution

        resource_usage, resource_used_total, resource_remaining = extract_resource_usage(
            active_methods, num_methods, method_outputs, lp_solution_for_usage, method_costs,
            resource_caps, custom_batch_specs, group_to_custom_batches,
            custom_pair_to_index, standard_resources,
            resources_no_batches, resources_with_batches, custom_batch_names,
            pool_of_entity, pool_pair_to_index,
        )

        resource_actual_gain = run_perturbation(
            constraint_labels, A_ub, b_ub, objective_coefficients, bounds,
            num_methods, total_removed,
        )
        
        return True, {
            "total_removed":        total_removed,
            "method_usage":         method_usage,
            "resource_usage":       resource_usage,
            "resource_used_total":  resource_used_total,
            "resource_remaining":   resource_remaining,
            "resource_actual_gain": resource_actual_gain,
        }

    except Exception as e:
        return False, {"message": str(e)}