"""Utils for HPO with hpopt."""

# Copyright (C) 2021-2022 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
#

import json
import logging
import re
import shutil
from copy import deepcopy
from enum import Enum
from functools import partial
from inspect import isclass
from math import floor
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union

import torch
import yaml

from otx.api.configuration.helper import create
from otx.api.entities.datasets import DatasetEntity
from otx.api.entities.model import ModelEntity
from otx.api.entities.model_template import TaskType
from otx.api.entities.subset import Subset
from otx.api.entities.task_environment import TaskEnvironment
from otx.api.entities.train_parameters import TrainParameters, UpdateProgressCallback
from otx.cli.utils.importing import get_impl_class
from otx.cli.utils.io import read_model, save_model_data
from otx.core.data.adapter import get_dataset_adapter
from otx.hpo import HyperBand, TrialStatus, run_hpo_loop

logger = logging.getLogger(__name__)


def _check_hpo_enabled_task(task_type):
    return task_type in [
        TaskType.CLASSIFICATION,
        TaskType.DETECTION,
        TaskType.SEGMENTATION,
        TaskType.INSTANCE_SEGMENTATION,
        TaskType.ROTATED_DETECTION,
        TaskType.ANOMALY_CLASSIFICATION,
        TaskType.ANOMALY_DETECTION,
        TaskType.ANOMALY_SEGMENTATION,
    ]


class TaskManager:
    """Task utility class to give common interface from different task.

    Args:
        task_type (TaskType): otx task type
    """

    def __init__(self, task_type: TaskType):
        self._task_type = task_type

    @property
    def task_type(self):
        """Task_type property."""
        return self._task_type

    def is_mmcv_framework_task(self) -> bool:
        """Check task is run on mmcv.

        Returns:
            bool: whether task is run on mmcv
        """
        return self.is_cls_framework_task() or self.is_det_framework_task() or self.is_seg_framework_task()

    def is_cls_framework_task(self) -> bool:
        """Check that task is run on mmcls framework.

        Returns:
            bool: whether task is run on mmcls
        """
        return self._task_type == TaskType.CLASSIFICATION

    def is_det_framework_task(self) -> bool:
        """Check that task is one of a task run on mmdet framework.

        Returns:
            bool: whether task is run on mmdet
        """
        return self._task_type in [
            TaskType.DETECTION,
            TaskType.INSTANCE_SEGMENTATION,
            TaskType.ROTATED_DETECTION,
        ]

    def is_seg_framework_task(self) -> bool:
        """Check that task is run on mmseg framework.

        Returns:
            bool: whether tasks is run on mmseg
        """
        return self._task_type == TaskType.SEGMENTATION

    def is_anomaly_framework_task(self) -> bool:
        """Check taht task is run on anomalib.

        Returns:
            bool: whether task is run on anomalib
        """
        return self._task_type in [
            TaskType.ANOMALY_CLASSIFICATION,
            TaskType.ANOMALY_DETECTION,
            TaskType.ANOMALY_SEGMENTATION,
        ]

    def get_batch_size_name(self) -> str:
        """Give an proper batch size name depending on framework.

        Returns:
            str: batch size name
        """
        if self.is_mmcv_framework_task():
            batch_size_name = "learning_parameters.batch_size"
        elif self.is_anomaly_framework_task():
            batch_size_name = "learning_parameters.train_batch_size"
        else:
            raise RuntimeError(f"There is no information about {self._task_type} batch size name")

        return batch_size_name

    def get_epoch_name(self) -> str:
        """Give an proper epoch name depending on framework.

        Returns:
            str: epoch name
        """
        if self.is_mmcv_framework_task():
            epoch_name = "num_iters"
        elif self.is_anomaly_framework_task():
            epoch_name = "max_epochs"
        else:
            raise RuntimeError(f"There is no information about {self._task_type} epoch name")

        return epoch_name

    def copy_weight(self, src: Union[str, Path], det: Union[str, Path]):
        """Copy all model weights from work directory.

        Args:
            src (Union[str, Path]): path where model weights are saved
            det (Union[str, Path]): path to save model weights
        """
        src = Path(src)
        det = Path(det)
        if self.is_mmcv_framework_task():
            for weight_candidate in src.rglob("*epoch*.pth"):
                if not (weight_candidate.is_symlink() or (det / weight_candidate.name).exists()):
                    shutil.copy(weight_candidate, det)
        # TODO need to implement after anomaly task supports resume

    def get_latest_weight(self, workdir: Union[str, Path]) -> Optional[str]:
        """Get latest model weight from all weights.

        Args:
            workdir (Union[str, Path]): path where model weights are saved

        Returns:
            Optional[str]: latest model weight path. If not found, than return None value.
        """
        latest_weight = None
        workdir = Path(workdir)
        if self.is_mmcv_framework_task():
            pattern = re.compile(r"(\d+)\.pth")
            current_latest_epoch = -1
            latest_weight = None

            for weight_name in workdir.rglob("epoch_*.pth"):
                ret = pattern.search(str(weight_name))
                if ret is not None:
                    epoch = int(ret.group(1))
                    if current_latest_epoch < epoch:
                        current_latest_epoch = epoch
                        latest_weight = str(weight_name)
        # TODO need to implement after anomaly task supports resume

        return latest_weight


