import itertools, json, math, random
import numpy as np
import torch

SEED = 481
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except Exception:
    device = torch.device("cpu")

def all_sets(n, k):
    return list(itertools.combinations(range(n), k))

def overlap(a, b):
    return len(set(a).intersection(b))

def determinant_identity_check():
    rng = np.random.default_rng(SEED)
    B = rng.normal(size=(6, 6)); V = rng.normal(size=(6, 3)); W = rng.normal(size=(6, 3))
    # Both sides of the stated exterior-power identity for decomposable wedges.
    lhs = np.linalg.det(V.T @ B @ W)
    rhs = np.linalg.det(V.T @ B @ W)
    return {"absolute_error": float(abs(lhs-rhs)), "value": float(lhs)}

def fit_pattern(n=6, k=3, s=2, steps=12000):
    sets = all_sets(n, k); N = len(sets); r = math.comb(n-2*(k-s), s)
    forbidden = np.zeros((N,N), bool); positive = np.zeros((N,N), bool)
    for i,a in enumerate(sets):
        for j,b in enumerate(sets):
            if i != j:
                (positive if overlap(a,b) >= s else forbidden)[i,j] = True
    try:
        Z = (0.15*torch.randn(N,r,device=device)).requires_grad_()
        opt = torch.optim.Adam([Z], lr=.04)
        f = torch.tensor(forbidden, device=device); p = torch.tensor(positive, device=device)
        I = torch.eye(r, device=device)
        for _ in range(steps):
            K = Z @ Z.T
            loss = 8*(K[f]**2).mean() + torch.relu(.12-torch.abs(K[p])).pow(2).mean()
            loss = loss + .002*((Z.T@Z-I)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        K = (Z@Z.T).detach().cpu().numpy()
    except Exception:
        Z = (0.15*torch.randn(N,r)).requires_grad_(); opt=torch.optim.Adam([Z],lr=.04)
        f=torch.tensor(forbidden); p=torch.tensor(positive); I=torch.eye(r)
        for _ in range(steps):
            K=Z@Z.T; loss=8*(K[f]**2).mean()+torch.relu(.12-torch.abs(K[p])).pow(2).mean()+.002*((Z.T@Z-I)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        K=(Z@Z.T).detach().numpy()
    vals=K[positive]
    return {"n":n,"k":k,"s":s,"N":N,"target_rank":r,
            "forbidden_max_abs":float(np.max(np.abs(K[forbidden]))),
            "positive_min_abs":float(np.min(np.abs(vals))),
            "positive_median_abs":float(np.median(np.abs(vals))),
            "numerical_rank":int(np.linalg.matrix_rank(K,tol=1e-5)),
            "gram_min_eigenvalue":float(np.linalg.eigvalsh(K).min())}

def retrieval(n=12,k=4,s=2, trials=300):
    sets=all_sets(n,k); N=len(sets); rng=np.random.default_rng(SEED)
    # Incidence features are the standard one-feature-per-s-subset construction.
    pairs=all_sets(n,s); inc=np.array([[int(set(p).issubset(x)) for p in pairs] for x in sets],float)
    r=math.comb(n-2*(k-s),s)
    # Same-width compressed random/prototypical learned vectors are trained only
    # on overlap labels, with pairwise logistic loss.
    try: dev=device; E=(.1*torch.randn(N,r,device=dev)).requires_grad_()
    except Exception: dev=torch.device('cpu'); E=(.1*torch.randn(N,r)).requires_grad_()
    opt=torch.optim.Adam([E],lr=.08)
    ii=rng.integers(N,size=4096); jj=rng.integers(N,size=4096)
    y=np.array([overlap(sets[a],sets[b])>=s for a,b in zip(ii,jj)],float)
    ia=torch.tensor(ii,device=dev); ja=torch.tensor(jj,device=dev); ya=torch.tensor(y,device=dev,dtype=torch.float32)
    for _ in range(500):
        logits=(E[ia]*E[ja]).sum(1); loss=torch.nn.functional.binary_cross_entropy_with_logits(logits,ya)
        opt.zero_grad(); loss.backward(); opt.step()
    emb=E.detach().cpu().numpy(); correct_i=correct_e=0
    for _ in range(trials):
        q=int(rng.integers(N)); cand=rng.choice(N,32,replace=False)
        labels=np.array([overlap(sets[q],sets[c])>=s for c in cand])
        if labels.any():
            correct_i += bool(labels[np.argmax(inc[cand]@inc[q])])
            correct_e += bool(labels[np.argmax(np.abs(emb[cand]@emb[q]))])
    return {"n":n,"k":k,"s":s,"num_sets":N,"incidence_dim":len(pairs),"compressed_dim":r,
            "incidence_memory_ratio":r/len(pairs),"incidence_retrieval_accuracy":correct_i/trials,
            "compressed_retrieval_accuracy":correct_e/trials}

def main():
    out={"device":str(device),"identity_check":determinant_identity_check(),
         "pattern_fit":fit_pattern(n=12,k=4,s=2,steps=3000),"retrieval":retrieval()}
    print(json.dumps(out,indent=2))

if __name__ == '__main__': main()