"""Matched compact direct classifiers, including source-only calibration."""
from __future__ import annotations
import torch
from torch import nn
from .factorial import hypotheses


def train(x,h,sizes,seed,nonlinear,steps=400):
    torch.cuda.manual_seed(seed);d=x.shape[1];out=sum(sizes)
    if nonlinear:
        width=(19984-out)//(d+1+out)
        net=nn.Sequential(nn.Linear(d,width,device=x.device),nn.SiLU(),nn.Linear(width,out,device=x.device))
    else:net=nn.Linear(d,out,device=x.device)
    opt=torch.optim.AdamW(net.parameters(),lr=.003,weight_decay=.01)
    counts=[torch.bincount(h[:,i],minlength=k).double()+1 for i,k in enumerate(sizes)]
    with torch.enable_grad():
        for _ in range(steps):
            opt.zero_grad(set_to_none=True);logits=net(x.float());loss=x.new_zeros(());offset=0
            for i,size in enumerate(sizes):
                weights=(1/counts[i]).float();weights/=weights.mean()
                loss=loss+nn.functional.cross_entropy(logits[:,offset:offset+size],h[:,i],weight=weights)/len(sizes);offset+=size
            loss.backward();opt.step()
    return net


def fit(x,h,sizes,kind,seed,steps=400):
    nonlinear=kind=='factorized_classifier';n=len(x);out=sum(sizes)
    network=train(x,h,sizes,seed,nonlinear,steps)
    oof=x.new_zeros((n,out));fold=torch.arange(n,device=x.device)%4
    for f in range(4):
        partial=train(x[fold!=f],h[fold!=f],sizes,seed+101+f,nonlinear,steps)
        with torch.no_grad():oof[fold==f]=partial(x[fold==f].float()).double()
    temperatures=[];offset=0
    for f,size in enumerate(sizes):
        options=x.new_tensor([.5,1.,2.,4.]);logits=oof[:,None,offset:offset+size]/options[None,:,None]
        loss=-logits.log_softmax(-1).gather(-1,h[:,f,None,None].expand(-1,len(options),1)).squeeze(-1).mean(0)
        temperatures.append(options[loss.argmin()]);offset+=size
    parameters=sum(p.numel() for p in network.parameters());assert parameters+len(sizes)<=20000
    return dict(kind=kind,sizes=tuple(sizes),state_dict=network.state_dict(),
        nonlinear=nonlinear,input_dimensions=x.shape[1],grid=hypotheses(sizes,x.device),
        temperatures=torch.stack(temperatures),neural_parameters=parameters+len(sizes),adaptation_fits=5)


def posterior(model,x,observed):
    # The registered direct-classifier comparison is full-evidence only.
    assert len(list(observed))==x.shape[1]
    state=model['state_dict']
    if model['nonlinear']:
        hidden=nn.functional.silu(x.float()@state['0.weight'].T+state['0.bias'])
        logits=hidden@state['2.weight'].T+state['2.bias']
    else:logits=x.float()@state['weight'].T+state['bias']
    out=logits.new_zeros((len(x),len(model['grid'])));offset=0
    for f,size in enumerate(model['sizes']):
        logp=(logits[:,offset:offset+size]/model['temperatures'][f]).log_softmax(-1)
        out+=logp[:,model['grid'][:,f]];offset+=size
    return out.double().softmax(-1)