#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""Collection of functions and modules to help development in PyTorch."""

__name__ = "torchwrench"
__author__ = "Étienne Labbé (Labbeti)"
__author_email__ = "labbeti.pub@gmail.com"
__license__ = "MIT"
__maintainer__ = "Étienne Labbé (Labbeti)"
__status__ = "Development"
__version__ = "0.8.0"

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    # Re-import torchwrench modules for language servers
    from . import core as core
    from . import extras as extras
    from . import hub as hub
    from . import nn as nn
    from . import optim as optim
    from . import serialization as serialization
    from . import types as types
    from . import utils as utils
    from .core.dtype import (
        complex64,
        complex128,
        double,
        float,
        float16,
        float32,
        float64,
        half,
        int,
        int8,
        int16,
        int32,
        int64,
        long,
        short,
        uint8,
    )
    from .core.dtype_enum import (
        DTypeEnum,
    )
    from .core.make import (
        CUDA_IF_AVAILABLE,
        Device,
        DeviceLike,
        DTypeLike,
        Generator,
        GeneratorLike,
    )
    from .core.version import get_githash_full, get_githash_short
    from .extras.numpy import ndarray_to_tensor, tensor_to_ndarray, to_ndarray
    from .extras.yaml.yaml import (
        dump_yaml,
        dumps_yaml,
        load_yaml,
        loads_yaml,
        read_yaml,
        save_yaml,
    )
    from .hub.download import download_file
    from .hub.paths import get_cache_dir, get_tmp_dir
    from .nn.functional import (
        activity_to_segments,
        activity_to_segments_list,
        all_eq,
        all_ne,
        arange,
        as_device,
        as_dtype,
        as_generator,
        as_tensor,
        average_power,
        cat,
        cat_padded_batch,
        checksum_any,
        concat,
        count_parameters,
        crop_dim,
        crop_dims,
        deep_equal,
        empty,
        equal,
        find,
        flatten,
        full,
        generate_square_subsequent_mask,
        get_default_device,
        get_default_dtype,
        get_default_generator,
        get_inverse_perm,
        get_ndim,
        get_perm_indices,
        get_shape,
        identity,
        index_to_name,
        index_to_onehot,
        indices_to_multihot,
        indices_to_multinames,
        initial_seed,
        insert_at_indices,
        is_complex,
        is_convertible_to_tensor,
        is_floating_point,
        is_full,
        is_sorted,
        is_stackable,
        is_tensor,
        is_unique,
        lengths_to_non_pad_mask,
        lengths_to_pad_mask,
        lengths_to_ratios,
        log_softmax_multidim,
        manual_seed,
        masked_equal,
        masked_mean,
        masked_sum,
        matmul,
        move_to,
        move_to_rec,
        mse,
        multi_indices_to_multihot,
        multi_indices_to_multinames,
        multihot_to_indices,
        multihot_to_multi_indices,
        multihot_to_multinames,
        multilabel_to_powerset,
        multinames_to_indices,
        multinames_to_multi_indices,
        multinames_to_multihot,
        name_to_index,
        name_to_onehot,
        ndim,
        nelement,
        no_grad,
        non_pad_mask_to_lengths,
        non_pad_mask_to_ratios,
        one_hot,
        onehot_to_index,
        onehot_to_name,
        ones,
        pad_and_crop_dim,
        pad_and_stack_rec,
        pad_dim,
        pad_dims,
        pad_mask_to_lengths,
        pad_mask_to_ratios,
        powerset_to_multilabel,
        probs_to_index,
        probs_to_indices,
        probs_to_multi_indices,
        probs_to_multihot,
        probs_to_multinames,
        probs_to_name,
        probs_to_onehot,
        prod,
        rand,
        randint,
        randn,
        randperm,
        randperm_diff,
        ranks,
        ratios_to_lengths,
        ratios_to_non_pad_mask,
        ratios_to_pad_mask,
        recursive_to,
        remove_at_indices,
        repeat_interleave_nd,
        resample_nearest_freqs,
        resample_nearest_rates,
        resample_nearest_steps,
        rmse,
        seed,
        segments_list_to_activity,
        segments_to_activity,
        segments_to_segments_list,
        set_default_dtype,
        set_default_generator,
        shape,
        shuffled,
        softmax_multidim,
        split,
        squeeze,
        squeeze_,
        squeeze_copy,
        stack,
        tensor_to_lengths,
        tensor_to_non_pad_mask,
        tensor_to_pad_mask,
        tensor_to_tensors_list,
        tensors_list_to_lengths,
        to_item,
        to_tensor,
        top_k,
        top_p,
        topk,
        transform_drop,
        unsqueeze,
        unsqueeze_,
        unsqueeze_copy,
        view_as_complex,
        view_as_real,
        where,
        zeros,
    )
    from .serialization.common import as_builtin
    from .serialization.csv import (
        dump_csv,
        dumps_csv,
        load_csv,
        loads_csv,
        read_csv,
        save_csv,
    )
    from .serialization.dump_fn import dump_to, save_to
    from .serialization.json import (
        dump_json,
        dumps_json,
        load_json,
        loads_json,
        read_json,
        save_json,
    )
    from .serialization.load_fn import load_from, read_from
    from .serialization.pickle import (
        dump_pickle,
        dumps_pickle,
        load_pickle,
        loads_pickle,
        read_pickle,
        save_pickle,
    )
    from .serialization.torch import dump_torch, load, load_torch, save
    from .types import Tensor
    from .types.guards import (
        is_builtin_collection,
        is_builtin_number,
        is_builtin_obj,
        is_builtin_scalar,
        is_dataclass_instance,
        is_namedtuple_instance,
        is_number_like,
        is_scalar_like,
        is_tensor_or_array,
        is_typed_dict,
        isinstance_generic,
    )
    from .types.tensor_subclasses import (
        BoolTensor,
        BoolTensor0D,
        BoolTensor1D,
        BoolTensor2D,
        BoolTensor3D,
        ByteTensor,
        ByteTensor0D,
        ByteTensor1D,
        ByteTensor2D,
        ByteTensor3D,
        CDoubleTensor,
        CDoubleTensor0D,
        CDoubleTensor1D,
        CDoubleTensor2D,
        CDoubleTensor3D,
        CFloatTensor,
        CFloatTensor0D,
        CFloatTensor1D,
        CFloatTensor2D,
        CFloatTensor3D,
        CharTensor,
        CharTensor0D,
        CharTensor1D,
        CharTensor2D,
        CharTensor3D,
        ComplexFloatingTensor,
        ComplexFloatingTensor0D,
        ComplexFloatingTensor1D,
        ComplexFloatingTensor2D,
        ComplexFloatingTensor3D,
        DoubleTensor,
        DoubleTensor0D,
        DoubleTensor1D,
        DoubleTensor2D,
        DoubleTensor3D,
        FloatingTensor,
        FloatingTensor0D,
        FloatingTensor1D,
        FloatingTensor2D,
        FloatingTensor3D,
        FloatTensor,
        FloatTensor0D,
        FloatTensor1D,
        FloatTensor2D,
        FloatTensor3D,
        HalfTensor,
        HalfTensor0D,
        HalfTensor1D,
        HalfTensor2D,
        HalfTensor3D,
        IntegralTensor,
        IntegralTensor0D,
        IntegralTensor1D,
        IntegralTensor2D,
        IntegralTensor3D,
        IntTensor,
        IntTensor0D,
        IntTensor1D,
        IntTensor2D,
        IntTensor3D,
        LongTensor,
        LongTensor0D,
        LongTensor1D,
        LongTensor2D,
        LongTensor3D,
        RealTensor,
        RealTensor0D,
        RealTensor1D,
        RealTensor2D,
        RealTensor3D,
        ShortTensor,
        ShortTensor0D,
        ShortTensor1D,
        ShortTensor2D,
        ShortTensor3D,
        SignedIntegerTensor,
        SignedIntegerTensor0D,
        SignedIntegerTensor1D,
        SignedIntegerTensor2D,
        SignedIntegerTensor3D,
        Tensor0D,
        Tensor1D,
        Tensor2D,
        Tensor3D,
        UnsignedIntegerTensor,
        UnsignedIntegerTensor0D,
        UnsignedIntegerTensor1D,
        UnsignedIntegerTensor2D,
        UnsignedIntegerTensor3D,
    )

