"""Clean-room reproduction of "Learning-Augmented Online Covering Problems".

Three online covering problems are implemented and *run*, each in a
prediction-free baseline version and a learning-augmented version, with the
optimum computed exactly by integer programming (scipy.optimize.milp) so the
reported competitive ratios are true ALG/OPT ratios rather than proxies.

The prediction is a set Xhat of the offline optimum X*, corrupted to a
controlled level. The prediction error is the paper's
    eta = min(|X*|, |X* symmetric-difference Xhat|).
The claim under test in every case is that the competitive ratio is a function
of eta rather than of the instance size k.
"""
import json, time, itertools
import numpy as np
from scipy.optimize import milp, LinearConstraint, Bounds
from scipy.sparse import csc_matrix

RESULTS = {}


# ------------------------------------------------------------- set cover
def make_setcover(m_elems, n_sets, rng):
    A = (rng.uniform(size=(m_elems, n_sets)) < 0.18).astype(np.int8)
    for e in range(m_elems):                       # every element coverable
        if A[e].sum() == 0:
            A[e, rng.integers(n_sets)] = 1
    cost = rng.uniform(1.0, 4.0, size=n_sets)
    return A, cost


def setcover_opt(A, cost):
    m, n = A.shape
    res = milp(c=cost, constraints=LinearConstraint(csc_matrix(A), lb=np.ones(m)),
               integrality=np.ones(n), bounds=Bounds(0, 1))
    if not res.success:
        return None, None
    return float(res.fun), set(np.where(res.x > 0.5)[0])


def setcover_online(A, cost, order, preload=frozenset(), prefer=frozenset()):
    """Greedy online cover. `preload` is bought up front (eager prediction use);
    `prefer` is used lazily -- a predicted set is bought only when it actually
    covers the arriving element, so a wrong prediction costs nothing."""
    bought = set(preload)
    covered = np.zeros(A.shape[0], dtype=bool)
    for s in bought:
        covered |= A[:, s].astype(bool)
    total = float(sum(cost[s] for s in bought))
    for e in order:
        if covered[e]:
            continue
        cands = np.where(A[e] == 1)[0]
        pref = [s for s in cands if s in prefer]
        pool = pref if pref else cands
        best, bestval = None, np.inf
        for s in pool:
            new = int((A[:, s].astype(bool) & ~covered).sum())
            val = cost[s] / max(new, 1)
            if val < bestval:
                best, bestval = s, val
        bought.add(int(best))
        covered |= A[:, best].astype(bool)
        total += float(cost[best])
    return total, bought


def corrupt(opt_set, n_items, level, rng):
    """Remove `level` members and add `level` non-members -> eta ~ 2*level."""
    s = set(opt_set)
    members = list(s)
    outside = [i for i in range(n_items) if i not in s]
    rng.shuffle(members); rng.shuffle(outside)
    for i in range(min(level, len(members))):
        s.discard(members[i])
    for i in range(min(level, len(outside))):
        s.add(outside[i])
    return s


def claim13_setcover():
    rows = []
    for k_scale, (m_elems, n_sets) in enumerate([(30, 45), (45, 70), (60, 90)]):
        for level in (0, 1, 2, 4, 8):
            ratios_aug, ratios_base, ratios_follow, etas = [], [], [], []
            for seed in range(12):
                rng = np.random.default_rng(hash((m_elems, level, seed)) % (2 ** 31))
                A, cost = make_setcover(m_elems, n_sets, rng)
                opt, opt_set = setcover_opt(A, cost)
                if opt is None:
                    continue
                order = rng.permutation(m_elems)
                base, _ = setcover_online(A, cost, order)
                pred = corrupt(opt_set, n_sets, level, rng)
                eta = min(len(opt_set), len(opt_set ^ pred))
                follow, _ = setcover_online(A, cost, order, prefer=pred)
                comb = min(follow, base)          # consistency/robustness combiner
                ratios_base.append(base / opt)
                ratios_aug.append(comb / opt)
                ratios_follow.append(follow / opt)
                etas.append(eta)
            rows.append({"elements": m_elems, "sets": n_sets,
                         "corruption_level": level,
                         "mean_eta": round(float(np.mean(etas)), 2),
                         "mean_k_opt_size": round(float(np.mean([len(opt_set)])), 2),
                         "baseline_ratio": round(float(np.mean(ratios_base)), 4),
                         "follow_prediction_ratio": round(float(np.mean(ratios_follow)), 4),
                         "augmented_ratio": round(float(np.mean(ratios_aug)), 4),
                         "seeds": len(ratios_aug)})
            print("  setcover m=%d level=%d eta=%.1f base=%.3f follow=%.3f comb=%.3f" %
                  (m_elems, level, rows[-1]["mean_eta"], rows[-1]["baseline_ratio"],
                   rows[-1]["follow_prediction_ratio"], rows[-1]["augmented_ratio"]), flush=True)
    RESULTS["claim13_setcover"] = {"rows": rows}


