"""Strong independent-factor controls sharing the same numerical evidence."""
from __future__ import annotations
import torch
from torch.nn.functional import one_hot
from .factorial import hypotheses


def fit(x,h,sizes,kind):
    d=x.shape[1];eye=torch.eye(d,device=x.device,dtype=x.dtype);parts=[]
    if kind=='factorized_ridge':
        z=torch.cat((x,x.new_ones((len(x),1))),-1)
        penalty=eye.new_zeros((d+1,d+1));penalty.diagonal().fill_(8.);penalty[-1,-1]=.01
        inverse=torch.linalg.inv(z.T@z+penalty)
    for i,size in enumerate(sizes):
        label=h[:,i];count=torch.bincount(label,minlength=size).double()
        sums=x.new_zeros((size,d)).scatter_add_(0,label[:,None].expand_as(x),x)
        mean=(sums+x.mean(0))/(count[:,None]+1)
        residual=x-mean[label]
        scatter=x.new_zeros((size,d,d)).scatter_add_(0,label[:,None,None].expand(-1,d,d),residual[:,:,None]*residual[:,None,:])
        pool=scatter.sum(0)/max(1,len(x)-size)
        if kind=='factorized_nb':
            var=(scatter.diagonal(dim1=-2,dim2=-1)+4*pool.diagonal()[None])/(count[:,None]+4)
            covariance=torch.diag_embed(var)+.05*eye
        elif kind=='factorized_lda':
            covariance=(.5*pool+.5*torch.diag_embed(pool.diagonal())+.05*eye)[None].expand(size,-1,-1)
        elif kind=='factorized_kernel':
            overall=(x-x.mean(0)).T@(x-x.mean(0))/max(1,len(x)-1)
            covariance=(.5*overall+.5*eye)*.7**2
        elif kind=='factorized_ridge':
            # Class-balanced targets impose the same uniform ontology prior.
            target=one_hot(label,size).double()/(count[None]+1)
            weights=inverse@z.T@target
            covariance=eye
        else:raise ValueError(kind)
        parts.append(dict(mean=mean,cov=covariance,count=count,
            centers=x,labels=label,weights=weights if kind=='factorized_ridge' else None,size=size))
    return dict(parts=parts,sizes=tuple(sizes),kind=kind,grid=hypotheses(sizes,x.device),neural_parameters=0)


def posterior(model,x,observed):
    ix=list(observed);n=len(x);logs=[]
    for part in model['parts']:
        size=part['size']
        if not ix:
            p=x.new_full((n,size),1/size)
        elif model['kind']=='factorized_kernel':
            cov=part['cov'][ix][:,ix];chol=torch.linalg.cholesky(cov)
            means=torch.cat((part['centers'],part['mean']),0)
            delta=x[:,None,ix]-means[None,:,ix]
            white=torch.linalg.solve_triangular(chol,delta.reshape(-1,len(ix)).T,upper=False)
            loglike=-.5*white.square().sum(0).reshape(n,-1)
            result=[]
            for a in range(size):
                components=torch.cat((torch.where(part['labels']==a)[0],part['labels'].new_tensor([len(part['centers'])+a])))
                weight=x.new_ones(len(components));weight[-1]=.5
                result.append(torch.logsumexp(loglike[:,components]+weight.log()[None],-1)-(part['count'][a]+.5).log())
            p=torch.stack(result,-1).softmax(-1)
        elif model['kind']=='factorized_ridge':
            # Missing coordinates contribute no invented observed value; the full
            # data diagnostic comparison uses every coordinate. Acquisition evaluates
            # only world models, rather than claiming these are missing-data trained.
            scores=x[:,ix]@part['weights'][ix]+part['weights'][-1]
            p=scores.clamp_min(1e-4);p/=p.sum(-1,keepdim=True)
        else:
            cov=part['cov'][:,ix][:,:,ix];chol=torch.linalg.cholesky(cov)
            delta=(x[:,None,ix]-part['mean'][None,:,ix]).permute(1,2,0)
            white=torch.linalg.solve_triangular(chol,delta,upper=False)
            mahal=white.square().sum(1).T
            logdet=chol.diagonal(dim1=-2,dim2=-1).log().sum(-1)
            p=(-.5*mahal-logdet[None]).softmax(-1)
        logs.append(p.clamp_min(1e-300).log()[:,model['grid'][:,len(logs)]])
    return torch.stack(logs).sum(0).softmax(-1)