class TaskEnvironmentManager:
    """OTX environment utility class to set or get a value from environment class.

    Args:
        environment (TaskEnvironment): OTX task environment
    """

    def __init__(self, environment: TaskEnvironment):
        self._environment = environment
        self.task = TaskManager(environment.model_template.task_type)

    @property
    def environment(self):
        """Environment property."""
        return self._environment

    def get_task(self) -> TaskType:
        """Get task type of environment.

        Returns:
            TaskType: task type
        """
        return self._environment.model_template.task_type

    def get_model_template(self):
        """Get model template."""
        return self._environment.model_template

    def get_model_template_path(self) -> str:
        """Get model template path.

        Returns:
            str: path of model template
        """
        return self._environment.model_template.model_template_path

    def set_hyper_parameter_using_str_key(self, hyper_parameter: Dict[str, Any]):
        """Set hyper parameter to environment using string key hyper_parameter.

        Set hyper parameter to environment. Argument `hyper_parameter` is a dictionary which has string key.
        For example, hyper_parameter has a key "a.b.c", then value is set at env_hp.a.b.c.

        Args:
            hyper_parameter (Dict[str, Any]): hyper parameter to set which has a string format
        """
        env_hp = self._environment.get_hyper_parameters()  # type: ignore

        for param_key, param_val in hyper_parameter.items():
            splited_param_key = param_key.split(".")

            target = env_hp
            for val in splited_param_key[:-1]:
                target = getattr(target, val)
            setattr(target, splited_param_key[-1], param_val)

    def get_dict_type_hyper_parameter(self) -> Dict[str, Any]:
        """Get dictionary type hyper parmaeter of environment.

        Returns:
            Dict[str, Any]: dictionary type hyper parameter of environment
        """
        learning_parameters = self._environment.get_hyper_parameters().learning_parameters  # type: ignore
        learning_parameters = self._convert_parameter_group_to_dict(learning_parameters)
        hyper_parameter = {f"learning_parameters.{key}": val for key, val in learning_parameters.items()}
        return hyper_parameter

    def _convert_parameter_group_to_dict(self, parameter_group) -> Dict[str, Any]:
        """Convert parameter group to dictionary.

        Args:
            parameter_group : parameter gruop

        Returns:
            Dict[str, Any]: parameter group converted to dictionary
        """
        groups = getattr(parameter_group, "groups", None)
        parameters = getattr(parameter_group, "parameters", None)

        total_arr = []
        for val in [groups, parameters]:
            if val is not None:
                total_arr.extend(val)
        if not total_arr:
            return parameter_group

        ret = {}
        for key in total_arr:
            val = self._convert_parameter_group_to_dict(getattr(parameter_group, key))
            if not (isclass(val) or isinstance(val, Enum)):
                ret[key] = val

        return ret

    def get_max_epoch(self) -> int:
        """Get max epoch from environment.

        Returns:
            int: max epoch of environment
        """
        return getattr(
            self._environment.get_hyper_parameters().learning_parameters, self.task.get_epoch_name()  # type: ignore
        )

    def save_initial_weight(self, save_path: Union[Path, str]) -> bool:
        """Save an initial model weight.

        Args:
            save_path (Union[str, Path]): path to save initial model weight

        Returns:
            bool: whether model weight is saved successfully
        """
        save_path = Path(save_path)
        dir_path = save_path.parent
        if self._environment.model is None:
            # if task isn't anomaly, then save model weight during first trial
            if self.task.is_anomaly_framework_task():
                task = self.get_train_task()
                model = self.get_new_model_entity()
                task.save_model(model)
                save_model_data(model, str(dir_path))
                (dir_path / "weights.pth").rename(save_path)
                return True
        else:
            save_model_data(self._environment.model, str(dir_path))
            (dir_path / "weights.pth").rename(save_path)
            return True
        return False

    def get_train_task(self):
        """Get OTX train task instance.

        Returns:
           OTX task: OTX train task instance
        """
        impl_class = get_impl_class(self._environment.model_template.entrypoints.base)
        return impl_class(task_environment=self._environment)

    def get_batch_size_name(self) -> str:
        """Get proper batch size name depending on task.

        Returns:
            str: batch size name
        """
        return self.task.get_batch_size_name()

    def load_model_weight(self, model_weight_path: str, dataset: DatasetEntity):
        """Set model weight on environment to load the weight during training.

        Args:
            model_weight_path (str): model weight to load during training
            dataset (DatasetEntity): dataset for training a model
        """
        self._environment.model = read_model(self._environment.get_model_configuration(), model_weight_path, dataset)

    def resume_model_weight(self, model_weight_path: str, dataset: DatasetEntity):
        """Set model weight on environment to resume the weight during training.

        Args:
            model_weight_path (str): model weight to resume during training
            dataset (DatasetEntity): dataset for training a model
        """
        self.load_model_weight(model_weight_path, dataset)
        self._environment.model.model_adapters["resume"] = True  # type: ignore

    def get_new_model_entity(self, dataset=None) -> ModelEntity:
        """Get new model entity using environment.

        Args:
            dataset (Optional[DatasetEntity]): OTX dataset

        Returns:
            ModelEntity: new model entity
        """
        return ModelEntity(
            dataset,
            self._environment.get_model_configuration(),
        )

    def set_epoch(self, epoch: int):
        """Set epoch on environment.

        Args:
            epoch (int): epoch to set
        """
        hyper_parameter = {f"learning_parameters.{self.task.get_epoch_name()}": epoch}
        self.set_hyper_parameter_using_str_key(hyper_parameter)


