import json, math, random, time
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED=251
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 entropy(p):
    if p<=0 or p>=1: return 0.0
    return -p*math.log2(p)-(1-p)*math.log2(1-p)
def insertion_capacity(delta):
    return (1+delta)*(1-entropy(delta/(1+delta)))

def markov_sample(n,q,rng):
    z=np.empty(n,dtype=np.int64); z[0]=rng.integers(2)
    for i in range(1,n): z[i]=z[i-1] if rng.random()>=q else 1-z[i-1]
    return z

def is_subsequence(c,y):
    j=0
    for x in y:
        if j<len(c) and c[j]==x: j+=1
    return j==len(c)

def insert_bits(c,delta,rng):
    k=int(round(delta*len(c))); positions=sorted(rng.choice(len(c)+k,k,replace=False)) if k else []
    out=[]; j=0; posset=set(positions)
    for i in range(len(c)+k):
        if i in posset: out.append(int(rng.integers(2)))
        else: out.append(int(c[j])); j+=1
    return np.array(out,dtype=np.int64)

def codebook(kind,n,M,q,rng):
    if kind=='iid': return [rng.integers(2,size=n,dtype=np.int64) for _ in range(M)]
    return [markov_sample(n,q,rng) for _ in range(M)]

def list_probe(codes,delta,trials,rng):
    lists=[]
    for _ in range(trials):
        c=codes[int(rng.integers(len(codes)))]; y=insert_bits(c,delta,rng)
        lists.append(sum(is_subsequence(x,y) for x in codes))
    return float(np.mean(lists)), int(np.max(lists)), float(np.mean(np.array(lists)>1))

def make_data(codes,delta,count,rng):
    xs=[]; ys=[]
    for _ in range(count):
        label=int(rng.integers(len(codes))); xs.append(insert_bits(codes[label],delta,rng)); ys.append(label)
    return torch.tensor(np.stack(xs),dtype=torch.long), torch.tensor(ys,dtype=torch.long)

class Decoder(nn.Module):
    def __init__(self,M,h=48):
        super().__init__(); self.emb=nn.Embedding(2,16); self.rnn=nn.GRU(16,h,batch_first=True); self.fc=nn.Linear(h,M)
    def forward(self,x):
        _,h=self.rnn(self.emb(x)); return self.fc(h[-1])

def train_eval(codes,delta,rng):
    # Fixed small setup, identical optimizer and number of updates.
    global device
    tr_x,tr_y=make_data(codes,delta,3200,rng)
    te={d:make_data(codes,d,1000,rng) for d in [0,.25,.5]}
    def run(dev):
        model=Decoder(len(codes)).to(dev)
        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        model.train(); bs=64
        for epoch in range(12):
            perm=torch.randperm(len(tr_y))
            for st in range(0,len(tr_y),bs):
                ix=perm[st:st+bs]
                opt.zero_grad(set_to_none=True)
                loss=nn.functional.cross_entropy(model(tr_x[ix].to(dev)),tr_y[ix].to(dev))
                loss.backward(); opt.step()
        out={}; model.eval()
        with torch.no_grad():
            for d,(x,y) in te.items():
                pred=model(x.to(dev)).argmax(1).cpu()
                out[str(d)]=float((pred==y).float().mean())
        return out
    try:
        return run(device)
    except Exception as exc:
        print('device failure, falling back to CPU:', repr(exc))
        if device.type == 'cuda':
            torch.cuda.empty_cache()
        device=torch.device('cpu')
        return run(device)

def main():
    delta=.25; n=12; M=8; q=.25; rng=np.random.default_rng(SEED)
    cap=insertion_capacity(delta); rate=math.log2(M)/n
    mathcheck={'delta':delta,'capacity_formula':cap,'code_rate':rate,'safe_rate_(eps=.1)':.9*cap,'rate_satisfies_margin':rate<=.9*cap}
    # Several independently drawn codebooks reduce dependence on one random draw.
    rows=[]
    for kind,kq in [('iid',.5),('markov',.25),('markov',.4)]:
        lp=[]; acc=[]; t0=time.time()
        for rep in range(2):
            rr=np.random.default_rng(SEED+100*rep+(0 if kind=='iid' else 1))
            codes=codebook(kind,n,M,kq,rr)
            lp.append(list_probe(codes,delta,800,rr))
            acc.append(train_eval(codes,delta,rr))
        rows.append({'kind':kind,'q':kq,'list_mean':float(np.mean([x[0] for x in lp])),'list_max_mean':float(np.mean([x[1] for x in lp])),'collision_fraction':float(np.mean([x[2] for x in lp])),'accuracy':{d:float(np.mean([a[d] for a in acc])) for d in ['0','0.25','0.5']},'seconds':time.time()-t0})
    result={'seed':SEED,'device':str(device),'mathcheck':mathcheck,'results':rows,'notes':'iid q=.5 is the symmetric iid Bernoulli source; Markov q=.25 has positive persistence. Lists are sampled by inserting random bits and counting codewords that are subsequences.'}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()