#!/bin/python3

import numpy as np

from smac.scenario.scenario import Scenario
from smac.tae.execute_ta_run import StatusType

from smac.runhistory.runhistory import RunHistory
from smac.runhistory.runhistory2epm import RunHistory2EPM4LogCost, RunHistory2EPM4Cost
from smac.utils.util_funcs import get_types

from smac.epm.rf_with_instances import RandomForestWithInstances
from smac.epm.rfr_imputator import RFRImputator


def convert_data_for_epm(scenario: Scenario, runhistory: RunHistory, logger=None):
    """
    converts data from runhistory into EPM format

    Parameters
    ----------
    scenario: Scenario
        smac.scenario.scenario.Scenario Object
    runhistory: RunHistory
        smac.runhistory.runhistory.RunHistory Object with all necessary data

    Returns
    -------
    X: np.array
        X matrix with configuartion x features for all observed samples
    y: np.array
        y matrix with all observations
    types: np.array
        types of X cols -- necessary to train our RF implementation
    """
    types, bounds = get_types(scenario.cs, scenario.feature_array)
    model = RandomForestWithInstances(types, bounds)

    params = scenario.cs.get_hyperparameters()
    num_params = len(params)

    run_obj = scenario.run_obj

    if run_obj == "runtime":
        # if we log the performance data,
        # the RFRImputator will already get
        # log transform data from the runhistory
        cutoff = np.log10(scenario.cutoff)
        threshold = np.log10(scenario.cutoff *
                             scenario.par_factor)

        imputor = RFRImputator(rng=np.random.RandomState(42),
                               cutoff=cutoff,
                               threshold=threshold,
                               model=model,
                               change_threshold=0.01,
                               max_iter=10)
        # TODO: Adapt runhistory2EPM object based on scenario
        rh2EPM = RunHistory2EPM4LogCost(scenario=scenario,
                                        num_params=num_params,
                                        success_states=[
                                            StatusType.SUCCESS, ],
                                        impute_censored_data=True,
                                        impute_state=[
                                            StatusType.TIMEOUT, ],
                                        imputor=imputor)
        X, Y = rh2EPM.transform(runhistory)
    else:
        rh2EPM = RunHistory2EPM4Cost(scenario=scenario,
                                     num_params=num_params,
                                     success_states=None,
                                     impute_censored_data=False,
                                     impute_state=None)
        X, Y = rh2EPM.transform(runhistory)

    return X, Y, types