class HpoRunner:
    """Class which is in charge of preparing and running HPO.

    Args:
        environment (TaskEnvironment): OTX environment
        train_dataset_size (int): train dataset size
        val_dataset_size (int): validation dataset size
        hpo_workdir (Union[str, Path]): work directory for HPO
        hpo_time_ratio (int, optional): time ratio to use for HPO compared to training time. Defaults to 4.
    """

    # pylint: disable=too-many-instance-attributes

    def __init__(
        self,
        environment: TaskEnvironment,
        train_dataset_size: int,
        val_dataset_size: int,
        hpo_workdir: Union[str, Path],
        hpo_time_ratio: int = 4,
    ):
        if train_dataset_size <= 0:
            raise ValueError(f"train_dataset_size should be bigger than 0. Your value is {train_dataset_size}")
        if val_dataset_size <= 0:
            raise ValueError(f"val_dataset_size should be bigger than 0. Your value is {val_dataset_size}")
        if hpo_time_ratio < 1:
            raise ValueError(f"hpo_time_ratio shouldn't be smaller than 1. Your value is {hpo_time_ratio}")

        self._environment = TaskEnvironmentManager(environment)
        self._hpo_workdir: Path = Path(hpo_workdir)
        self._hpo_time_ratio = hpo_time_ratio
        self._hpo_config: Dict = self._set_hpo_config()
        self._train_dataset_size = train_dataset_size
        self._val_dataset_size = val_dataset_size
        self._fixed_hp: Dict[str, Any] = {}
        self._initial_weight_name = "initial_weight.pth"

        self._align_batch_size_search_space_to_dataset_size()

    def _set_hpo_config(self):
        hpo_config_path = Path(self._environment.get_model_template_path()).parent / "hpo_config.yaml"
        with hpo_config_path.open("r") as f:
            hpopt_cfg = yaml.safe_load(f)

        return hpopt_cfg

    def _align_batch_size_search_space_to_dataset_size(self):
        batch_size_name = self._environment.get_batch_size_name()

        if batch_size_name in self._hpo_config["hp_space"]:
            if "range" in self._hpo_config["hp_space"][batch_size_name]:
                max_val = self._hpo_config["hp_space"][batch_size_name]["range"][1]
                min_val = self._hpo_config["hp_space"][batch_size_name]["range"][0]
                if max_val > self._train_dataset_size:
                    max_val = self._train_dataset_size
                    self._hpo_config["hp_space"][batch_size_name]["range"][1] = max_val
            else:
                max_val = self._hpo_config["hp_space"][batch_size_name]["max"]
                min_val = self._hpo_config["hp_space"][batch_size_name]["min"]

                if max_val > self._train_dataset_size:
                    max_val = self._train_dataset_size
                    self._hpo_config["hp_space"][batch_size_name]["max"] = max_val

            # If trainset size is lower than min batch size range,
            # fix batch size to trainset size
            if min_val >= max_val:
                logger.info(
                    "Train set size is equal or lower than batch size range. Batch size is fixed to train set size."
                )
                del self._hpo_config["hp_space"][batch_size_name]
                self._fixed_hp[batch_size_name] = self._train_dataset_size
                self._environment.set_hyper_parameter_using_str_key(self._fixed_hp)

    def run_hpo(self, train_func: Callable, data_roots: Dict[str, str]) -> Dict[str, Any]:
        """Run HPO and provides optimized hyper parameters.

        Args:
            train_func (Callable): training model function
            data_roots (Dict[str, str]): dataset path of each dataset type

        Returns:
            Dict[str, Any]: optimized hyper parameters
        """
        self._environment.save_initial_weight(self._get_initial_model_weight_path())
        hpo_algo = self._get_hpo_algo()
        resource_type = "gpu" if torch.cuda.is_available() else "cpu"
        run_hpo_loop(
            hpo_algo,
            partial(
                train_func,
                model_template=self._environment.get_model_template(),
                data_roots=data_roots,
                task_type=self._environment.get_task(),
                hpo_workdir=self._hpo_workdir,
                initial_weight_name=self._initial_weight_name,
                metric=self._hpo_config["metric"],
            ),
            resource_type,  # type: ignore
        )
        best_config = hpo_algo.get_best_config()
        self._restore_fixed_hp(best_config["config"])
        hpo_algo.print_result()

        return best_config

    def _restore_fixed_hp(self, hyper_parameter: Dict[str, Any]):
        for key, val in self._fixed_hp.items():
            hyper_parameter[key] = val

    def _get_hpo_algo(self):
        hpo_algo_type = self._hpo_config.get("search_algorithm", "asha")

        if hpo_algo_type == "asha":
            hpo_algo = self._prepare_asha()
        elif hpo_algo_type == "smbo":
            hpo_algo = self._prepare_smbo()
        else:
            raise ValueError(f"Supported HPO algorithms are asha and smbo. your value is {hpo_algo_type}.")

        return hpo_algo

    def _prepare_asha(self):
        args = {
            "search_space": self._hpo_config["hp_space"],
            "save_path": str(self._hpo_workdir),
            "maximum_resource": self._hpo_config.get("maximum_resource"),
            "minimum_resource": self._hpo_config.get("minimum_resource"),
            "mode": self._hpo_config.get("mode", "max"),
            "num_workers": 1,
            "num_full_iterations": self._environment.get_max_epoch(),
            "full_dataset_size": self._train_dataset_size,
            "non_pure_train_ratio": self._val_dataset_size / (self._train_dataset_size + self._val_dataset_size),
            "metric": self._hpo_config.get("metric", "mAP"),
            "expected_time_ratio": self._hpo_time_ratio,
            "prior_hyper_parameters": self._get_default_hyper_parameters(),
            "asynchronous_bracket": True,
            "asynchronous_sha": torch.cuda.device_count() != 1,
        }

        logger.debug(f"ASHA args = {args}")

        return HyperBand(**args)

    def _prepare_smbo(self):
        raise NotImplementedError

    def _get_default_hyper_parameters(self):
        default_hyper_parameters = {}
        hp_from_env = self._environment.get_dict_type_hyper_parameter()

        for key, val in hp_from_env.items():
            if key in self._hpo_config["hp_space"]:
                default_hyper_parameters[key] = val

        if not default_hyper_parameters:
            return None
        return default_hyper_parameters

    def _get_initial_model_weight_path(self):
        return self._hpo_workdir / self._initial_weight_name


