"""Make the decomposition satisfy property (B) of Definition 1.

(B): if c(S_i) < 2 c(S_{i-1}) then 8 c(S_{i-1}) < c(S_{i+1}).

The greedy partial-cover oracle satisfies (A) but not (B), because (B) constrains
the COST GROWTH PATTERN across consecutive pieces and a cost-per-element rule
does not control it. Fix: after building the greedy pieces, repeatedly MERGE an
adjacent pair that violates (B). Merging preserves (A) -- a union of consecutive
pieces still covers at least half the remaining prediction at its start -- and
strictly reduces the piece count, so the procedure terminates.
"""
import json, numpy as np
from ice_exp import (make_instance, ilp_cover, ilp_partial, greedy_partial,
                     online_greedy, ice)

RES = json.load(open("ice_results.json"))


def build_pieces(A, cost, Xhat):
    rem = sorted(Xhat); pieces = []
    while rem:
        need = int(np.ceil(len(rem)/2))
        S, cov = greedy_partial(A, cost, rem, need)
        if not S: break
        Xi = sorted(cov & set(rem))
        pieces.append({"S": list(S), "X": Xi, "c": float(cost[list(S)].sum()),
                       "need_at_start": need, "rem_at_start": len(rem)})
        rem = [e for e in rem if e not in cov]
    return pieces


def violates_B(p, i):
    """(B) applies at index i when c_i < 2 c_{i-1}; it then demands 8 c_{i-1} < c_{i+1}."""
    if i < 1 or i+1 >= len(p): return False
    return (p[i]["c"] < 2*p[i-1]["c"]) and not (8*p[i-1]["c"] < p[i+1]["c"])


def enforce_B(pieces, cost, max_iter=200):
    p = [dict(x) for x in pieces]
    for _ in range(max_iter):
        bad = [i for i in range(len(p)) if violates_B(p, i)]
        if not bad: break
        i = bad[0]
        merged = {"S": sorted(set(p[i]["S"]) | set(p[i+1]["S"])),
                  "X": sorted(set(p[i]["X"]) | set(p[i+1]["X"])),
                  "need_at_start": p[i]["need_at_start"],
                  "rem_at_start": p[i]["rem_at_start"]}
        merged["c"] = float(cost[merged["S"]].sum())
        p = p[:i] + [merged] + p[i+2:]
        if len(p) <= 2: break
    return p


def check(p):
    okA = all(len(x["X"]) >= x["need_at_start"] for x in p)
    okB = not any(violates_B(p, i) for i in range(len(p)))
    return okA, okB


def run():
    rows = []
    for (m, n) in ((40, 60), (60, 90), (80, 120)):
        for corrupt in (0.0, 0.2, 0.5, 1.0):
            rng = np.random.default_rng(m*10+int(corrupt*10))
            A, cost = make_instance(m, n, rng)
            req = list(rng.permutation(m)[:m//2])
            optcost, _ = ilp_cover(A, cost, req)
            Xhat = set(req); nswap = int(corrupt*len(req))
            if nswap:
                drop = set(rng.choice(list(Xhat), size=min(nswap, len(Xhat)), replace=False))
                Xhat -= drop
                pool = [e for e in range(m) if e not in req]
                Xhat |= set(rng.choice(pool, size=min(nswap, len(pool)), replace=False))
            eta = min(len(req), len(set(req) ^ Xhat))
            raw = build_pieces(A, cost, sorted(Xhat))
            fixed = enforce_B(raw, cost)
            aR, bR = check(raw); aF, bF = check(fixed)
            b_ice, _ = ice(A, cost, req, [x["S"] for x in fixed], [x["c"] for x in fixed])
            c_ice = float(cost[b_ice == 1].sum())
            b_base, _ = online_greedy(A, cost, req, np.zeros(n, dtype=int))
            c_base = float(cost[b_base == 1].sum())
            # measured alpha on the merged pieces, against exact ILP partial optima
            alphas = []
            rem = sorted(Xhat)
            for x in fixed:
                C1 = ilp_partial(A, cost, rem, min(x["need_at_start"], len(rem)))
                if C1 and C1 > 0: alphas.append(x["c"]/C1)
                rem = [e for e in rem if e not in set(x["X"])]
                if not rem: break
            rows.append({"m": m, "corrupt": corrupt, "eta": eta,
                         "pieces_raw": len(raw), "pieces_after_merge": len(fixed),
                         "A_raw": aR, "B_raw": bR, "A_fixed": aF, "B_fixed": bF,
                         "OPT": round(optcost, 3), "ICE_ratio": round(c_ice/optcost, 4),
                         "baseline_ratio": round(c_base/optcost, 4),
                         "max_alpha": (round(max(alphas), 3) if alphas else None)})
            print("  m=%-3d cor=%.1f eta=%-3d pieces %d->%d  (A,B) raw=(%s,%s) fixed=(%s,%s)  ICE=%.3f base=%.3f alpha<=%s"
                  % (m, corrupt, eta, len(raw), len(fixed), aR, bR, aF, bF,
                     c_ice/optcost, c_base/optcost, rows[-1]["max_alpha"]), flush=True)
    RES["claim2_ice_property_B"] = {
        "rows": rows,
        "B_holds_after_merge_all": all(r["B_fixed"] for r in rows),
        "A_preserved_all": all(r["A_fixed"] for r in rows),
        "B_held_before_merge_any": any(r["B_raw"] for r in rows),
        "max_alpha_overall": max(r["max_alpha"] for r in rows if r["max_alpha"]),
        "n_cells": len(rows)}
    json.dump(RES, open("ice_results.json", "w"), indent=1)


if __name__ == "__main__":
    run(); print("DONE")