import torch


def init_network_weights_from_pretraining(net, dataset, normal_cls, seed):
    """
    Initialize the DeepSVDD or DeepMSVDD encoder network weights from the encoder weights of the pretraining CAE.
    Similar to https://github.com/lukasruff/Deep-SAD-PyTorch/blob/master/src/DeepSAD.py#L116
    """
    net_dict = net.state_dict()
    ae_net_dict = torch.load('./trained_models/{}_{}_{}.pt'.format(dataset, normal_cls, seed))
    # Filter out decoder network keys
    ae_net_dict = {k: v for k, v in ae_net_dict.items() if k in net_dict}
    # Overwrite values in the existing state_dict
    net_dict.update(ae_net_dict)
    # Load the new state_dict
    net.load_state_dict(net_dict)
    return net