def run_hpo(
    args, environment: TaskEnvironment, dataset: DatasetEntity, data_roots: Dict[str, str]
) -> Optional[TaskEnvironment]:
    """Run HPO and load optimized hyper parameter and best HPO model weight.

    Args:
        args: arguments passed to otx train
        environment (TaskEnvironment): otx task environment
        dataset (DatasetEntity): dataset to use for training
        data_roots (Dict[str, str]): dataset path of each dataset type
    """
    task_type = environment.model_template.task_type
    if not _check_hpo_enabled_task(task_type):
        logger.warning(
            "Currently supported task types are classification, detection, segmentation and anomaly"
            f"{task_type} is not supported yet."
        )
        return None

    hpo_save_path = (Path(args.save_model_to).parent / "hpo").absolute()
    hpo_runner = HpoRunner(
        environment,
        len(dataset.get_subset(Subset.TRAINING)),
        len(dataset.get_subset(Subset.VALIDATION)),
        hpo_save_path,
        args.hpo_time_ratio,
    )

    logger.info("started hyper-parameter optimization")
    best_config = hpo_runner.run_hpo(run_trial, data_roots)
    logger.info("completed hyper-parameter optimization")

    env_manager = TaskEnvironmentManager(environment)
    env_manager.set_hyper_parameter_using_str_key(best_config["config"])
    best_hpo_weight = get_best_hpo_weight(hpo_save_path, best_config["id"])
    if best_hpo_weight is None:
        logger.warning("Can not find the best HPO weight. Best HPO wegiht won't be used.")
    else:
        logger.debug(f"{best_hpo_weight} will be loaded as best HPO weight")
        env_manager.load_model_weight(best_hpo_weight, dataset)

    return env_manager.environment


