from typing import Any
from typing import Dict
from typing import List
from typing import Union
from typing import Optional

from .monitors import BasicMonitor
from ...types import configs_type
from ...trainer import Trainer
from ...protocol import TqdmSettings
from ...protocol import OptimizerPack
from ...protocol import TrainerMonitor
from ...protocol import _IMetric
from ...protocol import _MultipleMetrics
from ...protocol import TrainerCallback
from .callbacks.general import _LogMetricsMsgCallback


def make_trainer(
    state_config: Optional[Dict[str, Any]] = None,
    *,
    workplace: str,
    num_epoch: int = 40,
    max_epoch: int = 1000,
    fixed_epoch: Optional[int] = None,
    fixed_steps: Optional[int] = None,
    log_steps: Optional[int] = None,
    valid_portion: float = 1.0,
    amp: bool = False,
    clip_norm: float = 0.0,
    metric_names: Optional[Union[str, List[str]]] = None,
    metric_configs: configs_type = None,
    metric_weights: Optional[Dict[str, float]] = None,
    use_losses_as_metrics: Optional[bool] = None,
    loss_metrics_weights: Optional[Dict[str, float]] = None,
    recompute_train_losses_in_eval: bool = True,
    monitor_names: Optional[Union[str, List[str]]] = None,
    monitor_configs: configs_type = None,
    callback_names: Optional[Union[str, List[str]]] = None,
    callback_configs: configs_type = None,
    lr: Optional[float] = None,
    optimizer_name: Optional[str] = None,
    scheduler_name: Optional[str] = None,
    optimizer_config: Optional[Dict[str, Any]] = None,
    scheduler_config: Optional[Dict[str, Any]] = None,
    update_scheduler_per_epoch: bool = False,
    optimizer_settings: Optional[Dict[str, Dict[str, Any]]] = None,
    use_zero: bool = False,
    data_info_name: str = "data_info",
    metrics_log_file: str = "metrics.txt",
    finetune_config: Optional[Dict[str, Any]] = None,
    tqdm_settings: Optional[Dict[str, Any]] = None,
) -> Trainer:
    if state_config is None:
        state_config = {}
    if log_steps is not None:
        state_config.setdefault("num_step_per_log", log_steps)
        state_config.setdefault("snapshot_start_step", log_steps)
        state_config.setdefault("num_step_per_snapshot", log_steps)
    # metrics
    metrics: Optional[Union[_IMetric, _MultipleMetrics]]
    if metric_names is None:
        metrics = None
    else:
        metrics = _IMetric.fuse(
            metric_names,
            metric_configs,
            metric_weights=metric_weights,
        )
    # monitor
    monitors: Optional[List[TrainerMonitor]]
    if monitor_names is None:
        monitors = None
    else:
        monitors = TrainerMonitor.make_multiple(monitor_names, monitor_configs)
    # callback
    callbacks: Optional[List[TrainerCallback]]
    if callback_names is None:
        callbacks = None
    else:
        callbacks = TrainerCallback.make_multiple(callback_names, callback_configs)
    # optimizer
    if optimizer_settings is None:
        optimizer_packs = None
    else:
        optimizer_packs = []
        for key, settings in optimizer_settings.items():
            optimizer = settings.get("optimizer")
            if optimizer is None:
                raise ValueError(f"optimizer must be provided (key={key})")
            optimizer_packs.append(
                OptimizerPack(
                    key,
                    optimizer,
                    settings.get("scheduler"),
                    settings.get("optimizer_config"),
                    settings.get("scheduler_config"),
                )
            )
    # tqdm
    if tqdm_settings is None:
        tqdm_settings = {}
    use_tqdm = tqdm_settings.setdefault("use_tqdm", False)
    # epoch
    if fixed_epoch is not None:
        num_epoch = max_epoch = fixed_epoch
    if max_epoch < num_epoch:
        raise ValueError("`max_epoch` should not be smaller than `num_epoch`")
    # default behaviors
    if monitors is None:
        monitors = [BasicMonitor()]
    if callbacks is None:
        callbacks = [_LogMetricsMsgCallback(not use_tqdm)]
    return Trainer(
        state_config,
        num_epoch=num_epoch,
        max_epoch=max_epoch,
        fixed_steps=fixed_steps,
        valid_portion=valid_portion,
        amp=amp,
        clip_norm=clip_norm,
        metrics=metrics,
        use_losses_as_metrics=use_losses_as_metrics,
        loss_metrics_weights=loss_metrics_weights,
        recompute_train_losses_in_eval=recompute_train_losses_in_eval,
        monitors=monitors,
        callbacks=callbacks,
        lr=lr,
        optimizer_name=optimizer_name,
        scheduler_name=scheduler_name,
        optimizer_config=optimizer_config,
        scheduler_config=scheduler_config,
        update_scheduler_per_epoch=update_scheduler_per_epoch,
        optimizer_packs=optimizer_packs,
        use_zero=use_zero,
        workplace=workplace,
        data_info_name=data_info_name,
        metrics_log_file=metrics_log_file,
        finetune_config=finetune_config,
        tqdm_settings=TqdmSettings(
            use_tqdm,
            tqdm_settings.setdefault("use_step_tqdm", use_tqdm),
            tqdm_settings.setdefault("use_tqdm_in_validation", False),
            tqdm_settings.setdefault("in_distributed", False),
            tqdm_settings.setdefault("tqdm_position", 0),
            tqdm_settings.setdefault("tqdm_desc", "epoch"),
        ),
    )


__all__ = [
    "make_trainer",
]
