# Copyright (c) 2021 Graphcore Ltd. All rights reserved.

__version__ = "@VERSION@"

__all__ = [
    # context.py
    "get_current_graph", "get_main_graph", "get_ir", "gcg", "gmg", "gir",
    "in_sequence", "name_scope", "ipu", "io_tiles",
    # dtypes.py
    "dtype", "bool", "int8", "int16", "int32", "int64", "uint8", "uint16",
    "uint32", "uint64", "half", "float8_143", "float8_152", "float16", "float",
    "float32", "double", "float64", "complex64", "complex128",
    # graph.py
    "Graph",
    # ir.py
    "Ir",
    # module.py
    "Module",
    # random_seed.py,
    "create_seeds",
    # remote_buffer.py
    "RemoteBuffer", "remote_buffer",
    # replica_grouping.py
    "ReplicaGrouping",
    # streams.py
    "HostToDeviceStream", "DeviceToHostStream", "h2d_stream", "d2h_stream",
    # tensor.py
    "Tensor", "variable", "remote_variable", "remote_replica_sharded_variable",
    "replica_sharded_variable", "constant", "graph_input", "graph_output",
    "TensorByRef", "TensorSpec",
    # tensorlocation.py
    "TensorStorage", "TensorLocation", "ReplicatedTensorSharding", "TileSet",
    "ExecutionContext",
    # session.py
    "Session"
]

_import_exc = None
try:
    from .context import (get_current_graph, get_main_graph, get_ir, gcg, gmg, gir,
        in_sequence, name_scope, ipu, io_tiles)
    from .dtypes import (dtype, bool, int8, int16, int32, int64, uint8,
        uint16, uint32, uint64, half, float8_143, float8_152, float16,
        float, float32, double, float64, complex64, complex128)
    from .graph import Graph
    from .module import Module
    from .ir import Ir
    from .random_seed import create_seeds
    from .remote_buffer import (RemoteBuffer, remote_buffer)
    from .replica_grouping import (ReplicaGrouping)
    from .streams import (HostToDeviceStream, DeviceToHostStream,
        h2d_stream, d2h_stream)
    from .tensor import (Tensor, variable, remote_variable,
        remote_replica_sharded_variable, replica_sharded_variable,
        replica_sharded_buffer, constant,
        graph_input, graph_output, TensorByRef, TensorSpec)
    from .tensorlocation import (TensorStorage, TensorLocation,
        ReplicatedTensorSharding, TileSet, ExecutionContext)
    from .session import Session

    from . import transforms
except ModuleNotFoundError as exc:
    _import_exc = exc

if _import_exc:
    msg = (
        f"{_import_exc}.\n"
        "Hint: You might be running a Python version that PopXL doesn't support."
    )
    raise RuntimeError(msg)

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    # Static Type Checking requires "import .. as"
    # however this causes a circular import at runtime
    import popxl.transforms as transforms