def get_best_hpo_weight(hpo_dir: Union[str, Path], trial_id: Union[str, Path]) -> Optional[str]:
    """Get best model weight path of the HPO trial.

    Args:
        hpo_dir (Union[str, Path]): HPO work directory path
        trial_id (Union[str, Path]): trial id

    Returns:
        Optional[str]: best HPO model weight
    """
    hpo_dir = Path(hpo_dir)
    trial_output_files = list(hpo_dir.rglob(f"{trial_id}.json"))
    if not trial_output_files:
        return None
    trial_output_file = trial_output_files[0]

    with trial_output_file.open("r") as f:
        trial_output = json.load(f)

    best_epochs = []
    best_score = None
    for eph, score in trial_output["score"].items():
        if best_score is None:
            best_score = score
            best_epochs.append(eph)
        elif best_score < score:
            best_score = score
            best_epochs = [eph]
        elif best_score == score:
            best_epochs.append(eph)

    best_weight = None
    for best_epoch in best_epochs:
        best_weight_path = list(hpo_dir.glob(f"weight/{trial_id}/*epoch*{best_epoch}*"))
        if best_weight_path:
            best_weight = str(best_weight_path[0])

    return best_weight


class Trainer:
    """Class which prepares and trains a model given hyper parameters.

    Args:
        hp_config (Dict[str, Any]): hyper parameter to use on training
        report_func (Callable): function to report score
        model_template: model template
        data_roots (Dict[str, str]): dataset path of each dataset type
        task_type (TaskType): OTX task type
        hpo_workdir (Union[str, Path]): work directory for HPO
        initial_weight_name (str): initial model weight name for each trials to load
        metric (str): metric name
    """

    # pylint: disable=too-many-arguments, too-many-instance-attributes

    def __init__(
        self,
        hp_config: Dict[str, Any],
        report_func: Callable,
        model_template,
        data_roots: Dict[str, str],
        task_type: TaskType,
        hpo_workdir: Union[str, Path],
        initial_weight_name: str,
        metric: str,
    ):
        self._hp_config = hp_config
        self._report_func = report_func
        self._model_template = model_template
        self._data_roots = data_roots
        self._task = TaskManager(task_type)
        self._hpo_workdir: Path = Path(hpo_workdir)
        self._initial_weight_name = initial_weight_name
        self._metric = metric
        self._epoch = floor(self._hp_config["configuration"]["iterations"])
        del self._hp_config["configuration"]["iterations"]

    def run(self):
        """Run each training of each trial with given hyper parameters."""
        hyper_parameters = self._prepare_hyper_parameter()
        dataset_adapter = self._prepare_dataset_adapter()

        dataset = dataset_adapter.get_otx_dataset()
        dataset = HpoDataset(dataset, self._hp_config)

        label_schema = dataset_adapter.get_label_schema()

        environment = self._prepare_environment(hyper_parameters, label_schema)
        self._set_hyper_parameter(environment)

        need_to_save_initial_weight = False
        resume_weight_path = self._get_resume_weight_path()
        if resume_weight_path is not None:
            environment.resume_model_weight(resume_weight_path, dataset)
        else:
            initial_weight = self._load_fixed_initial_weight()
            if initial_weight is not None:
                environment.load_model_weight(str(initial_weight), dataset)
            else:
                need_to_save_initial_weight = True

        task = environment.get_train_task()
        if need_to_save_initial_weight:
            self._add_initial_weight_saving_hook(task)

        output_model = environment.get_new_model_entity(dataset)
        score_report_callback = self._prepare_score_report_callback(task)
        task.train(dataset=dataset, output_model=output_model, train_parameters=score_report_callback)
        self._finalize_trial(task)

    def _prepare_hyper_parameter(self):
        return create(self._model_template.hyper_parameters.data)

    def _prepare_dataset_adapter(self):
        dataset_adapter = get_dataset_adapter(
            self._task.task_type,
            self._model_template.hyper_parameters.parameter_overrides["algo_backend"]["train_type"]["default_value"],
            train_data_roots=self._data_roots["train_subset"]["data_root"],
            val_data_roots=self._data_roots["val_subset"]["data_root"] if "val_subset" in self._data_roots else None,
            unlabeled_data_roots=self._data_roots["unlabeled_subset"]["data_root"]
            if "unlabeled_subset" in self._data_roots
            else None,
        )

        return dataset_adapter

    def _set_hyper_parameter(self, environment: TaskEnvironmentManager):
        environment.set_hyper_parameter_using_str_key(self._hp_config["configuration"])
        environment.set_epoch(self._epoch)

    def _prepare_environment(self, hyper_parameters, label_schema):
        enviroment = TaskEnvironment(
            model=None,
            hyper_parameters=hyper_parameters,
            label_schema=label_schema,
            model_template=self._model_template,
        )

        return TaskEnvironmentManager(enviroment)

    def _get_resume_weight_path(self):
        trial_work_dir = self._get_weight_dir_path()
        if not trial_work_dir.exists():
            return None
        return self._task.get_latest_weight(trial_work_dir)

    def _load_fixed_initial_weight(self):
        initial_weight_path = self._get_initial_weight_path()
        if initial_weight_path.exists():
            return initial_weight_path
        return None

    def _add_initial_weight_saving_hook(self, task):
        initial_weight_path = self._get_initial_weight_path()
        task.update_override_configurations(
            {
                "custom_hooks": [
                    dict(
                        type="SaveInitialWeightHook",
                        save_path=initial_weight_path.parent,
                        file_name=initial_weight_path.name,
                    )
                ]
            }
        )

    def _prepare_score_report_callback(self, task) -> TrainParameters:
        return TrainParameters(False, HpoCallback(self._report_func, self._metric, self._epoch, task))

    def _get_initial_weight_path(self) -> Path:
        return self._hpo_workdir / self._initial_weight_name

    def _finalize_trial(self, task):
        weight_dir_path = self._get_weight_dir_path()
        weight_dir_path.mkdir(parents=True, exist_ok=True)
        self._task.copy_weight(task.project_path, weight_dir_path)
        self._report_func(0, 0, done=True)

    def _get_weight_dir_path(self) -> Path:
        return self._hpo_workdir / "weight" / self._hp_config["id"]


