import numpy as np
from scipy.optimize import linprog


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.

    LP variables:
      y_m    = CO₂ removed by method m  (indices 0..M-1)
      z_{m,n} = CO₂ drawn from custom batch n by method m  (indices M..N-1)
    """

    resource_caps = dict(resource_caps)
    method_costs  = {m: dict(v) for m, v in method_costs.items()}
    print("first enabled_methods ----------- ")

    # ── Apply manual allocation overrides ─────────────────────────────────────
    enabled_methods = enabled_methods or {}
    print(enabled_methods)
    for method, resources in method_costs.items():
        for resource in list(resources):
            if not enabled_methods.get(resource, {}).get(method, True):
                method_costs[method][resource] = 0.0

    # ── Custom batch metadata ─────────────────────────────────────────────────
    custom_batch_info = {}
    for batch in (custom_resources or []):
        custom_batch_info[batch["name"]] = {
            "group":           batch["group"],
            "amount":          float(batch["amount"]),
            "primary_methods": set(batch.get("methods", {}).keys()),
        }
    custom_names   = set(custom_batch_info)
    ref_to_batches = {}
    for bname, info in custom_batch_info.items():
        ref_to_batches.setdefault(info["group"], []).append(bname)
        if bname not in resource_caps:
            resource_caps[bname] = info["amount"]

    # ── Inject primary batch coefficients (fallback: use batch median if coef=0)
    for batch in (custom_resources or []):
        bname = batch["name"]
        for method, coefs in batch["methods"].items():
            if method_costs.get(method, {}).get(bname, 0.0) == 0.0:
                method_costs.setdefault(method, {})[bname] = coefs["median"]

    # ── Debug custom resources ────────────────────────────────────────────────
    print(f"\n{'='*50}")
    print(f"custom_resources received: {len(custom_resources or [])} batch(es)")
    for batch in (custom_resources or []):
        r_key = batch["name"]
        print(f"  → '{r_key}' | amount: {resource_caps.get(r_key)} | methods: {list(batch['methods'].keys())}")
        for method in batch["methods"]:
            coef = method_costs.get(method, {}).get(r_key)
            print(f"      method_costs['{method}']['{r_key}'] = {coef}")
    print(method_costs)
    print(f"Total resources in matrix: {len(resource_caps)}")
    print(f"Total resources in matrix: {resource_caps}")
    print(f"Resources: {list(resource_caps.keys())}")
    print(f"{'='*50}\n")

    # ── Active methods ────────────────────────────────────────────────────────
    methods    = [m for m in method_costs if method_constraints.get(m, {}).get("active", True)]
    M          = len(methods)
    method_idx = {m: j for j, m in enumerate(methods)}
    if not M:
        return False, {"message": "No active methods."}

    # ── Split variables z_{m,n} ───────────────────────────────────────────────
    z_pairs = []
    for bname, info in custom_batch_info.items():
        ref             = info["group"]
        primary_methods = info["primary_methods"]
        if resource_caps.get(bname, 0.0) <= 0:
            continue
        for m in methods:
            if m not in primary_methods:
                continue
            c_n = method_costs.get(m, {}).get(bname, 0.0)
            if c_n > 0:
                z_pairs.append((m, bname))

    Z     = len(z_pairs)
    z_idx = {pair: M + i for i, pair in enumerate(z_pairs)}
    N     = M + Z

    print(f"z_pairs ({Z}): {z_pairs}")

    # ── Resource categorisation ───────────────────────────────────────────────
    standard_resources     = [r for r in resource_caps if r not in custom_names]
    resources_with_batches = [r for r in standard_resources if r in ref_to_batches]
    resources_no_batches   = [r for r in standard_resources if r not in ref_to_batches]

    # ── Build LP constraint matrix ────────────────────────────────────────────
    A_rows, b_rows, constraint_labels = [], [], []

    # (A) Standard resources — no custom substitutes
    for r in resources_no_batches:
        row = np.zeros(N)
        for j, m in enumerate(methods):
            row[j] = method_costs.get(m, {}).get(r, 0.0)
        if np.any(row[:M] != 0):
            A_rows.append(row)
            b_rows.append(resource_caps[r])
            constraint_labels.append(("resource", r))

    # (B) Standard resources — with custom substitutes (subtract z contribution)
    # Σ_m cost(m,r) · (y_m − Σ_n z_{m,n}) ≤ cap(r)
    for r in resources_with_batches:
        row = np.zeros(N)
        for j, m in enumerate(methods):
            c_s_m = method_costs.get(m, {}).get(r, 0.0)
            row[j] = c_s_m
            for bname in ref_to_batches.get(r, []):
                if (m, bname) in z_idx:
                    row[z_idx[(m, bname)]] -= c_s_m
        if np.any(row != 0):
            A_rows.append(row)
            b_rows.append(resource_caps[r])
            constraint_labels.append(("resource", r))

    # (C) Custom batch capacity: Σ_m cost(m,n) · z_{m,n} ≤ Q_n
    for bname, info in custom_batch_info.items():
        if resource_caps.get(bname, 0.0) <= 0:
            continue
        row = np.zeros(N)
        for j, m in enumerate(methods):
            if (m, bname) not in z_idx:
                continue
            c_n_m = method_costs.get(m, {}).get(bname, 0.0)
            if c_n_m > 0:
                row[z_idx[(m, bname)]] = c_n_m
        if np.any(row != 0):
            A_rows.append(row)
            b_rows.append(resource_caps.get(bname, 0.0))
            constraint_labels.append(("custom_batch", bname))

    # (D1) Σ z_{m,n} ≤ y_m  — keeps standard resource usage non-negative
    # (D2) y_m ≤ Σ z_{m,n}  — only when c_s=0: entire output must come from batch
    #      D1 + D2 together force y_m = Σz when the standard resource is unavailable
    coupling_done = set()
    for (m, bname) in z_pairs:
        ref = custom_batch_info[bname]["group"]
        key = (m, ref)
        if key in coupling_done:
            continue
        coupling_done.add(key)
        j     = method_idx[m]
        c_s_m = method_costs.get(m, {}).get(ref, 0.0)
        batch_cols = [z_idx[(m, bn)] for bn in ref_to_batches.get(ref, []) if (m, bn) in z_idx]

        row = np.zeros(N)
        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(N)
            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 caps
    for j, m in enumerate(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(N)
            row[j]  = 1.0
            A_rows.append(row)
            b_rows.append(cap_value)
            constraint_labels.append(("method_cap", m))
        else:
            cap_ratio  = cap_value / 100.0
            row        = np.zeros(N)
            row[:M]   -= cap_ratio
            row[j]    += 1.0
            A_rows.append(row)
            b_rows.append(0.0)
            constraint_labels.append(("method_cap", m))

    A_ub = np.array(A_rows) if A_rows else np.empty((0, N))
    b_ub = np.array(b_rows)

    # Objective: maximize Σ y_m  →  minimize −Σ y_m
    c_obj         = np.zeros(N)
    c_obj[:M]     = -1.0
    bounds        = [(0, None)] * N

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

        sol = result.x
        y   = sol[:M]

        total_removed = float(y.sum())
        method_usage  = {methods[j]: float(y[j]) for j in range(M) if y[j] > 1e-9}

        # ── Resource usage ────────────────────────────────────────────────────
        all_resource_names = standard_resources + list(custom_names)
        resource_usage     = {r: {} for r in all_resource_names}

        for r in resources_no_batches:
            for j, m in enumerate(methods):
                used = method_costs.get(m, {}).get(r, 0.0) * y[j]
                if used > 1e-9:
                    resource_usage[r][m] = float(used)

        for r in resources_with_batches:
            for j, m in enumerate(methods):
                c_s_m = method_costs.get(m, {}).get(r, 0.0)
                z_sum = sum(
                    sol[z_idx[(m, bn)]]
                    for bn in ref_to_batches.get(r, [])
                    if (m, bn) in z_idx
                )
                used = c_s_m * (y[j] - z_sum)
                if used > 1e-9:
                    resource_usage[r][m] = float(used)

        for bname, info in custom_batch_info.items():
            primary_methods = info["primary_methods"]
            for j, m in enumerate(methods):
                if m not in primary_methods:
                    continue
                c_n_m = method_costs.get(m, {}).get(bname, 0.0)
                if c_n_m <= 0:
                    continue
                used = c_n_m * sol[z_idx[(m, bname)]] if (m, bname) in z_idx else 0.0
                if used > 1e-9:
                    resource_usage[bname][m] = float(used)

        resource_used_total = {r: sum(v.values()) for r, v in resource_usage.items()}

        all_caps           = {**resource_caps,
                              **{bn: info["amount"] for bn, info in custom_batch_info.items()}}
        resource_remaining = {r: all_caps.get(r, 0.0) - resource_used_total.get(r, 0.0)
                              for r in all_caps}

        # ── Perturbation analysis: marginal CO₂ gain per +1 unit of each resource
        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(
                    c_obj, A_ub=A_ub, b_ub=b_perturbed, bounds=bounds, method="highs"
                )
                if perturbed.success:
                    new_total = float(perturbed.x[:M].sum())
                    resource_actual_gain[cid] = max(new_total - total_removed, 0.0)
            except Exception:
                pass

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