"""Encoder construction matching nnssl's ResEncL preset, not a Transformers AutoModel.

Architecture reference: MIC-DKFZ/nnssl, architectures/architecture_registry.py.
Its local source carries CC-BY-SA-4.0; this does not assign a license to weights.
"""
from pathlib import Path
import torch
from torch import nn
from dynamic_network_architectures.architectures.unet import ResidualEncoderUNet
from safetensors.torch import load_model

def build_encoder():
    network = ResidualEncoderUNet(
        input_channels=1, n_stages=6,
        features_per_stage=[32,64,128,256,320,320],conv_op=nn.Conv3d,
        kernel_sizes=[[3,3,3]]*6,
        strides=[[1,1,1]]+[[2,2,2]]*5,
        n_blocks_per_stage=[1,3,4,6,6,6],num_classes=1,
        n_conv_per_stage_decoder=[1]*5,conv_bias=True,
        norm_op=nn.InstanceNorm3d,norm_op_kwargs={'eps':1e-5,'affine':True},
        nonlin=nn.LeakyReLU,nonlin_kwargs={'inplace':True},deep_supervision=False,
    )
    return network.encoder.eval()

def load_encoder(directory):
    encoder=build_encoder()
    load_model(encoder,str(Path(directory)/'encoder.safetensors'),strict=True)
    return encoder.eval()

if __name__=='__main__':
    import sys
    torch.set_num_threads(2)
    model=load_encoder(sys.argv[1])
    with torch.inference_mode():
        out=model(torch.randn(1,1,64,64,64))
    tensors=out if isinstance(out,(list,tuple)) else [out]
    assert all(torch.isfinite(x).all() for x in tensors)
    print([list(x.shape) for x in tensors])