def run_trial(
    hp_config: Dict[str, Any],
    report_func: Callable,
    model_template,
    data_roots: Dict[str, str],
    task_type: TaskType,
    hpo_workdir: Union[str, Path],
    initial_weight_name: str,
    metric: str,
):
    """Function to train a model given hyper parameters.

    Args:
        hp_config (Dict[str, Any]): hyper parameter to use on training
        report_func (Callable): function to report score
        model_template: model template
        data_roots (Dict[str, str]): dataset path of each dataset type
        task_type (TaskType): OTX task type
        hpo_workdir (Union[str, Path]): work directory for HPO
        initial_weight_name (str): initial model weight name for each trials to load
        metric (str): metric name
    """
    # pylint: disable=too-many-arguments
    trainer = Trainer(
        hp_config, report_func, model_template, data_roots, task_type, hpo_workdir, initial_weight_name, metric
    )
    trainer.run()


class HpoCallback(UpdateProgressCallback):
    """Callback class to report score to HPO.

    Args:
        report_func (Callable): function to report score
        metric (str): metric name
        max_epoch (int): max_epoch
        task: OTX train task
    """

    def __init__(self, report_func: Callable, metric: str, max_epoch: int, task):
        if max_epoch <= 0:
            raise ValueError(f"max_epoch should be bigger than 0. Current value is {max_epoch}.")

        super().__init__()
        self._report_func = report_func
        self.metric = metric
        self._max_epoch = max_epoch
        self._task = task

    def __call__(self, progress: Union[int, float], score: Optional[float] = None):
        """When callback is called, report a score to HPO algorithm."""
        if score is not None:
            epoch = round(self._max_epoch * progress / 100)
            logger.debug(f"In hpo callback : {score} / {progress} / {epoch}")
            if self._report_func(score=score, progress=epoch) == TrialStatus.STOP:
                self._task.cancel_training()

    def __deepcopy__(self, memo):
        """Prevent repot_func from deepcopied."""
        args = [self.metric, self._max_epoch, self._task]
        copied_args = deepcopy(args, memo)
        return self.__class__(self._report_func, *copied_args)


