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_optimization.

    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):
    """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}}.

    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.
    """
    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 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_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,
):
    """Assemble the LP inequality constraint matrix (A_ub · x ≤ b_ub).

    Builds five 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}.
      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}.

    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.
    """
    A_rows, b_rows, constraint_labels = [], [], []

    # ── CONSTRAINT GROUPS ─────────────────────────────────────────────────────
    # Each group adds one row per resource/method to 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."
    #   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.
    for r in resources_no_batches:
        row = np.zeros(num_lp_vars)
        for j, m in enumerate(active_methods):
            row[j] = method_costs.get(m, {}).get(r, 0.0)
        if np.any(row[:num_methods] != 0):
            A_rows.append(row)
            b_rows.append(resource_caps[r])
            constraint_labels.append(("resource", r))

    # (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 = np.zeros(num_lp_vars)
        for j, m in enumerate(active_methods):
            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
        if np.any(row != 0):
            A_rows.append(row)
            b_rows.append(resource_caps[r])
            constraint_labels.append(("resource", r))

    # (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 = np.zeros(num_lp_vars)
        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
        if np.any(row != 0):
            A_rows.append(row)
            b_rows.append(resource_caps.get(batch_name, 0.0))
            constraint_labels.append(("custom_batch", batch_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)))

    # (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,
):
    """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.

    Returns:
        tuple[dict, dict, dict]:
            resource_usage       – {resource: {method: qty_used}}.
            resource_used_total  – {resource: total_qty_used}.
            resource_remaining   – {resource: qty_remaining}.
    """
    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):
            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):
            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
            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)

    _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 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 run_optimization(
    resource_caps,
    method_constraints,
    method_costs,
    custom_resources=None,
    enabled_methods=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 (two 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
                (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.

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

    resource_caps = dict(resource_caps)
    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},
            }

    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

    custom_usage_pairs, custom_pair_to_index, num_lp_vars = build_lp_indices(
        active_methods, num_methods, custom_batch_specs, resource_caps, method_costs
    )

    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,
        )
    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
        }

        resource_usage, resource_used_total, resource_remaining = 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,
        )

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