"""Stream a concept--episode random walk over native nonnegative addresses.

This is an engineering adaptation of associative indexing, not HippoRAG's
OpenIE graph or a proof of biological indexing. Records contain model-derived
coordinates only. Physical weights, source text and generated labels are not
inputs. The caller supplies one immutable, temporally filtered snapshot.
"""
import math
import torch


def concept_entries(codes,ports,vocab_size,device):
    indices=[];values=[];present=0
    for i,name in enumerate(ports):
        if name not in codes:continue
        code=codes[name];idx=code.indices.to(device=device,dtype=torch.long)
        value=code.coefficients.to(device=device,dtype=torch.float32)
        if (idx<0).any() or (idx>=vocab_size).any() or not torch.isfinite(value).all() or (value<0).any():
            raise ValueError('invalid associative concept coordinates')
        total=value.sum(-1,keepdim=True);valid=total.squeeze(-1)>0
        if not valid.any():continue
        # Every active position has equal mass; long episodes do not get an
        # automatic popularity advantage simply because they contain more rows.
        weights=value[valid]/total[valid]
        coordinates,inverse=torch.unique((idx[valid]+i*vocab_size).flatten(),return_inverse=True)
        mass=torch.zeros(len(coordinates),device=device)
        mass.scatter_add_(0,inverse,(weights/valid.sum()).flatten())
        indices.append(coordinates);values.append(mass)
        present+=1
    if not present:
        return torch.empty(0,device=device,dtype=torch.long),torch.empty(0,device=device)
    return torch.cat(indices),torch.cat(values)/present


def concept_distribution(codes,ports,vocab_size,device):
    indices,values=concept_entries(codes,ports,vocab_size,device)
    return torch.zeros(len(ports)*vocab_size,device=device).scatter_add_(0,indices,values)


@torch.no_grad()
def associative_query(records,query,ports,vocab_size,*,steps,restart=.5,device='cpu'):
    """Return concept arrival mass per degree; revisit all episodes each step.

    For episode distributions A[:,j], T=A diag(1/degree) is applied as two
    sparse passes. No episode-by-episode adjacency matrix is constructed.
    Restart anchors each round to the original query; dangling mass returns
    there too. `records()` must yield the same eligible snapshot every time.
    """
    if type(steps) is not int or steps<1 or not math.isfinite(restart) or not 0<restart<=1:
        raise ValueError('positive association steps and restart in (0,1] required')
    seed=concept_distribution(query,ports,vocab_size,device)
    degree=torch.zeros_like(seed);count=0
    for codes in records():
        indices,mass=concept_entries(codes,ports,vocab_size,device)
        degree.scatter_add_(0,indices,mass);count+=1
    current=seed;history=[]
    for _ in range(steps):
        normalized=torch.where(degree>0,current/degree.clamp_min(1e-30),0.)
        updated=torch.zeros_like(current);visited=0
        for codes in records():
            indices,mass=concept_entries(codes,ports,vocab_size,device)
            arrival=torch.dot(mass,normalized[indices])
            updated.scatter_add_(0,indices,mass*arrival);visited+=1
        if visited!=count:raise RuntimeError('associative index snapshot changed during query')
        dangling=current[degree==0].sum()
        updated=restart*seed+(1-restart)*(updated+dangling*seed)
        history.append(float((updated-current).abs().sum()));current=updated
    score_vector=torch.where(degree>0,current/degree.clamp_min(1e-30),0.)
    return score_vector,dict(steps=steps,restart=restart,eligible_units=count,
        concept_state_dimensions=len(seed),iteration_l1_changes=history,
        indexed_weight_tensors_loaded=0,source='native_nonnegative_coordinates',
        algorithm='query-anchored concept-episode diffusion',
        scope='Associative candidate routing, not ordered hash-chain execution or a recall guarantee.')