from __future__ import annotations

from typing import Any, Optional

from .callbacks import EarlyStopping as EarlyStopping
from .callbacks import config_callbacks as config_callbacks
from .model_summary import summary as summary

def to_list(value: Any): ...
def to_numpy(var: Any): ...
def flatten_list(l: Any): ...
def restore_flatten_list(l: Any, splits: Any): ...
def extract_args(func: Any): ...
def wait_server_ready(endpoints: Any) -> None: ...
def init_communicator(
    program: Any, rank: Any, nranks: Any, wait_port: Any, current_endpoint: Any, endpoints: Any
) -> None: ...
def prepare_distributed_context(place: Optional[Any] = ...): ...

class StaticGraphAdapter:
    model: Any = ...
    def __init__(self, model: Any) -> None: ...
    @property
    def mode(self): ...
    @mode.setter
    def mode(self, value: Any) -> None: ...
    def train_batch(self, inputs: Any, labels: Optional[Any] = ..., update: bool = ...): ...
    def eval_batch(self, inputs: Any, labels: Optional[Any] = ...): ...
    def predict_batch(self, inputs: Any): ...
    def parameters(self, *args: Any, **kwargs: Any): ...
    def save(self, path: Any) -> None: ...
    def load(self, param_state_pairs: Any, optim_state: Any) -> None: ...
    def prepare(self) -> None: ...

class DynamicGraphAdapter:
    model: Any = ...
    ddp_model: Any = ...
    def __init__(self, model: Any) -> None: ...
    @property
    def mode(self): ...
    @mode.setter
    def mode(self, value: Any) -> None: ...
    def train_batch(self, inputs: Any, labels: Optional[Any] = ..., update: bool = ...): ...
    def eval_batch(self, inputs: Any, labels: Optional[Any] = ...): ...
    def predict_batch(self, inputs: Any): ...
    def parameters(self, *args: Any, **kwargs: Any): ...
    def save(self, path: Any) -> None: ...
    def load(self, param_state_pairs: Any, optim_state: Any, scaler_state: Optional[Any] = ...): ...
    def prepare(self) -> None: ...

class Model:
    mode: str = ...
    network: Any = ...
    stop_training: bool = ...
    def __init__(self, network: Any, inputs: Optional[Any] = ..., labels: Optional[Any] = ...) -> None: ...
    def train_batch(self, inputs: Any, labels: Optional[Any] = ..., update: bool = ...): ...
    def eval_batch(self, inputs: Any, labels: Optional[Any] = ...): ...
    def predict_batch(self, inputs: Any): ...
    def save(self, path: Any, training: bool = ...) -> None: ...
    def load(self, path: Any, skip_mismatch: bool = ..., reset_optimizer: bool = ...): ...
    def parameters(self, *args: Any, **kwargs: Any): ...
    def prepare(
        self,
        optimizer: Optional[Any] = ...,
        loss: Optional[Any] = ...,
        metrics: Optional[Any] = ...,
        amp_configs: Optional[Any] = ...,
    ) -> None: ...
    num_iters: Any = ...
    def fit(
        self,
        train_data: Optional[Any] = ...,
        eval_data: Optional[Any] = ...,
        batch_size: int = ...,
        epochs: int = ...,
        eval_freq: int = ...,
        log_freq: int = ...,
        save_dir: Optional[Any] = ...,
        save_freq: int = ...,
        verbose: int = ...,
        drop_last: bool = ...,
        shuffle: bool = ...,
        num_workers: int = ...,
        callbacks: Optional[Any] = ...,
        accumulate_grad_batches: int = ...,
        num_iters: Optional[Any] = ...,
    ) -> None: ...
    def evaluate(
        self,
        eval_data: Any,
        batch_size: int = ...,
        log_freq: int = ...,
        verbose: int = ...,
        num_workers: int = ...,
        callbacks: Optional[Any] = ...,
        num_iters: Optional[Any] = ...,
    ): ...
    def predict(
        self,
        test_data: Any,
        batch_size: int = ...,
        num_workers: int = ...,
        stack_outputs: bool = ...,
        verbose: int = ...,
        callbacks: Optional[Any] = ...,
    ): ...
    def summary(self, input_size: Optional[Any] = ..., dtype: Optional[Any] = ...): ...
