import torch


def load_pretrained(pretrained_network):
    """Loads parameters of neural network pretrained using spectrai.

    Arguments:
        pretrained_network: dictionary containing the following fields
            network: network state_dict
            optimizer: optimizer state_dict
            scheduler: scheduler state_dict
            criterion: criterion
            train_metrics: dictionary of training metrics
            val_metrics: dictionary of validation metrics
            test_metrics: dictionary of test metrics
            achitecture: network achitecture [UNet/ResUNet/ResNet/RCAN]
            task: deep learning task [Calibration/Classification/Denoising/Segmentation/Super-Resolution]
            channels: spectrum length (int)
            normalization: normalization layer [None/BatchNorm/LayerNorm/InstanceNorm/GroupNorm]
            activation: activation function [ReLU/LeakyReLU/PReLU]
            classes: classification/segmentation classes else ignored
            input_image_size: input image size (int)
            target_image_size: target image size (int)
            data_format: data format ['Image: H, W, C'/'Image: C, H, W'/Spectra]
    Returns:
        output: dictionary containing the following fields
            achitecture: network achitecture [UNet/ResUNet/ResNet/RCAN]
            task: deep learning task [Calibration/Classification/Denoising/Segmentation/Super-Resolution]
            channels: spectrum length (int)
            normalization: normalization layer [None/BatchNorm/LayerNorm/InstanceNorm/GroupNorm]
            activation: activation function [ReLU/LeakyReLU/PReLU]
            classes: classification/segmentation classes else ignored
            input_image_size: input image size (int)
            target_image_size: target image size (int)
            data_format: data format ['Image: H, W, C'/'Image: C, H, W'/Spectra]
    """
    input_data = torch.load(pretrained_network)
    architecture = input_data['architecture']
    task = input_data['task']
    channels = input_data['channels']
    normalization = input_data['normalization']
    activation = input_data['activation']
    classes = input_data['classes']
    input_image_size = input_data['input_image_size']
    target_image_size = input_data['target_image_size']
    data_format = input_data['data_format']

    output = {'architecture':architecture, 'task':task, 'channels':channels,
            'normalization':normalization, 'activation':activation, 'classes':classes, 
            'input_image_size':input_image_size, 'target_image_size':target_image_size, 
            'data_format':data_format}

    return output