"""Two-tower click model trained on events before cutoff: user = position-weighted sum of
history item embeddings; score = <user, item_out> + item bias. Trained with BCE per candidate
plus an in-event softmax term (with a learned 'no click' option)."""
import torch, numpy as np, polars as pl, time
torch.set_num_threads(4)
H=20
def pack(E):
    """events -> padded history matrix [n,H] (newest last, right-aligned), mask, plus delta-time buckets"""
    n=len(E); M=np.full((n,H),-1,np.int64); DT=np.zeros((n,H),np.float32)
    hl=E['h'].to_list(); htl=E['h_ts'].to_list(); ts=E['ts'].to_numpy()
    for i,(h,t) in enumerate(zip(hl,htl)):
        if h:
            h=h[-H:]; t=t[-H:]; M[i,H-len(h):]=h; DT[i,H-len(h):]=np.log1p((ts[i]-np.array(t))/3600.)
    return M,DT
class TT(torch.nn.Module):
    def __init__(s,NI,d=64):
        super().__init__()
        s.ein=torch.nn.Embedding(NI+1,d,padding_idx=NI); s.eout=torch.nn.Embedding(NI,d); s.b=torch.nn.Embedding(NI,1)
        s.pw=torch.nn.Parameter(torch.zeros(H)); s.tw=torch.nn.Parameter(torch.zeros(1)); s.b0=torch.nn.Parameter(torch.tensor(-2.7))
        s.null=torch.nn.Linear(d,1)
        for e in (s.ein,s.eout): torch.nn.init.normal_(e.weight,0,0.05)
        torch.nn.init.zeros_(s.b.weight); s.NI=NI
    def user(s,M,DT):
        mask=(M>=0); Mi=torch.where(mask,M,torch.full_like(M,s.NI))
        w=torch.exp(s.pw[None,:]-torch.nn.functional.softplus(s.tw)*DT)*mask
        u=(s.ein(Mi)*w[...,None]).sum(1)/(w.sum(1,keepdim=True)+1e-6)
        return u
    def forward(s,M,DT,C):
        u=s.user(M,DT); return (u[:,None,:]*s.eout(C)).sum(-1)+s.b(C)[...,0]+s.b0, s.null(u)[:,0]
def make_xy(E,R,with_y=True):
    R=R.join(E.select('eid').with_row_index('ev'),on='eid').sort('ev')
    g=R.group_by('ev',maintain_order=True).agg(pl.col('ii'),*( [pl.col('clicked')] if with_y else []))
    C=np.array(g['ii'].to_list(),np.int64); Y=np.array(g['clicked'].to_list(),np.float32) if with_y else None
    return C,Y,R
def train(P_ev,P_rows,NI,epochs=3,d=64,bs=512,lr=3e-3,wd=1e-6,seed=0,log=print):
    torch.manual_seed(seed); np.random.seed(seed)
    M,DT=pack(P_ev); C,Y,_=make_xy(P_ev,P_rows)
    M,DT,C,Y=map(torch.from_numpy,(M,DT,C,Y))
    m=TT(NI,d); opt=torch.optim.Adam(m.parameters(),lr=lr)
    n=len(M); t=time.time()
    for ep in range(epochs):
        perm=torch.randperm(n); tot=0
        for k in range(0,n,bs):
            idx=perm[k:k+bs]; s,nl=m(M[idx],DT[idx],C[idx]); y=Y[idx]
            bce=torch.nn.functional.binary_cross_entropy_with_logits(s,y)
            logits=torch.cat([s,nl[:,None]],1); tgt=torch.where(y.sum(1)>0,y.argmax(1),torch.full((len(y),),10))
            ce=torch.nn.functional.cross_entropy(logits,tgt)
            reg=wd*(m.ein.weight[M[idx].clamp(min=0)].pow(2).sum()+m.eout.weight[C[idx]].pow(2).sum())/len(idx)
            loss=bce+0.1*ce+reg; opt.zero_grad(); loss.backward(); opt.step(); tot+=loss.item()*len(idx)
        log(f'tt epoch {ep} loss {tot/n:.4f} {time.time()-t:.0f}s')
    return m
@torch.no_grad()
def score(m,T_ev,T_rows):
    """returns T_rows (eid,ii) joined with tt_score, tt_null, tt_unorm"""
    m.eval(); M,DT=pack(T_ev); C,_,R=make_xy(T_ev,T_rows,with_y=False)
    s,nl=m(torch.from_numpy(M),torch.from_numpy(DT),torch.from_numpy(C)); u=m.user(torch.from_numpy(M),torch.from_numpy(DT))
    out=pl.DataFrame({'eid':np.repeat(T_ev['eid'].to_numpy(),10),'ii':C.ravel().astype(np.int32),'tt_score':s.numpy().ravel(),
                      'tt_null':np.repeat(nl.numpy(),10),'tt_unorm':np.repeat(u.norm(dim=1).numpy(),10)})
    return out