import json, math, random
from pathlib import Path
import numpy as np

SEED=1363

def seed(s=SEED):
    random.seed(s); np.random.seed(s)


def topk_counts(M, rho, trials=20000, seed0=SEED):
    rng=np.random.default_rng(seed0)
    K=int(math.floor(rho*2*M))
    counts=np.empty(trials,dtype=int)
    for t in range(trials):
        scores=rng.standard_normal(2*M)
        counts[t]=np.sum(np.argpartition(scores,-K)[-K:] < M)
    mean=K/2.0
    var=K*(M/(2*M))*(1-M/(2*M))*((2*M-K)/(2*M-1)) if 2*M>1 else 0
    return K, counts, mean, var


def representation_check(M=12, K=9):
    ok=True; examples=[]
    for bits in range(1<<M):
        real=np.array([(bits>>j)&1 for j in range(M)], dtype=int)
        r=int(real.sum())
        if r<=K:
            aug=np.r_[real, np.zeros(M,dtype=int)]
            aug[M:M+(K-r)]=1
            ok &= (aug[:M].tolist()==real.tolist() and int(aug.sum())==K)
            if r in (0,K) and len(examples)<2: examples.append((r,int(aug[M:].sum())))
    return ok, examples


def score_training(d=20, n=600, rho_aug=.30, steps=180, seed0=0, double=True):
    import torch
    torch.manual_seed(seed0); np.random.seed(seed0)
    X=torch.randn(n,d)
    teacher=torch.randn(d,1)
    y=(X@teacher + .25*torch.randn(n,1)>0).float().squeeze()
    W=torch.randn(d,1)/math.sqrt(d)
    M=d
    K=max(1,int(math.floor(rho_aug*(2*M if double else M))))
    sr=torch.randn(M,requires_grad=True)
    sd=torch.randn(M,requires_grad=True) if double else None
    opt=torch.optim.SGD([sr]+([] if sd is None else [sd]),lr=.35)
    for _ in range(steps):
        opt.zero_grad()
        sa=torch.cat([sr,sd]) if double else sr
        hard=torch.zeros_like(sa); hard[torch.topk(sa,K).indices]=1
        h=(hard-sa).detach()+sa
        he=h[:M]
        pred=(X*(W.squeeze()*he)).sum(1)
        loss=torch.nn.functional.binary_cross_entropy_with_logits(pred,y)
        loss.backward(); opt.step()
    with torch.no_grad():
        sa=torch.cat([sr,sd]) if double else sr
        ind=torch.topk(sa,K).indices
        real=(ind<M)
        r=int(real.sum()); dummy=int((~real).sum())
        mask=torch.zeros(M)
        mask[ind[ind<M]]=1
        pred=(X*(W.squeeze()*mask)).sum(1)
        acc=float(((pred>0)==(y>.5)).float().mean())
    return {'acc':acc,'r':r,'dummy':dummy,'eff_density':r/M,'K':K}


def main():
    seed(); out={'predictions':{},'representation':{},'mini_experiment':{}}
    rows=[]
    for M in (20,50,100):
        for rho in (.2,.5,.8):
            K,c,mu,var=topk_counts(M,rho)
            rows.append({'M':M,'rho_aug':rho,'K':K,'pred_mean_r':mu,'obs_mean_r':float(c.mean()),
                         'pred_std_r':math.sqrt(var),'obs_std_r':float(c.std()),
                         'pred_eff_density':mu/M,'obs_eff_density':float(c.mean()/M),
                         'pred_dummy_fraction':.5,'obs_dummy_fraction':float(1-c.mean()/K)})
    out['predictions']['iid_topk_sweep']=rows
    ok,ex=representation_check()
    out['representation']={'M':12,'K':9,'all_masks_at_most_K_represented':ok,'examples_r_dummy':ex}
    reps=[]
    for s in range(5):
        reps.append({'seed':s,'baseline':score_training(rho_aug=.30,seed0=s,double=False),
                     'double':score_training(rho_aug=.30,seed0=s,double=True)})
    out['mini_experiment']['runs']=reps
    for kind in ('baseline','double'):
        vals=[x[kind]['acc'] for x in reps]; dens=[x[kind]['eff_density'] for x in reps]
        out['mini_experiment'][kind]={'mean_acc':float(np.mean(vals)),'std_acc':float(np.std(vals)),
             'mean_eff_density':float(np.mean(dens)),'std_eff_density':float(np.std(dens)), 'K':reps[0][kind]['K']}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()