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

# %% auto 0
__all__ = ['get_train_n_test_data', 'df2ds', 'peek', 'find_hyperparameters', 'create_tuner_stats']

# %% ../../nbs/Experiments.ipynb 3
import shutil
import urllib.request
from contextlib import contextmanager
from datetime import datetime
from os import environ
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import *

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pytest
import seaborn as sns
import tensorflow as tf
from keras_tuner import (
    BayesianOptimization,
    HyperModel,
    HyperParameters,
    Objective,
    Tuner,
)
from numpy.typing import ArrayLike, NDArray
from tensorflow.keras import Model
from tensorflow.keras.backend import count_params
from tensorflow.keras.layers import Concatenate, Dense, Dropout, Input
from tensorflow.keras.optimizers.experimental import AdamW
from tensorflow.types.experimental import TensorLike
from tqdm import tqdm

from .layers import MonoDense

# %% ../../nbs/Experiments.ipynb 7
class _DownloadProgressBar(tqdm):
    def update_to(
        self, b: int = 1, bsize: int = 1, tsize: Optional[int] = None
    ) -> None:
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)


def _download_url(url: str, output_path: Path) -> None:
    with _DownloadProgressBar(
        unit="B", unit_scale=True, miniters=1, desc=url.split("/")[-1]
    ) as t:
        # nosemgrep: python.lang.security.audit.dynamic-urllib-use-detected.dynamic-urllib-use-detected
        urllib.request.urlretrieve(
            url, filename=output_path, reporthook=t.update_to
        )  # nosec

# %% ../../nbs/Experiments.ipynb 8
def _get_data_path(data_path: Optional[Union[Path, str]] = None) -> Path:
    if data_path is None:
        data_path = "./data"
    return Path(data_path)


def _download_data(
    dataset_name: str,
    data_path: Optional[Union[Path, str]] = "data",
    force_download: bool = False,
) -> None:
    data_path = _get_data_path(data_path)
    data_path.mkdir(exist_ok=True, parents=True)

    for prefix in ["train", "test"]:
        filename = f"{prefix}_{dataset_name}.csv"
        if not (data_path / filename).exists() or force_download:
            with TemporaryDirectory() as d:
                _download_url(
                    f"https://zenodo.org/record/7968969/files/{filename}",
                    Path(d) / filename,
                )
                shutil.copyfile(Path(d) / filename, data_path / filename)
        else:
            print(f"Upload skipped, file {(data_path / filename).resolve()} exists.")

# %% ../../nbs/Experiments.ipynb 10
def _sanitize_col_names(df: pd.DataFrame) -> pd.DataFrame:
    columns = {c: c.replace(" ", "_") for c in df}
    df = df.rename(columns=columns)
    return df