else:
    import lazy_loader as lazy

    __getattr__, __dir__, _ = lazy.attach(
        __name__,
        submodules=[
            "core",
            "extras",
            "hub",
            "nn",
            "optim",
            "serialization",
            "types",
            "utils",
        ],
        submod_attrs={
            "core.dtype_enum": [
                "DTypeEnum",
            ],
            "core.dtype": [
                "complex64",
                "complex128",
                "double",
                "float",
                "float16",
                "float32",
                "float64",
                "half",
                "int",
                "int8",
                "int16",
                "int32",
                "int64",
                "long",
                "short",
                "uint8",
            ],
            "core.make": [
                "CUDA_IF_AVAILABLE",
                "Device",
                "DeviceLike",
                "DTypeLike",
                "Generator",
                "GeneratorLike",
            ],
            "core.version": ["get_githash_full", "get_githash_short"],
            "extras.numpy": ["ndarray_to_tensor", "tensor_to_ndarray", "to_ndarray"],
            "extras.yaml.yaml": [
                "dump_yaml",
                "dumps_yaml",
                "load_yaml",
                "loads_yaml",
                "read_yaml",
                "save_yaml",
            ],
            "hub.download": ["download_file"],
            "hub.paths": ["get_cache_dir", "get_tmp_dir"],
            "nn.functional": [
                "log_softmax_multidim",
                "softmax_multidim",
                "checksum_any",
                "crop_dim",
                "crop_dims",
                "get_inverse_perm",
                "get_perm_indices",
                "insert_at_indices",
                "randperm_diff",
                "remove_at_indices",
                "as_device",
                "as_dtype",
                "as_generator",
                "get_default_device",
                "get_default_dtype",
                "get_default_generator",
                "set_default_dtype",
                "set_default_generator",
                "generate_square_subsequent_mask",
                "lengths_to_non_pad_mask",
                "lengths_to_pad_mask",
                "lengths_to_ratios",
                "masked_equal",
                "masked_mean",
                "masked_sum",
                "non_pad_mask_to_lengths",
                "non_pad_mask_to_ratios",
                "pad_mask_to_lengths",
                "pad_mask_to_ratios",
                "ratios_to_lengths",
                "ratios_to_non_pad_mask",
                "ratios_to_pad_mask",
                "tensor_to_lengths",
                "tensor_to_non_pad_mask",
                "tensor_to_pad_mask",
                "tensor_to_tensors_list",
                "tensors_list_to_lengths",
                "index_to_name",
                "index_to_onehot",
                "name_to_index",
                "name_to_onehot",
                "one_hot",
                "onehot_to_index",
                "onehot_to_name",
                "probs_to_index",
                "probs_to_name",
                "probs_to_onehot",
                "indices_to_multihot",
                "indices_to_multinames",
                "multi_indices_to_multihot",
                "multi_indices_to_multinames",
                "multihot_to_indices",
                "multihot_to_multi_indices",
                "multihot_to_multinames",
                "multinames_to_indices",
                "multinames_to_multi_indices",
                "multinames_to_multihot",
                "probs_to_indices",
                "probs_to_multi_indices",
                "probs_to_multihot",
                "probs_to_multinames",
                "arange",
                "empty",
                "full",
                "ones",
                "rand",
                "randn",
                "randint",
                "randperm",
                "zeros",
                "average_power",
                "cat",
                "concat",
                "count_parameters",
                "deep_equal",
                "find",
                "get_ndim",
                "get_shape",
                "mse",
                "ndim",
                "nelement",
                "prod",
                "ranks",
                "rmse",
                "shape",
                "stack",
                "cat_padded_batch",
                "pad_and_stack_rec",
                "pad_dim",
                "pad_dims",
                "multilabel_to_powerset",
                "powerset_to_multilabel",
                "all_eq",
                "all_ne",
                "is_complex",
                "is_convertible_to_tensor",
                "is_floating_point",
                "is_full",
                "is_sorted",
                "is_stackable",
                "is_tensor",
                "is_unique",
                "activity_to_segments",
                "activity_to_segments_list",
                "segments_list_to_activity",
                "segments_to_activity",
                "segments_to_segments_list",
                "as_tensor",
                "flatten",
                "identity",
                "move_to",
                "move_to_rec",
                "pad_and_crop_dim",
                "recursive_to",
                "repeat_interleave_nd",
                "resample_nearest_freqs",
                "resample_nearest_rates",
                "resample_nearest_steps",
                "shuffled",
                "squeeze",
                "squeeze_",
                "squeeze_copy",
                "to_item",
                "to_tensor",
                "top_k",
                "top_p",
                "topk",
                "transform_drop",
                "unsqueeze",
                "unsqueeze_",
                "unsqueeze_copy",
                "view_as_complex",
                "view_as_real",
                "equal",
                "initial_seed",
                "manual_seed",
                "matmul",
                "no_grad",
                "seed",
                "split",
                "where",
            ],
            "serialization.common": [
                "as_builtin",
            ],
            "serialization.csv": [
                "dump_csv",
                "dumps_csv",
                "load_csv",
                "loads_csv",
                "read_csv",
                "save_csv",
            ],
            "serialization.dump_fn": [
                "save_to",
                "dump_to",
            ],
            "serialization.json": [
                "dump_json",
                "dumps_json",
                "load_json",
                "loads_json",
                "read_json",
                "save_json",
            ],
            "serialization.load_fn": [
                "load_from",
                "read_from",
            ],
            "serialization.pickle": [
                "dump_pickle",
                "dumps_pickle",
                "load_pickle",
                "loads_pickle",
                "read_pickle",
                "save_pickle",
            ],
            "serialization.torch": ["dump_torch", "load_torch", "save", "load"],
            "types": [
                "Tensor",
            ],
            "types.guards": [
                "is_number_like",
                "is_scalar_like",
                "is_tensor_or_array",
                "is_builtin_collection",
                "is_builtin_number",
                "is_builtin_obj",
                "is_builtin_scalar",
                "is_dataclass_instance",
                "is_namedtuple_instance",
                "is_typed_dict",
                "isinstance_generic",
            ],
            "types.tensor_subclasses": [
                "BoolTensor",
                "BoolTensor0D",
                "BoolTensor1D",
                "BoolTensor2D",
                "BoolTensor3D",
                "ByteTensor",
                "ByteTensor0D",
                "ByteTensor1D",
                "ByteTensor2D",
                "ByteTensor3D",
                "CDoubleTensor",
                "CDoubleTensor0D",
                "CDoubleTensor1D",
                "CDoubleTensor2D",
                "CDoubleTensor3D",
                "CFloatTensor",
                "CFloatTensor0D",
                "CFloatTensor1D",
                "CFloatTensor2D",
                "CFloatTensor3D",
                "CharTensor",
                "CharTensor0D",
                "CharTensor1D",
                "CharTensor2D",
                "CharTensor3D",
                "ComplexFloatingTensor",
                "ComplexFloatingTensor0D",
                "ComplexFloatingTensor1D",
                "ComplexFloatingTensor2D",
                "ComplexFloatingTensor3D",
                "DoubleTensor",
                "DoubleTensor0D",
                "DoubleTensor1D",
                "DoubleTensor2D",
                "DoubleTensor3D",
                "FloatingTensor",
                "FloatingTensor0D",
                "FloatingTensor1D",
                "FloatingTensor2D",
                "FloatingTensor3D",
                "FloatTensor",
                "FloatTensor0D",
                "FloatTensor1D",
                "FloatTensor2D",
                "FloatTensor3D",
                "HalfTensor",
                "HalfTensor0D",
                "HalfTensor1D",
                "HalfTensor2D",
                "HalfTensor3D",
                "IntegralTensor",
                "IntegralTensor0D",
                "IntegralTensor1D",
                "IntegralTensor2D",
                "IntegralTensor3D",
                "IntTensor",
                "IntTensor0D",
                "IntTensor1D",
                "IntTensor2D",
                "IntTensor3D",
                "LongTensor",
                "LongTensor0D",
                "LongTensor1D",
                "LongTensor2D",
                "LongTensor3D",
                "RealTensor",
                "RealTensor0D",
                "RealTensor1D",
                "RealTensor2D",
                "RealTensor3D",
                "ShortTensor",
                "ShortTensor0D",
                "ShortTensor1D",
                "ShortTensor2D",
                "ShortTensor3D",
                "SignedIntegerTensor",
                "SignedIntegerTensor0D",
                "SignedIntegerTensor1D",
                "SignedIntegerTensor2D",
                "SignedIntegerTensor3D",
                "Tensor0D",
                "Tensor1D",
                "Tensor2D",
                "Tensor3D",
                "UnsignedIntegerTensor",
                "UnsignedIntegerTensor0D",
                "UnsignedIntegerTensor1D",
                "UnsignedIntegerTensor2D",
                "UnsignedIntegerTensor3D",
            ],
        },
    )

from pythonwrench.semver import Version

version = __version__
version_info = Version(__version__)

del TYPE_CHECKING, Version
