"""Sanity checks of the efficient in-context learners of Section 5 on small random instances:
A_LP (Theorem 5.2) is consistent with every realizable context of BP_k and LD_k, and A_l1 (Theorem 5.5(b))
has empirical error at most that of the best h_w. Exact identities are checked in check_inference_linear_risk.py.
Run: python -I check_learners.py
"""
import itertools
import random

import numpy as np
from scipy.optimize import linprog

random.seed(3)


def rand_sparse(n, k):
    r = random.randint(1, k)
    S = random.sample(range(n), r)
    x = np.zeros(n)
    for i in S:
        x[i] = random.choice((-1, 1))
    return x


def h(cls, sigma, x):
    ip = float(sigma @ x)
    if cls == "BP":
        return int(ip >= 0)
    return int(any(x[i] != 0 and x[i] == sigma[i] for i in range(len(x))))


def a_lp(cls, ctx, n):
    """Solve (2): <v,x>+c(x) >= 0 for y=1 and <= -1 for y=0; return v or None."""
    A, b = [], []
    for x, y in ctx:
        c = 0.0 if cls == "BP" else float(np.count_nonzero(x) - 2)
        if y == 1:
            A.append(-x)
            b.append(c)
        else:
            A.append(x)
            b.append(-1.0 - c)
    res = linprog(np.zeros(n), A_ub=np.array(A), b_ub=np.array(b), bounds=[(None, None)] * n, method="highs")
    return res.x if res.status == 0 else None


def check_lp(trials=300):
    for _ in range(trials):
        cls = random.choice(["BP", "LD"])
        n = random.randint(2, 9)
        k = random.randint(1, n)
        sigma = np.array([random.choice((-1, 1)) for _ in range(n)])
        ctx = []
        for _ in range(random.randint(1, 25)):
            x = rand_sparse(n, k)
            ctx.append((x, h(cls, sigma, x)))
        v = a_lp(cls, ctx, n)
        assert v is not None
        for x, y in ctx:
            c = 0.0 if cls == "BP" else float(np.count_nonzero(x) - 2)
            assert int(v @ x + c >= -1e-7) == y
    print("A_LP: feasible and consistent on", trials, "realizable contexts of BP_k and LD_k")


def fourier_features(R, table, N, kappa):
    feats = {}
    r = len(R)
    for size in range(r + 1):
        for T in itertools.combinations(R, size):
            pos = [R.index(i) for i in T]
            s = sum(table[beta] * (-1) ** sum(beta[p] for p in pos) for beta in itertools.product((0, 1), repeat=r))
            feats[T] = s / 2 ** r
    return feats


def check_l1(trials=60):
    for _ in range(trials):
        N = random.randint(2, 6)
        kappa = random.randint(1, min(3, N))
        index = [T for j in range(kappa + 1) for T in itertools.combinations(range(N), j)]
        pos = {T: i for i, T in enumerate(index)}
        D = len(index)
        domain = []
        for _ in range(8):
            r = random.randint(0, kappa)
            R = tuple(sorted(random.sample(range(N), r)))
            table = {beta: random.randint(0, 1) for beta in itertools.product((0, 1), repeat=r)}
            psi = np.zeros(D)
            for T, f in fourier_features(R, table, N, kappa).items():
                psi[pos[T]] = f
            domain.append((R, table, psi))
        m = random.randint(5, 20)
        ctx = [(random.randrange(8), random.randint(0, 1)) for _ in range(m)]
        # L1 regression: variables v (D) and slacks t (m); minimize sum t, |<v,psi_i> - y_i| <= t_i
        c = np.concatenate([np.zeros(D), np.ones(m)])
        A, b = [], []
        for i, (x, y) in enumerate(ctx):
            psi = domain[x][2]
            row = np.concatenate([psi, -np.eye(m)[i]])
            A.append(row)
            b.append(y)
            row2 = np.concatenate([-psi, -np.eye(m)[i]])
            A.append(row2)
            b.append(-y)
        res = linprog(c, A_ub=np.array(A), b_ub=np.array(b), bounds=[(None, None)] * D + [(0, None)] * m, method="highs")
        v = res.x[:D]
        z = [float(domain[x][2] @ v) for x, _ in ctx]
        cands = [0.0, 1.0] + [zz for zz in z if 0 <= zz <= 1]
        best = min(sum(int(zz >= th) != y for zz, (_, y) in zip(z, ctx)) for th in cands)
        best_h = min(
            sum(domain[x][1][tuple(w[i] for i in domain[x][0])] != y for x, y in ctx)
            for w in itertools.product((0, 1), repeat=N)
        )
        assert best <= best_h + 1e-9, (best, best_h)
    print("A_l1: chosen classifier never makes more context errors than the best h_w, on", trials, "instances")


if __name__ == "__main__":
    check_lp()
    check_l1()