class HpoDataset:
    """Wrapper class for DatasetEntity of dataset. It's used to make subset during HPO.

    Args:
        fullset: full dataset
        config (Optional[Dict[str, Any]], optional): hyper parameter trial config
        indices (Optional[List[int]]): dataset index. Defaults to None.
    """

    def __init__(self, fullset, config: Optional[Dict[str, Any]] = None, indices: Optional[List[int]] = None):
        self.fullset = fullset
        self.indices = indices
        if config is not None:
            subset_ratio = config["train_environment"]["subset_ratio"]
            self.subset_ratio = 1 if subset_ratio is None else subset_ratio

    def __len__(self) -> int:
        """Get length of subset."""
        if self.indices is None:
            return len(self.fullset)
        return len(self.indices)

    def __getitem__(self, indx) -> dict:
        """Get dataset at index."""
        if self.indices is None:
            return self.fullset[indx]
        return self.fullset[self.indices[indx]]

    def __getattr__(self, name):
        """When trying to get other attributes, not dataset, get values from fullset."""
        if name == "__setstate__":
            raise AttributeError(name)
        return getattr(self.fullset, name)

    def get_subset(self, subset: Subset):
        """Get subset according to subset_ratio if training dataset is requested.

        Args:
            subset (Subset): which subset to get

        Returns:
            HpoDataset: subset wrapped by HpoDataset
        """
        dataset = self.fullset.get_subset(subset)
        if subset != Subset.TRAINING or self.subset_ratio > 0.99:
            return dataset

        indices = torch.randperm(len(dataset), generator=torch.Generator().manual_seed(42))
        indices = indices.tolist()  # type: ignore
        indices = indices[: int(len(dataset) * self.subset_ratio)]

        return HpoDataset(dataset, config=None, indices=indices)