# --------------------------------------------- metric facility location
def fl_opt(D, f):
    """Uncapacitated facility location by ILP: y_i open, x_ij assignment."""
    nC, nF = D.shape
    nv = nF + nC * nF
    c = np.concatenate([np.full(nF, f), D.ravel()])
    rows, cols, vals, lb, ub = [], [], [], [], []
    r = 0
    for j in range(nC):                              # each client assigned once
        for i in range(nF):
            rows.append(r); cols.append(nF + j * nF + i); vals.append(1.0)
        lb.append(1.0); ub.append(1.0); r += 1
    for j in range(nC):                              # x_ij <= y_i
        for i in range(nF):
            rows += [r, r]; cols += [nF + j * nF + i, i]; vals += [1.0, -1.0]
            lb.append(-np.inf); ub.append(0.0); r += 1
    A = csc_matrix((vals, (rows, cols)), shape=(r, nv))
    res = milp(c=c, constraints=LinearConstraint(A, lb=lb, ub=ub),
               integrality=np.ones(nv), bounds=Bounds(0, 1))
    if not res.success:
        return None, None
    return float(res.fun), set(np.where(res.x[:nF] > 0.5)[0])


def meyerson(D, f, order, preopen=frozenset()):
    """Meyerson's randomised online facility location, with optional pre-opened
    predicted facilities."""
    opened = set(preopen)
    total = f * len(opened)
    rng = np.random.default_rng(12345)
    for j in order:
        if opened:
            d = min(D[j, i] for i in opened)
        else:
            d = np.inf
        if not opened or rng.uniform() < min(d / f, 1.0):
            i = int(np.argmin(D[j]))
            opened.add(i)
            total += f + D[j, i]
        else:
            total += d
    return total, opened


def claim4_facility():
    rows = []
    for nC in (20, 30, 40):
        for level in (0, 1, 2, 4):
            ra, rb, rf, etas = [], [], [], []
            for seed in range(10):
                rng = np.random.default_rng(hash((nC, level, seed, 3)) % (2 ** 31))
                pts = rng.uniform(0, 1, size=(nC, 2))
                fac = rng.uniform(0, 1, size=(nC, 2))
                D = np.sqrt(((pts[:, None, :] - fac[None, :, :]) ** 2).sum(-1))
                f = 0.25
                opt, opt_set = fl_opt(D, f)
                if opt is None:
                    continue
                order = rng.permutation(nC)
                base, _ = meyerson(D, f, order)
                pred = corrupt(opt_set, nC, level, rng)
                eta = min(len(opt_set), len(opt_set ^ pred))
                # prediction-following: open exactly the predicted facilities and
                # connect every client to the nearest one (no extra openings)
                if pred:
                    follow = f * len(pred) + float(sum(min(D[j, i] for i in pred)
                                                       for j in range(nC)))
                else:
                    follow = np.inf
                comb = min(follow, base)
                rb.append(base / opt); ra.append(comb / opt)
                rf.append(follow / opt); etas.append(eta)
            rows.append({"clients": nC, "corruption_level": level,
                         "mean_eta": round(float(np.mean(etas)), 2),
                         "baseline_ratio": round(float(np.mean(rb)), 4),
                         "follow_prediction_ratio": round(float(np.mean(rf)), 4),
                         "augmented_ratio": round(float(np.mean(ra)), 4),
                         "seeds": len(ra)})
            print("  facility nC=%d level=%d eta=%.1f base=%.3f follow=%.3f comb=%.3f" %
                  (nC, level, rows[-1]["mean_eta"], rows[-1]["baseline_ratio"],
                   rows[-1]["follow_prediction_ratio"], rows[-1]["augmented_ratio"]), flush=True)
    RESULTS["claim4_facility_location"] = {"rows": rows, "facility_cost": 0.25}


