import json
import itertools
import numpy as np


def make_problem(n, rho, seed=0):
    rng = np.random.default_rng(seed)
    d = rng.uniform(-1.0, 1.0, n)
    Q = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            Q[i, j] = Q[j, i] = 0.35 * rho ** (j-i)
    return d, Q


def objective(mask, d, Q):
    z = np.asarray(mask, dtype=float)
    return float(d @ z + 0.5 * z @ Q @ z)


def exact_optimum(d, Q, k):
    n = len(d); best = float('inf'); best_mask = None
    for comb in itertools.combinations(range(n), k):
        z = np.zeros(n, dtype=np.int8); z[list(comb)] = 1
        val = objective(z, d, Q)
        if val < best:
            best, best_mask = val, z
    return best, best_mask


def boundary_dp(d, Q, k, eta):
    # Approximate decision diagram. The state retains selected count and the
    # last eta decisions; interactions outside that boundary are truncated.
    # Each surviving state also stores a feasible backtracked mask.
    n = len(d); width = min(eta, n)
    states = {(0, 0): (0.0, np.zeros(0, dtype=np.int8))}
    counts = [1]
    for i in range(n):
        nxt = {}
        for (used, hist), (cost, prefix) in states.items():
            for bit in (0, 1):
                nu = used + bit
                if nu > k: continue
                add = d[i] * bit
                for dist in range(1, min(width, i) + 1):
                    if bit and ((hist >> (dist-1)) & 1):
                        add += Q[i-dist, i]
                nh = ((hist << 1) | bit) & ((1 << width) - 1) if width else 0
                key = (nu, nh); val = cost + add
                if key not in nxt or val < nxt[key][0]:
                    nxt[key] = (val, np.append(prefix, bit))
        states = nxt; counts.append(len(states))
    candidates = [(v[0], v[1]) for key, v in states.items() if key[0] == k]
    _, mask = min(candidates, key=lambda x: x[0])
    return objective(mask, d, Q), counts, max(counts), mask


def magnitude(d, Q, k):
    z = np.zeros(len(d), dtype=np.int8)
    z[np.argsort(d)[:k]] = 1
    return objective(z, d, Q)


def path_boundary(n, eta):
    return min(n, 2 * eta)


def grid_boundary(side, eta):
    n = side * side; best = 0
    for cut in range(n + 1):
        near = set()
        for a in range(cut):
            ra, ca = divmod(a, side)
            for b in range(cut, n):
                rb, cb = divmod(b, side)
                if abs(ra-rb) + abs(ca-cb) <= eta:
                    near.update((a, b))
        best = max(best, len(near))
    return best


def main():
    n, k, seed = 20, 10, 7
    d0, Q0 = make_problem(n, 0.8, seed)
    exact, exact_mask = exact_optimum(d0, Q0, k)
    eta_sweep = []
    for eta in range(0, 11):
        approx, counts, peak, mask = boundary_dp(d0, Q0, k, eta)
        eta_sweep.append({'eta': eta, 'full_objective_gap': approx-exact,
                          'peak_states': peak, 'predicted_boundary': path_boundary(n, eta),
                          'theorem_upper_bound': (n+1) * 2**path_boundary(n, eta)})

    rho_sweep = []
    for rho in [0.0, 0.2, 0.4, 0.6, 0.8, 0.9]:
        d, Q = make_problem(n, rho, seed)
        ex, _ = exact_optimum(d, Q, k)
        ap, _, peak, _ = boundary_dp(d, Q, k, 3)
        rho_sweep.append({'rho': rho, 'eta': 3, 'full_objective_gap': ap-ex,
                          'peak_states': peak, 'omitted_tail_scale': rho**4})

    size_sweep = []
    for nn in [10, 14, 18, 22, 26, 30, 40, 50]:
        dd, QQ = make_problem(nn, 0.8, seed); kk = nn // 2
        _, _, peak, _ = boundary_dp(dd, QQ, kk, 3)
        size_sweep.append({'n': nn, 'eta': 3, 'peak_states': peak,
                           'states_per_n': peak / nn})

    graph = [{'eta': eta, 'path_boundary': path_boundary(n, eta),
              'grid8_boundary': grid_boundary(8, eta), 'dense_boundary': n,
              'path_state_bound': (n+1)*2**path_boundary(n, eta),
              'grid_state_bound': (n+1)*2**grid_boundary(8, eta)}
             for eta in [1, 2, 3, 4]]
    baseline = magnitude(d0, Q0, k)
    out = {'config': {'n': n, 'k': k, 'seed': seed},
           'exact_objective': exact, 'magnitude_objective': baseline,
           'magnitude_gap': baseline-exact, 'eta_sweep': eta_sweep,
           'rho_sweep': rho_sweep, 'size_sweep': size_sweep,
           'graph_boundary_sweep': graph}
    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()