# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/common.base_auto.ipynb.

# %% auto 0
__all__ = ['BaseAuto']

# %% ../../nbs/common.base_auto.ipynb 4
from copy import deepcopy
from os import cpu_count

import torch
from pytorch_lightning.callbacks import TQDMProgressBar
from ray import air, tune
from ray.tune.integration.pytorch_lightning import TuneReportCallback
from ray.tune.search.basic_variant import BasicVariantGenerator

from ..losses.pytorch import MAE

# %% ../../nbs/common.base_auto.ipynb 5
def train_tune(config_step, cls_model, dataset, val_size, test_size):
    metrics = {"loss": "ptl/val_loss"}
    callbacks = [TQDMProgressBar(), TuneReportCallback(metrics, on="validation_end")]
    if 'callbacks' in config_step.keys():
        callbacks += config_step['callbacks']
    config_step = {**config_step, **{'callbacks': callbacks}}
    model = cls_model(**config_step)
    model.fit(
        dataset,
        val_size=val_size, 
        test_size=test_size
    )

# %% ../../nbs/common.base_auto.ipynb 6
def tune_model(
        cls_model, 
        dataset, 
        val_size, 
        test_size,
        cpus,
        gpus,
        verbose,
        num_samples, 
        search_alg, 
        config
    ):
    train_fn_with_parameters = tune.with_parameters(
        train_tune,
        cls_model=cls_model,
        dataset=dataset,
        val_size=val_size,
        test_size=test_size,
    )

    # Device
    if gpus > 0:
        device_dict = {'gpu':gpus}
    else:
        device_dict = {'cpu':cpus}

    tuner = tune.Tuner(
        tune.with_resources(train_fn_with_parameters, device_dict),
        run_config=air.RunConfig(
            verbose=verbose,
            #checkpoint_config=air.CheckpointConfig(
                #num_to_keep=0,
                #keep_checkpoints_num=None
            #)
        ),
        tune_config=tune.TuneConfig(
            metric="loss",
            mode="min",
            num_samples=num_samples, 
            search_alg=search_alg
        ),
        param_space=config,
    )
    results = tuner.fit()
    return results

# %% ../../nbs/common.base_auto.ipynb 7
# If None overwrite with default


# %% ../../nbs/common.base_auto.ipynb 8
class BaseAuto:
    
    def __init__(self, 
                 cls_model,
                 h,
                 config, 
                 search_alg=BasicVariantGenerator(random_state=1),
                 num_samples=10,
                 cpus=cpu_count(),
                 gpus=torch.cuda.device_count(),
                 refit_wo_val=False,
                 verbose=False):
        
        config['h'] = h
        self.cls_model = cls_model
        self.h = h
        self.config = config
        self.num_samples = num_samples
        self.search_alg = search_alg
        self.cpus = cpus
        self.gpus = gpus
        self.refit_wo_val = refit_wo_val
        self.verbose = verbose
        self.loss = self.config.get('loss', MAE())
        
    def fit(self, dataset, val_size=0, test_size=0):
        #we need val_size > 0 to perform
        #hyperparameter selection.
        search_alg = deepcopy(self.search_alg)
        val_size = val_size if val_size > 0 else self.h
        results = tune_model(
            cls_model=self.cls_model,
            dataset=dataset,
            val_size=val_size, 
            test_size=test_size, 
            cpus=self.cpus,
            gpus=self.gpus,
            verbose=self.verbose,
            num_samples=self.num_samples, 
            search_alg=search_alg, 
            config=self.config
        )
        best_config = results.get_best_result().config
        self.model = self.cls_model(**best_config)
        self.model.fit(
            dataset=dataset, 
            val_size=val_size * (1 - self.refit_wo_val), 
            test_size=test_size,
        )
        self.results = results
        
    def predict(self, dataset, step_size=1, **data_kwargs):
        return self.model.predict(dataset=dataset, 
                                  step_size=step_size, **data_kwargs)
