verbose = True  # Need to be defined early on (for @bd.on of components)

from .matlab_engine import Matlab

from .util import (
    null_function,
    apply,
    Singleton,
    Null,
    interrupt_guard,
    filelock,
    SignalHandler,
    identity,
    make_grid,
    out_size,
    in_size,
    kernel_size,
    stride_size,
    padding_size,
    dilation_size,
    find_layers,
    Average,
    MeanVar,
    slide_window_,
    re_stride,
    moving_avg,
    moving_var,
    sub_avg,
    sub_var,
    has_nan,
    has_inf,
    replace_specials_,
    replace_inf_,
    replace_nan_,
    map_range,
    str_or_none,
    str2bool,
    str_is_int,
    is_tensor,
    is_cuda,
    is_array,
    to_array,
    to_tensor,
    permute,
    hwc2chw,
    chw2hwc,
    channel_flip,
    rgb2bgr,
    bgr2rgb,
    change_view,
    cv2torch,
    torch2cv,
    cv2plt,
    plt2cv,
    plt2torch,
    torch2plt,
    replicate,
    determine_view,
    view,
    default_view,
    forward_hook,
    backward_hook,
    parameters_hook,
    visualize_parameters,
    main_file_path,
    split_path,
    make_dir,
    copy_file_to_dir,
    process_path,
    write_string_to_file,
    number_file_if_exists,
    index_gauss,
    slice_gauss,
    index_uniform,
    slice_uniform,
    Timer,
    recurse_get_elements,
    timestamp,
    loader_process_pool,
    once,
    once_property,
)


from .io import (
    replicate_std_stream,
    silent,
    write,
    log,
    warn,
    error,
    print_model_info,
    print_model_cfg,
    print_separator,
    CSVLogger,
    BoardomLogger,
    compress,
    decompress,
    pack,
    unpack,
    write_pfm,
    load_pfm,
    imwrite,
    imread,
    load_encoded,
    decode_loaded,
    imshow,
    Video,
    VideoDisplay,
)

from .config import Config, autoconfig

from .data import (
    PlacesDataset,
    LoadedDataset,
    DirectoryDataset,
    grow_dataset,
    ListDataset,
)

from .plot import plot_csv

from .models import (
    Module,
    magic_module,
    magic_builder,
    magic_off,
    BaseCollection,
    Sequential,
    Split,
    Branch,
    Sum,
    Cat,
    Prod,
    CropCat2d,
    SingleSelect,
    MultiSelect,
    Parallel,
    Map,
    Pad2dSame,
    pad2dsame,
    Pad2dMultiple,
    pad2dmultiple,
    UnPad2dMultiple,
    unpad2dmultiple,
    Resize,
    conv,
    convsame,
    conv2dsame,
    Conv2dSame,
    ConvTranspose2dSame,
    SelfAttention,
    freeze_bn_running_stats,
    unfreeze_bn_running_stats,
    box_filter1d,
    box_filter2d,
    box_filternd,
    guided_filter,
    multi_guided_filter,
    GuidedFilter,
    LearnedMultiScaleGuidedFilter,
    PretrainedResnet,
    pretrained_resnet_layers,
    determine_channel_sizes,
    OriginalUNet,
    pretty_print,
    module_is_conv,
    module_is_linear,
    module_is_norm,
    param_is_bias,
    param_is_initialized,
    param_is_not_none,
    param_is_not_frozen,
    param_is_weight,
    Initializer,
    InitUniform,
    InitNormal,
    InitConstant,
    InitOnes,
    InitZeros,
    InitEye,
    InitDirac,
    InitXavierUniform,
    InitXavierNormal,
    InitKaimingUniform,
    InitKaimingNormal,
    InitOrthogonal,
    InitSparse,
    InitSELU,
    gaussian_kernel_nd,
    gaussian_blur2d,
    GaussianBlur2d,
    psnr,
    PSNR,
    SSIM,
    MultiLoss,
    CosineLoss,
    SSIMLoss,
    GramLoss,
    GramMatrix,
    AnisotropicTotalVariation,
    IsotropicTotalVariation,
    TotalVariation,
    StyleLoss,
    FeatureLoss,
    PerceptualLoss,
    FeatureMatchingLoss,
    PiecewiseMLP,
    Interp1d,
    is_frozen,
    is_trainable,
    freeze,
    frozen_parameters,
    named_frozen_parameters,
    frozen_modules,
    named_frozen_modules,
    trainable_parameters,
    named_trainable_parameters,
    count_parameters,
    spectral_norm,
    remove_spectral_norm,
)

from .external import accuracy, compose, fft_conv2d, FFTConv2d

from .multiprocessing import PersistentProcessPool, is_main_process, only_main_process

from .engine import (
    Engine,
    State,
    KeepDict,
    on,
    every,
    middleware,
    EventReturnValue,
    CleanupState,
)

from .components import (
    StateUtils,
    Device,
    SGDTrainer,
    DefaultSGDTrainer,
    SGDValidator,
    DefaultSGDValidator,
    Checkpoint,
    LoggerEngine,
    ImageSampler,
)

from .assets import assets

from .imaging import (
    culling,
    exposure,
    reinhard,
    drago,
    mantiuk,
    imfft,
    imnormalize,
    imdenormalize,
    normalize_torchvision_imagenet,
    denormalize_torchvision_imagenet,
    luminance,
    batch_luminance,
    resize_keep_ratio,
    pu_encode,
    #  rgb2xyz,
    #  batch_rgb2xyz,
    #  xyz2rgb,
    #  batch_xyz2rgb,
    #  xyz2lab,
    #  batch_xyz2lab,
    #  lab2xyz,
    #  batch_lab2xyz,
    #  lab2rgb,
    #  batch_lab2rgb,
    #  rgb2lab,
    #  batch_rgb2lab,
)

from .apis import (
    loader,
    torchvision_dataset,
    optimizer,
    scheduler,
    initialize,
    activation,
    normalization,
    residual,
    block,
    pretrained_unet,
    custom_unet,
    pretrained_resnet,
)

from .version import __version__


# Needs to be here for reload() to work well for tests
_CFG_OBJ_LIST = []
cfg = Config()
setup = cfg.setup
boardom_logger = BoardomLogger()
