"""Exact Top-k expected accuracy for shared-variance normal observations.

After removing the common quadratic term, posterior class ranks are ranks of
affine functions of one observation. For each class, every pairwise crossing
changes its rank once. Integrate its normal law over all Top-k intervals.
No hypothetical draws, observation bins, or future measurements are inputs.
"""
from __future__ import annotations
import torch


def expected_topk(prior,mean,variance,k=3):
    assert prior.device.type=='cuda' and prior.ndim==2 and mean.ndim==1
    n,h=prior.shape;assert len(mean)==h and variance.ndim==0
    k=min(k,h)
    m=mean/variance
    intercept=prior.clamp_min(1e-300).log()-mean.square()[None]/(2*variance)
    # Row i against competitor j. Positive slope difference means i overtakes j.
    difference=m[:,None]-m[None,:]
    nonzero=difference!=0
    crossing=(intercept[:,None,:]-intercept[:,:,None])/torch.where(nonzero,difference,1)[None]
    crossing=crossing.masked_fill(~nonzero[None],torch.inf)
    points,order=crossing.sort(-1)
    change=torch.where(difference>0,-1,1).long()[None].expand(n,-1,-1)
    change=torch.where(nonzero[None],change,0).gather(-1,order)
    # At -infinity, lower-slope lines win. Equal-slope ties use a canonical index;
    # their combined expected value is invariant because densities/weights match.
    index=torch.arange(h,device=prior.device)
    equal_better=(~nonzero[None])&((intercept[:,None,:]>intercept[:,:,None])|
        ((intercept[:,None,:]==intercept[:,:,None])&(index[None,None,:]<index[None,:,None])))
    initial=(difference>0).sum(-1)[None]+equal_better.sum(-1)
    rank=torch.cat((initial[:,:,None],initial[:,:,None]+change.cumsum(-1)),-1)
    boundaries=torch.cat((torch.full((n,h,1),-torch.inf,device=prior.device,dtype=prior.dtype),points,
                          torch.full((n,h,1),torch.inf,device=prior.device,dtype=prior.dtype)),-1)
    standardized=(boundaries-mean[None,:,None])/variance.sqrt()
    cdf=torch.special.ndtr(standardized);mass=(cdf[:,:,1:]-cdf[:,:,:-1]).clamp_min(0)
    return (prior*(mass*(rank<k)).sum(-1)).sum(-1)