# %% ../../nbs/Experiments.ipynb 12
def get_train_n_test_data(
    dataset_name: str,
    *,
    data_path: Optional[Union[Path, str]] = "./data",
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Download data

    Args:
        dataset_name: name of the dataset, one of "auto", "heart", compas", "blog", "loan"
        data_path: root directory where to download data to
    """
    data_path = _get_data_path(data_path)
    _download_data(dataset_name=dataset_name, data_path=data_path)

    dfx = [
        pd.read_csv(data_path / f"{prefix}_{dataset_name}.csv")
        for prefix in ["train", "test"]
    ]
    dfx = [_sanitize_col_names(df) for df in dfx]
    return dfx[0], dfx[1]

# %% ../../nbs/Experiments.ipynb 14
def df2ds(df: pd.DataFrame) -> tf.data.Dataset:
    """Converts DataFrame to Dataset

    Args:
        df: input DataFrame

    Returns:
        dataset
    """
    x = df.to_dict("list")
    y = x.pop("ground_truth")

    ds = tf.data.Dataset.from_tensor_slices((x, y))

    return ds


def peek(ds: tf.data.Dataset) -> tf.Tensor:
    """Returns the first element of the dataset

    Args:
        ds: dataset

    Returns:
        the first element of the dataset
    """
    for x in ds:
        return x

# %% ../../nbs/Experiments.ipynb 16
def _build_mono_model_f(
    *,
    monotonicity_indicator: Dict[str, int],
    final_activation: Union[str, Callable[[TensorLike], TensorLike]],
    loss: Union[str, Callable[[TensorLike, TensorLike], TensorLike]],
    metrics: Union[str, Callable[[TensorLike, TensorLike], TensorLike]],
    train_ds: tf.data.Dataset,
    batch_size: int,
    units: int,
    n_layers: int,
    activation: Union[str, Callable[[TensorLike], TensorLike]],
    learning_rate: float,
    weight_decay: float,
    dropout: float,
    decay_rate: float,
) -> Model:
    inputs = {k: Input(name=k, shape=(1,)) for k in monotonicity_indicator.keys()}
    outputs = MonoDense.create_type_2(
        inputs,
        units=units,
        final_units=1,
        activation=activation,
        n_layers=n_layers,
        monotonicity_indicator=monotonicity_indicator,
        is_convex=False,
        is_concave=False,
        dropout=dropout,
        final_activation=final_activation,
    )
    model = Model(inputs=inputs, outputs=outputs)

    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        learning_rate,
        decay_steps=len(train_ds.batch(batch_size)),
        decay_rate=decay_rate,
        staircase=True,
    )

    optimizer = AdamW(learning_rate=lr_schedule, weight_decay=weight_decay)
    model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

    return model

# %% ../../nbs/Experiments.ipynb 18
def _get_build_model_with_hp_f(
    build_model_f: Callable[[], Model],
    hp_params_f: Optional[Callable[[HyperParameters], Dict[str, Any]]] = None,
    **kwargs: Any,
) -> Callable[[HyperParameters], Model]:
    def build_model_with_hp_f(
        hp: HyperParameters,
        hp_params_f: Optional[
            Callable[[HyperParameters], Dict[str, Any]]
        ] = hp_params_f,
        kwargs: Dict[str, Any] = kwargs,
    ) -> Model:
        override_kwargs = hp_params_f(hp) if hp_params_f is not None else {}

        default_kwargs = dict(
            units=hp.Int("units", min_value=8, max_value=32, step=1),
            n_layers=hp.Int("n_layers", min_value=1, max_value=4),
            activation=hp.Choice("activation", values=["elu"]),
            learning_rate=hp.Float(
                "learning_rate", min_value=1e-3, max_value=0.3, sampling="log"
            ),
            weight_decay=hp.Float(
                "weight_decay", min_value=1e-1, max_value=0.3, sampling="log"
            ),
            dropout=hp.Float(
                "dropout", min_value=0.0, max_value=0.5, sampling="linear"
            ),
            decay_rate=hp.Float(
                "decay_rate", min_value=0.5, max_value=1.0, sampling="reverse_log"
            ),
        )

        default_kwargs.update(**override_kwargs)
        model = build_model_f(**default_kwargs, **kwargs)
        return model

    return build_model_with_hp_f


class TestHyperModel(HyperModel):
    def __init__(self, **kwargs: Any):
        self.kwargs = kwargs

    def build(self, hp: HyperParameters) -> Model:
        build_model_with_hp_f = _get_build_model_with_hp_f(
            _build_mono_model_f, **self.kwargs  # type: ignore
        )
        return build_model_with_hp_f(hp)

# %% ../../nbs/Experiments.ipynb 20
def find_hyperparameters(
    dataset_name: str,
    *,
    monotonicity_indicator: Dict[str, int],
    final_activation: Union[str, Callable[[TensorLike, TensorLike], TensorLike]],
    loss: Union[str, Callable[[TensorLike, TensorLike], TensorLike]],
    metrics: Union[str, Callable[[TensorLike, TensorLike], TensorLike]],
    hp_params_f: Optional[Callable[[HyperParameters], Dict[str, Any]]] = None,
    max_trials: int = 100,
    max_epochs: int = 50,
    batch_size: int = 8,
    objective: Union[str, Objective],
    direction: str,
    dir_root: Union[Path, str] = "tuner",
    seed: int = 42,
    executions_per_trial: int = 3,
    max_consecutive_failed_trials: int = 5,
    patience: int = 10,
) -> Tuner:
    """Search for optimal hyperparameters

    Args:
        monotonicity_indicator: monotonicity indicator as used in `MonoDense.__init__`
        final_activation:  final activation of the neural network
        loss: Tensorflow loss function
        metrics: Tensorflow metrics function
        hp_params_f: a function constructing sampling hyperparameters using Keras Tuner
        max_trials: maximum number of trials
        max_epochs: maximum number of epochs in each trial
        batch_size: batch size
        objective: objective, typically f"val_{metrics}"
        direction: direction of the objective, either "min" or "max"
        dir_root: root directory for storing Keras Tuner data
        seed: random seed used to guarantee reproducibility of results
        executions_per_trial: number of executions per trial. Set it to number higher than zero for small datasets
        max_consecutive_failed_trials: maximum number of failed trials as used in Keras Tuner
        patience: number of epoch with worse objective before stopping trial early

    Returns:
        An instance of Keras Tuner

    """
    tf.keras.utils.set_random_seed(seed)

    train_df, test_df = get_train_n_test_data(dataset_name)
    train_ds, test_ds = df2ds(train_df), df2ds(test_df)

    oracle = TestHyperModel(
        monotonicity_indicator=monotonicity_indicator,
        hp_params_f=hp_params_f,
        final_activation=final_activation,
        loss=loss,
        metrics=metrics,
        train_ds=train_ds,
        batch_size=batch_size,
    )

    tuner = BayesianOptimization(
        oracle,
        objective=Objective(objective, direction),
        max_trials=max_trials,
        seed=seed,
        directory=Path(dir_root),
        project_name=dataset_name,
        executions_per_trial=executions_per_trial,
        max_consecutive_failed_trials=max_consecutive_failed_trials,
    )

    stop_early = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=patience)

    tuner.search(
        train_ds.shuffle(len(train_ds)).batch(batch_size).prefetch(2),
        validation_data=test_ds.batch(256),
        callbacks=[stop_early],
        epochs=max_epochs,
    )

    return tuner

# %% ../../nbs/Experiments.ipynb 22
def _count_model_params(model: Model) -> int:
    return sum([sum([count_params(v) for v in l.variables]) for l in model.layers])


def _create_model_stats(
    tuner: Tuner,
    hp: Dict[str, Any],
    *,
    stats: Optional[pd.DataFrame] = None,
    max_epochs: int,
    num_runs: int,
    top_runs: int,
    batch_size: int,
    patience: int,
    verbose: int,
    train_ds: tf.data.Dataset,
    test_ds: tf.data.Dataset,
) -> pd.DataFrame:
    tf.keras.utils.set_random_seed(42)

    def model_stats(
        tuner: Tuner = tuner,
        hp: Dict[str, Any] = hp,
        max_epochs: int = max_epochs,
        batch_size: int = batch_size,
        patience: int = patience,
        verbose: int = verbose,
        train_ds: tf.data.Dataset = train_ds,
        test_ds: tf.data.Dataset = test_ds,
    ) -> float:
        model = tuner.hypermodel.build(hp)
        stop_early = tf.keras.callbacks.EarlyStopping(
            monitor="val_loss", patience=patience
        )
        history = model.fit(
            train_ds.shuffle(len(train_ds)).batch(batch_size).prefetch(2),
            epochs=max_epochs,
            validation_data=test_ds.batch(256),
            verbose=verbose,
            callbacks=[stop_early],
        )
        objective = history.history[tuner.oracle.objective.name]
        if tuner.oracle.objective.direction == "max":
            best_epoch = objective.index(max(objective))
        else:
            best_epoch = objective.index(min(objective))
        return objective[best_epoch]  # type: ignore

    xs = sorted(
        [model_stats() for _ in range(num_runs)],
        reverse=tuner.oracle.objective.direction == "max",
    )
    stats = pd.Series(xs[:top_runs])
    stats = stats.describe()
    stats = {
        f"{tuner.oracle.objective.name}_{k}": stats[k]
        for k in ["mean", "std", "min", "max"]
    }
    model = tuner.hypermodel.build(hp)
    stats_df = pd.DataFrame(
        dict(**hp.values, **stats, params=_count_model_params(model)),  # type: ignore
        index=[0],
    )
    return stats_df

# %% ../../nbs/Experiments.ipynb 23
def create_tuner_stats(
    tuner: Tuner,
    *,
    num_models: int = 10,
    max_epochs: int = 50,
    batch_size: int = 8,
    patience: int = 10,
    verbose: int = 0,
) -> pd.DataFrame:
    """Calculates statistics for the best models found by Keras Tuner

    Args:
        tuner: an instance of Keras Tuner
        num_models: number of best models to use for calculating statistics
        max_epochs: maximum number of epochs used in runs
        batch_size: batch_size
        patience: maximum number of epochs with worse objective before stopping trial early
        verbose: verbosity level of `Model.fit` function

    Returns:
        A dataframe with statistics
    """
    stats = None

    train_df, test_df = get_train_n_test_data(tuner.project_name)
    train_ds, test_ds = df2ds(train_df), df2ds(test_df)

    for hp in tuner.get_best_hyperparameters(num_trials=num_models):
        new_entry = _create_model_stats(
            tuner,
            hp,
            stats=stats,
            max_epochs=max_epochs,
            num_runs=10,
            top_runs=5,
            batch_size=batch_size,
            patience=patience,
            verbose=verbose,
            train_ds=train_ds,
            test_ds=test_ds,
        )
        if stats is None:
            stats = new_entry
        else:
            stats = pd.concat([stats, new_entry]).reset_index(drop=True)

        try:
            display(stats.sort_values(f"{tuner.oracle.objective.name}_mean"))  # type: ignore
        # nosemgrep
        except Exception as e:  # nosec
            pass

    return stats.sort_values(f"{tuner.oracle.objective.name}_mean")  # type: ignore