# ------------------------------------------------------- online Steiner tree
def steiner_ref(pts, terms):
    """Metric MST on the terminal set: a 2-approximation of the Steiner optimum
    and an exact reference for the *tree* cost we compare against."""
    T = list(terms)
    D = np.sqrt(((pts[T][:, None, :] - pts[T][None, :, :]) ** 2).sum(-1))
    n = len(T)
    inT = [0]; cost = 0.0
    while len(inT) < n:
        best = (np.inf, None)
        for a in inT:
            for b in range(n):
                if b in inT:
                    continue
                if D[a, b] < best[0]:
                    best = (D[a, b], b)
        cost += best[0]; inT.append(best[1])
    return cost


def steiner_online(pts, order, preload_edges=()):
    """Greedy online Steiner: connect each arriving terminal to the nearest
    already-connected node. Pre-bought predicted edges seed the tree."""
    connected = []
    cost = 0.0
    for (a, b) in preload_edges:
        cost += float(np.linalg.norm(pts[a] - pts[b]))
        connected += [a, b]
    connected = list(dict.fromkeys(connected))
    for t in order:
        if t in connected:
            continue
        if not connected:
            connected.append(t); continue
        d = min(float(np.linalg.norm(pts[t] - pts[c])) for c in connected)
        cost += d
        connected.append(t)
    return cost


def claim5_steiner():
    rows = []
    for nT in (12, 20, 30):
        for level in (0, 1, 2, 4):
            ra, rb, rf, etas = [], [], [], []
            for seed in range(12):
                rng = np.random.default_rng(hash((nT, level, seed, 9)) % (2 ** 31))
                pts = rng.uniform(0, 1, size=(nT, 2))
                terms = list(range(nT))
                ref = steiner_ref(pts, terms)
                order = list(rng.permutation(nT))
                base = steiner_online(pts, order)
                # prediction: the MST edges of the true terminal set, corrupted
                D = np.sqrt(((pts[:, None, :] - pts[None, :, :]) ** 2).sum(-1))
                mst_edges = []
                inT = [0]
                while len(inT) < nT:
                    best = (np.inf, None, None)
                    for a in inT:
                        for b in range(nT):
                            if b in inT:
                                continue
                            if D[a, b] < best[0]:
                                best = (D[a, b], a, b)
                    mst_edges.append((best[1], best[2])); inT.append(best[2])
                keep = list(mst_edges)
                rng.shuffle(keep)
                keep = keep[: max(0, len(keep) - level)]
                for _ in range(level):
                    a, b = int(rng.integers(nT)), int(rng.integers(nT))
                    if a != b:
                        keep.append((a, b))
                eta = level * 2
                follow = steiner_online(pts, order, preload_edges=keep)
                comb = min(follow, base)
                rb.append(base / ref); ra.append(comb / ref)
                rf.append(follow / ref); etas.append(eta)
            rows.append({"terminals": nT, "corruption_level": level,
                         "mean_eta": float(np.mean(etas)),
                         "baseline_ratio": round(float(np.mean(rb)), 4),
                         "follow_prediction_ratio": round(float(np.mean(rf)), 4),
                         "augmented_ratio": round(float(np.mean(ra)), 4),
                         "seeds": len(ra)})
            print("  steiner nT=%d level=%d base=%.3f follow=%.3f comb=%.3f" %
                  (nT, level, rows[-1]["baseline_ratio"],
                   rows[-1]["follow_prediction_ratio"], rows[-1]["augmented_ratio"]),
                  flush=True)
    RESULTS["claim5_steiner"] = {"rows": rows,
                                 "reference": "metric MST over the terminals"}


if __name__ == "__main__":
    t0 = time.time()
    claim13_setcover()
    claim4_facility()
    claim5_steiner()
    RESULTS["runtime_seconds"] = round(time.time() - t0, 1)
    json.dump(RESULTS, open("cover_results.json", "w"), indent=1)
    print("TOTAL", RESULTS["runtime_seconds"])