# Copyright (c) 2015, Ecole Polytechnique Federale de Lausanne, Blue Brain Project
# All rights reserved.
#
# This file is part of NeuroM <https://github.com/BlueBrain/NeuroM>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     1. Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#     2. Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#     3. Neither the name of the copyright holder nor the names of
#        its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Core code for morph_stats application."""
import os
import logging
from collections import defaultdict
from itertools import product
from pathlib import Path
import multiprocessing
from functools import partial
import warnings

import numpy as np
import pandas as pd
import pkg_resources

import neurom as nm
from neurom.exceptions import ConfigError
from neurom.features import NEURITEFEATURES, NEURONFEATURES
from neurom.fst._core import FstNeuron

L = logging.getLogger(__name__)

EXAMPLE_CONFIG = Path(pkg_resources.resource_filename(
    'neurom', 'config'), 'morph_stats.yaml')

def eval_stats(values, mode):
    """Extract a summary statistic from an array of list of values.

    Arguments:
        values: A numpy array of values
        mode: A summary stat to extract. One of:
            ['min', 'max', 'median', 'mean', 'std', 'raw', 'total']

    Note: fails silently if values is empty, and None is returned
    """
    if mode == 'raw':
        return values.tolist()
    if mode == 'total':
        mode = 'sum'

    try:
        return getattr(np, mode)(values, axis=0)
    except ValueError:
        pass

    return None

def _stat_name(feat_name, stat_mode):
    """Set stat name based on feature name and stat mode."""
    if feat_name[-1] == 's':
        feat_name = feat_name[:-1]
    if feat_name == 'soma_radii':
        feat_name = 'soma_radius'
    if stat_mode == 'raw':
        return feat_name

    return '%s_%s' % (stat_mode, feat_name)

def _run_extract_stats(nrn, config):
    """The function to be called by multiprocessing.Pool.imap_unordered."""
    if not isinstance(nrn, FstNeuron):
        nrn = nm.load_neuron(nrn)
    return nrn.name, extract_stats(nrn, config)

def extract_dataframe(neurons, config, n_workers=1):
    """Extract stats grouped by neurite type from neurons.

    Arguments:
        neurons: a neuron, population, neurite tree or list of neuron paths
        config (dict): configuration dict. The keys are:
            - neurite_type: a list of neurite types for which features are extracted
              If not provided, all neurite_type will be used
            - neurite: a dictionary {{neurite_feature: mode}} where:
                - neurite_feature is a string from NEURITEFEATURES
                - mode is an aggregation operation provided as a string such as:
                  ['min', 'max', 'median', 'mean', 'std', 'raw', 'total']
        n_workers (int): number of workers for multiprocessing (on collection of neurons)

    Returns:
        The extracted statistics

    Note:
        An example config can be found at:

    {config_path}
    """
    if isinstance(neurons, FstNeuron):
        neurons = [neurons]
    config = config.copy()

    # Only NEURITEFEATURES are considered since the dataframe is built by neurite_type
    # NEURONFEATURES are discarded
    if 'neuron' in config:
        del config['neuron']

    func = partial(_run_extract_stats, config=config)
    if n_workers == 1:
        stats = dict(map(func, neurons))
    else:
        if n_workers > os.cpu_count():
            warnings.warn(f'n_workers ({n_workers}) > os.cpu_count() ({os.cpu_count()}))')
        with multiprocessing.Pool(n_workers) as pool:
            stats = dict(pool.imap_unordered(func, neurons))

    columns = list(next(iter(next(iter(stats.values())).values())).keys())

    rows = [[name, neurite_type] + list(features.values())
            for name, data in stats.items()
            for neurite_type, features in data.items()]
    return pd.DataFrame(columns=['name', 'neurite_type'] + columns,
                        data=rows)

def extract_stats(neurons, config):
    """Extract stats from neurons.

    Arguments:
        neurons: a neuron, population, neurite tree or list of neuron paths/str
        config (dict): configuration dict. The keys are:
            - neurite_type: a list of neurite types for which features are extracted
              If not provided, all neurite_type will be used
            - neurite: a dictionary {{neurite_feature: mode}} where:
                - neurite_feature is a string from NEURITEFEATURES
                - mode is an aggregation operation provided as a string such as:
                  ['min', 'max', 'median', 'mean', 'std', 'raw', 'total']

    Returns:
        The extracted statistics

    Note:
        An example config can be found at:

    {config_path}
    """

    def _fill_stats_dict(data, stat_name, stat):
        """Insert the stat data in the dict.

        And if the stats is a 3D array, splits it into XYZ components.
        """
        if stat is None or not isinstance(stat, np.ndarray) or stat.shape not in ((3, ), ):
            data[stat_name] = stat
        else:
            for i, suffix in enumerate('XYZ'):
                compound_stat_name = stat_name + '_' + suffix
                data[compound_stat_name] = stat[i]

    stats = defaultdict(dict)

    for (feature_name, modes), neurite_type in product(config['neurite'].items(),
                                                       config.get('neurite_type',
                                                                  _NEURITE_MAP.keys())):
        neurite_type = _NEURITE_MAP[neurite_type]
        feature = nm.get(feature_name, neurons, neurite_type=neurite_type)
        for mode in modes:
            stat_name = _stat_name(feature_name, mode)
            stat = eval_stats(feature, mode)
            _fill_stats_dict(stats[neurite_type.name], stat_name, stat)

    for feature_name, modes in config.get('neuron', {}).items():
        feature = nm.get(feature_name, neurons)
        for mode in modes:
            stat_name = _stat_name(feature_name, mode)
            stat = eval_stats(feature, mode)
            _fill_stats_dict(stats, stat_name, stat)

    return dict(stats)

def get_header(results):
    """Extracts the headers, using the first value in the dict as the template."""
    ret = ['name', ]
    values = next(iter(results.values()))
    for k, v in values.items():
        if isinstance(v, dict):
            for metric in v.keys():
                ret.append('%s:%s' % (k, metric))
        else:
            ret.append(k)
    return ret

def generate_flattened_dict(headers, results):
    """Extract from results the fields in the headers list."""
    for name, values in results.items():
        row = []
        for header in headers:
            if header == 'name':
                row.append(name)
            elif ':' in header:
                neurite_type, metric = header.split(':')
                row.append(values[neurite_type][metric])
            else:
                row.append(values[header])
        yield row

_NEURITE_MAP = {
    'AXON': nm.AXON,
    'BASAL_DENDRITE': nm.BASAL_DENDRITE,
    'APICAL_DENDRITE': nm.APICAL_DENDRITE,
    'ALL': nm.ANY_NEURITE
}

def full_config():
    """Returns a config with all features, all modes, all neurite types."""
    modes = ['min', 'max', 'median', 'mean', 'std']
    return {
        'neurite': {feature: modes for feature in NEURITEFEATURES},
        'neuron': {feature: modes for feature in NEURONFEATURES},
        'neurite_type': list(_NEURITE_MAP.keys()),
    }

def sanitize_config(config):
    """Check that the config has the correct keys, add missing keys if necessary."""
    if 'neurite' in config:
        if 'neurite_type' not in config:
            raise ConfigError('"neurite_type" missing from config, but "neurite" set')
    else:
        config['neurite'] = {}

    if 'neuron' not in config:
        config['neuron'] = {}

    return config

extract_stats.__doc__ = extract_stats.__doc__.format(config_path=EXAMPLE_CONFIG)
extract_dataframe.__doc__ = extract_dataframe.__doc__.format(config_path=EXAMPLE_CONFIG)


# -*- coding: utf-8 -*-
# pyuic5 -o SegGroundClassUI.py SegGroundClassUI.ui
import sys
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

from SegGroundClassUI import Ui_MainWindow

from paddleseg.models import BiSeNetV2, UNet
import paddleseg.transforms as T
from paddleseg.core import infer
import paddle
import numpy as np
import SimpleITK as sitk
import os
import cv2
from utils.utils import readNii
from widgets.canvas import Canvas

class InferThread(QThread):
    """
    建立一个任务线程类,  推理任务
    """
    signal_infer_fail = pyqtSignal()  # 推理失败的信号
    signal_infer_result = pyqtSignal(np.ndarray)  # 这信号用来传递推理结果

    def __init__(self, sitkImage, model):
        super(InferThread, self).__init__()
        self.sitkImage = sitkImage
        self.model = model
        self.transforms = T.Compose([
            T.Resize(target_size=(512, 512)),
            T.Normalize()
        ])

    def run(self):  # 在启动线程后任务从这个函数里面开始执行
        try:
            data = readNii(self.sitkImage, 1500, -500)
            inferData = np.zeros_like(data)
            d, h, w = data.shape

            for i in range(d):
                img = data[i].copy()
                img = img.astype(np.float32)
                pre = self.nn_infer(self.model, img, self.transforms)
                inferData[i] = pre

            self.signal_infer_result.emit(inferData)
        except Exception as e:
            print(e)
            self.signal_infer_fail.emit()

    def nn_infer(self, model, im, transforms):
        # 预测结果

        img, _ = transforms(im)
        img = paddle.to_tensor(img[np.newaxis, :])
        pre = infer.inference(model, img)
        pred = paddle.argmax(pre, axis=1).numpy().reshape((512, 512))
        return pred.astype('uint8')

class MainWindow(QMainWindow,Ui_MainWindow):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.setupUi(self)
        self.canvas = Canvas(self.graphicsView)
        self.canvas.wheeled.connect(self.myWheel)

        self.initUI()
        self.setWindowTitle('结合PYQT5部署PaddleSeg模型')
        self.bn_open.clicked.connect(self.openFile)  # 打开 nii文件选择器
        self.bn_loadModel.clicked.connect(self.openModleFile)  # 打开模型文件选择器
        self.bn_infer.clicked.connect(self.infer)  # 推理按钮
        self.bn_output.clicked.connect(self.outputFile)

        self.sitkImage = object()
        self.npImage = object()
        self.currIndex = 0  # 记录当前第几层
        self.maxCurrIndex = 0  # 记录数据的最大层
        self.minCurrIndex = 0  # 记录数据的最小层，其实就是0
        self.baseFileName = ''
        self.isInferSucceed = False

        self.model = object()
        # 判断模型是否加载成功
        self.isModelReady = False

        # 宽宽窗位滑动条
        self.slider_ww.valueChanged.connect(self.resetWWWcAndShow)
        self.slider_wc.valueChanged.connect(self.resetWWWcAndShow)
        # 窗宽窗位下来框选择器
        self.cb_wwwc.currentIndexChanged.connect(self.resetWWWcAndShow)

        # 设置窗宽窗位文本框只能输入一定范围的整数
        intValidator = QIntValidator(self)
        intValidator.setRange(-2000, 2000)
        self.line_ww.setValidator(intValidator)
        self.line_ww.editingFinished.connect(self.resetWWWcAndShow)
        self.line_wc.setValidator(intValidator)
        self.line_wc.editingFinished.connect(self.resetWWWcAndShow)

        self.listWidget.itemDoubleClicked.connect(self.changeLayer)

    def initUI(self):
        try:
            self.wwwcList = {'软组织窗': [350, 80],
                             "纵隔窗": [300, 40],
                             "脑窗": [100, 40],
                             '肺窗': [1700, -700],
                             '骨窗': [1400, 350]}

            windowWidth = self.cb_wwwc.currentText()
            self.line_ww.setText(str(self.wwwcList[windowWidth][0]))
            self.line_wc.setText(str(self.wwwcList[windowWidth][1]))

            self.slider_ww.setValue(self.wwwcList[windowWidth][0])
            self.slider_wc.setValue(self.wwwcList[windowWidth][1])
            self.ww = self.wwwcList[windowWidth][0]
            self.wc = self.wwwcList[windowWidth][1]

            self.currWw = self.ww
            self.currWc = self.wc
        except Exception as e:
            print(e)

    def openFile(self):
        """
        打开医学影像文件选择器
        """
        try:
            filename, _ = QFileDialog.getOpenFileName(self,
                                                      "选取文件",
                                                      "./",
                                                      "Nii Files (*.nii);;Nii Files (*.nii.gz);;All Files (*)")
            if filename:
                self.listWidget.clear()  # 清空列表
                self.isInferSucceed = False
                self.text_loadModel.setText("数据加载成功！")

                self.baseFileName = os.path.basename(filename).split('.')[0]
                self.sitkImage = sitk.ReadImage(filename)

                self.npImage = readNii(self.sitkImage, self.ww, self.wc)
                print(self.npImage.shape)
                self.maxCurrIndex = self.npImage.shape[0]
                self.currIndex = int(self.maxCurrIndex / 2)
                self.showImg(self.npImage[self.currIndex])
        except Exception as e:
            print(e)

    def openModleFile(self):
        """
        打开模型文件选择器
        """
        filename, _ = QFileDialog.getOpenFileName(self, "选取文件", "./", "model Files (*.pdparams)")

        if filename:
            try:
                self.isInferSucceed = False
                self.text_loadModel.setText(" ")
                model_name = self.cb_modelList.currentText()
                num_class = int(self.spinBox_numClass.value())
                if model_name == "BiSeNetV2":
                    self.model = BiSeNetV2(num_classes=num_class)
                elif model_name == "UNet":
                    self.model = UNet(num_classes=num_class)
                para_state_dict = paddle.load(filename)
                self.model.set_dict(para_state_dict)
                self.text_loadModel.setText("模型加载成功！")
                self.isModelReady = True
            except Exception as e:
                self.text_loadModel.setText("模型加载失败！")
                print(e)

    def showImg(self, img):
        """
        显示图片
        """
        try:
            if img.ndim == 2:
                img = np.expand_dims(img, axis=2)
                img = np.concatenate((img, img, img), axis=-1).astype(np.uint8)
            elif img.ndim == 3:
                img = img.astype(np.uint8)
            qimage = QImage(img, img.shape[0], img.shape[1], img.shape[1] * 3, QImage.Format_RGB888)
            pixmap_imgSrc = QPixmap.fromImage(qimage)

            self.canvas.addScenes(pixmap_imgSrc,self.currIndex)
        except Exception as e:
            print(e)

    def resetWWWcAndShow(self):
        """
        有四个方式可以修改医学图像的窗宽窗位，
        每次修改后都会在界面呈现出来
        """
        if hasattr(self.sender(), "objectName"):
            objectName = self.sender().objectName()
        else:
            objectName = None
        try:

            if objectName == 'cb_wwwc':
                windowWidth = self.cb_wwwc.currentText()
                self.line_ww.setText(str(self.wwwcList[windowWidth][0]))
                self.line_wc.setText(str(self.wwwcList[windowWidth][1]))
                self.slider_ww.setValue(self.wwwcList[windowWidth][0])
                self.slider_wc.setValue(self.wwwcList[windowWidth][1])
                self.ww = self.wwwcList[windowWidth][0]
                self.wc = self.wwwcList[windowWidth][1]
                self.currWw = self.ww
                self.currWc = self.wc
            elif objectName == 'slider_ww' or objectName == 'slider_wc':
                self.currWw = self.slider_ww.value()
                self.currWc = self.slider_wc.value()
                self.line_ww.setText(str(self.currWw))
                self.line_wc.setText(str(self.currWc))
            elif objectName == 'line_ww' or objectName == 'line_wc':
                self.currWw = int(self.line_ww.text())
                self.currWc = int(self.line_wc.text())
                self.slider_ww.setValue(self.currWw)
                self.slider_wc.setValue(self.currWc)
            if self.maxCurrIndex != self.minCurrIndex:
                self.npImage = readNii(self.sitkImage, self.currWw, self.currWc)
                if self.isInferSucceed:
                    self.showImg(self.drawContours(self.npImage, self.inferData, self.currIndex))
                else:
                    self.showImg(self.npImage[self.currIndex])
        except Exception as e:
            print(e)

    def infer(self):
        """
        模型分割预测
        """
        if self.maxCurrIndex != self.minCurrIndex and self.isModelReady:
            self.bn_infer.setEnabled(True)
            # 创建推理线程
            self.infer_thread = InferThread(self.sitkImage, self.model)
            # 绑定推理失败的槽函数
            self.infer_thread.signal_infer_fail.connect(self.infer_fail)
            # 绑定推理成功的槽函数
            self.infer_thread.signal_infer_result.connect(self.infer_result)
            self.infer_thread.start()
            self.text_loadModel.setText("模型推理中！")

        else:
            QMessageBox.warning(self, "警告", "请加载模型或者加载数据再进行推理", QMessageBox.Yes, QMessageBox.Yes)

    def infer_result(self, inferData):
        """
        分割模型预测成功后，结果保存在self.inferData
        """
        # 推理成功，并显示结果
        try:
            self.inferData = inferData.astype(np.uint8)
            QMessageBox.information(self, "信息", "推理完成！", QMessageBox.Yes, QMessageBox.Yes)
            self.text_loadModel.setText("模型推理成功！")
            self.isInferSucceed = True
            self.infer_thread.quit()
            self.addListInfo(self.inferData)
            self.showImg(self.drawContours(self.npImage, self.inferData, self.currIndex))
        except Exception as e:
            print(e)

    def infer_fail(self):
        """
        推理失败的情况
        """
        QMessageBox.warning(self, "警告", "推理失败！", QMessageBox.Yes, QMessageBox.Yes)

    def outputFile(self):
        """
        保存模型预测结果为nii格式文件
        """
        try:
            if self.isInferSucceed:
                filedir = QFileDialog.getExistingDirectory(None, "文件保存", os.getcwd())
                if filedir:
                    #因为读取nii文件转换np文件时对数据进行上下翻转，才输入模型推理的
                    #所有现在保存回nii文件，要翻转回来。
                    #不知道为什么这个肺部CT数据用SimpleITK读取是上下翻转的。
                    self.inferData = np.flip(self.inferData, 1)
                    pre_sitkImage = sitk.GetImageFromArray(self.inferData)
                    pre_sitkImage.CopyInformation(self.sitkImage)
                    pre_sitkImage = sitk.Cast(pre_sitkImage, sitk.sitkUInt8)
                    save_path = os.path.join(filedir, self.baseFileName + '_mask.nii')
                    sitk.WriteImage(pre_sitkImage, save_path)
            else:
                QMessageBox.warning(self, "警告", "无进行过推理，无法保存！", QMessageBox.Yes, QMessageBox.Yes)
        except Exception as e:
            print(e)

    def drawContours(self, npImage, inferData, currIndex):
        """
        把mask转换成轮廓绘制在原图上
        """
        img = npImage[currIndex]
        img = np.expand_dims(img, axis=2)
        img = np.concatenate((img, img, img), axis=-1).astype(np.uint8)
        ret, thresh = cv2.threshold(inferData[currIndex], 0, 255, cv2.THRESH_BINARY)
        thresh = cv2.dilate(thresh, kernel=np.ones((5, 5), np.uint8), iterations=1)
        contours, hierarchy = cv2.findContours(thresh, 1, 2)
        # 这是画轮廓
        img = cv2.drawContours(img, contours, -1, (0, 255, 0), 1)

        return img

    def addListInfo(self, inferData):
        """
        增加列表信息
        """
        self.listWidget.clear()
        d, h, w = inferData.shape
        result = {}
        for i in range(d):
            img = inferData[i]
            if np.sum(img > 0) != 0:
                result[str(i)] = np.sum(img > 0)

        result = sorted(result.items(), key=lambda x: x[1], reverse=True)
        for key, value in result:
            self.listWidget.addItem("层 " + str(int(key) + 1))

    def changeLayer(self, item):
        """点击列表自动展示该层"""
        self.currIndex = int(item.text().split(' ')[1]) - 1
        if self.isInferSucceed:
            self.showImg(self.drawContours(self.npImage, self.inferData, self.currIndex))
        else:
            self.showImg(self.npImage[self.currIndex])

    """画布事件"""
    def myWheel(self,currIndex):
        """
        滚轮切换上下层
        """
        try:
            if self.maxCurrIndex != self.minCurrIndex:
                if currIndex > self.maxCurrIndex-1 :
                    self.currIndex = self.maxCurrIndex-1
                elif currIndex < 0:
                    self.currIndex = 0
                else:
                    self.currIndex = currIndex
                if self.isInferSucceed:
                    # self.npImage = self.drawContours(self.npImage, self.inferData)
                    self.showImg(self.drawContours(self.npImage, self.inferData, self.currIndex))
                else:
                    self.showImg(self.npImage[self.currIndex])
        except Exception as e:
            print(e)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    main = MainWindow()
    main.show()
    sys.exit(app.exec_())


# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# File:

import torch
import numpy as np
from typing import Dict
from torch import nn
# from typing import Any, Iterator, List, Union
# import pycocotools.mask as mask_utils
# from torch.nn import functional as F
from detectron2.layers import Conv2d, ShapeSpec, get_norm
# import fvcore.nn.weight_init as weight_init
from detectron2.modeling import ROI_HEADS_REGISTRY, StandardROIHeads
from detectron2.modeling.poolers import ROIPooler, MultiROIPooler

import fvcore.nn.weight_init as weight_init
from torch.nn import functional as F
# Is this function good to expose as an API?
from detectron2.modeling.roi_heads.roi_heads import select_foreground_proposals, select_proposals_with_visible_keypoints
from detectron2.modeling.roi_heads.keypoint_head import keypoint_rcnn_inference,keypoint_rcnn_loss
from .densepose_head import (
    build_densepose_data_filter,
    build_densepose_head,
    build_densepose_losses,
    build_densepose_predictor,
    densepose_inference,
    DensePoseInterLosses,
    DensePoseDataFilter,
    dp_keypoint_rcnn_loss,
    DensePoseKeypointsPredictor,
)
from .semantic_mask_head import (
    build_semantic_mask_data_filter,
    MIDPredictor,
    DpSemSegFPNHead
)

@ROI_HEADS_REGISTRY.register()
class DensePoseROIHeads(StandardROIHeads):
    """
    A Standard ROIHeads which contains an addition of DensePose head.
    """

    def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape)
        self._init_densepose_head(cfg, input_shape)
        self._init_dpsemseg_head(cfg)
        self._init_dp_keypoint_head(cfg,input_shape)

    def _init_dpsemseg_head(self, cfg):
        self.dp_semseg_on = cfg.MODEL.ROI_DENSEPOSE_HEAD.SEMSEG_ON

        if not self.dp_semseg_on:
            return
        input_shape = {
            name: ShapeSpec(
                channels=self.feature_channels[name], stride=self.feature_strides[name]
            )
            for name in self.in_features
        }
        self.common_stride = cfg.MODEL.SEM_SEG_HEAD.COMMON_STRIDE
        if self.use_mid:
            self.dp_semseg_head = MIDPredictor(cfg)
        else:
            self.dp_semseg_head = MIDPredictor(cfg)
        #     self.dp_semseg_head = DpSemSegFPNHead(cfg, input_shape)
        self.sem_mask_data_filter = build_semantic_mask_data_filter(cfg)
        # pooler_scales = [1.0 / self.feature_strides[self.in_features[0]]]
        # self.sematic_mask_pooler = MultiROIPooler(
        #         output_size=[[28,28]],
        #         scales=pooler_scales,
        #         sampling_ratio=cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO,
        #         pooler_type=cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE,
        #         canonical_level=2,
        #     )
        # self.sematic_mask_pooler = ROIPooler(
        #     output_size=cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION,
        #     scales=tuple(1.0 / self.feature_strides[k] for k in self.in_features),
        #     sampling_ratio=cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO,
        #     pooler_type=cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE,
        # )

    def _init_dp_keypoint_head(self, cfg,input_shape):
        # fmt: off

        self.dp_keypoint_on                         = cfg.MODEL.ROI_DENSEPOSE_HEAD.KPT_ON
        if not self.dp_keypoint_on:
            return
        self.normalize_loss_by_visible_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.NORMALIZE_LOSS_BY_VISIBLE_KEYPOINTS  # noqa
        self.keypoint_loss_weight                = cfg.MODEL.ROI_KEYPOINT_HEAD.LOSS_WEIGHT
        dp_pooler_resolution = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION
        # dp_multi_pooler_res        = ((28,28),(14,14),(14,14),(7,7))
        dp_pooler_scales = tuple(1.0 / self.feature_strides[k] for k in self.in_features)
        dp_pooler_sampling_ratio = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO
        dp_pooler_type = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE
        in_channels = [self.feature_channels[f] for f in self.in_features][0]
        if not self.densepose_on:
            self.use_mid = cfg.MODEL.ROI_DENSEPOSE_HEAD.MID_ON

            if cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME == 'DensePoseAMAHead':
                self.densepose_pooler = MultiROIPooler(
                    output_size=[[28, 28], [14, 14], [14, 14], [7, 7]],
                    # output_size=[[28, 28], [28, 28], [28, 28], [28, 28]],
                    scales=dp_pooler_scales,
                    sampling_ratio=dp_pooler_sampling_ratio,
                    pooler_type=dp_pooler_type,
                )
            else:
                self.densepose_pooler = ROIPooler(
                    output_size=dp_pooler_resolution,
                    scales=dp_pooler_scales,
                    sampling_ratio=dp_pooler_sampling_ratio,
                    pooler_type=dp_pooler_type,
                )
            self.densepose_head = build_densepose_head(cfg, in_channels)
            # print(self.densepose_head)
            self.keypoint_predictor = DensePoseKeypointsPredictor(cfg, self.densepose_head.n_out_channels)

    def _init_densepose_head(self, cfg, input_shape):
        # fmt: off
        self.cfg                   = cfg
        self.densepose_on          = cfg.MODEL.DENSEPOSE_ON and cfg.MODEL.ROI_DENSEPOSE_HEAD.RCNN_HEAD_ON
        if not self.densepose_on:
            return
        self.densepose_data_filter = build_densepose_data_filter(cfg)
        dp_pooler_resolution       = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION
        # dp_multi_pooler_res        = ((28,28),(14,14),(14,14),(7,7))
        dp_pooler_scales           = tuple(1.0 / self.feature_strides[k] for k in self.in_features)
        dp_pooler_sampling_ratio   = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO
        dp_pooler_type             = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE
        self.use_mid = cfg.MODEL.ROI_DENSEPOSE_HEAD.MID_ON

        self.inter_super_on = False
        self.inter_weight          = cfg.MODEL.ROI_DENSEPOSE_HEAD.INTER_WEIGHTS
        if cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME == 'DensePosePIDHead':
            self.inter_super_on = True

        # fmt: on
        in_channels = [self.feature_channels[f] for f in self.in_features][0]
        if cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME == 'DensePoseAMAHead':
            self.densepose_pooler = MultiROIPooler(
                # output_size=[[28,28],[14,14],[14,14],[7,7]],
                output_size=[[32, 32], [16, 16], [16, 16], [8, 8]],
                scales=dp_pooler_scales,
                sampling_ratio=dp_pooler_sampling_ratio,
                pooler_type=dp_pooler_type,
            )
        else:
            self.densepose_pooler = ROIPooler(
                output_size=dp_pooler_resolution,
                scales=dp_pooler_scales,
                sampling_ratio=dp_pooler_sampling_ratio,
                pooler_type=dp_pooler_type,
            )

        self.densepose_head = build_densepose_head(cfg, in_channels)
        self.densepose_predictor = build_densepose_predictor(
            cfg, self.densepose_head.n_out_channels
        )
        self.densepose_losses = build_densepose_losses(cfg)
        self.densepose_inter_losses = DensePoseInterLosses(cfg)

    def _forward_semsegs(self, features, instances, extra):

        if not self.dp_semseg_on:
            return
        if self.training:
            im_h, im_w = int(features[0].size(2)* self.common_stride), \
                         int(features[0].size(3)* self.common_stride)
            # proposals_with_targets, _ = select_foreground_proposals(instances, self.num_classes)
            # gt_sem_seg = self.sem_mask_data_filter(proposals_with_targets, im_h, im_w)
            gt_sem_seg = self.sem_mask_data_filter(extra, im_h, im_w)
            sem_seg_results, sem_seg_losses, latent_features = self.dp_semseg_head(features, gt_sem_seg)
            # sem_seg_results = gt_sem_seg.float().unsqueeze(1)
        else:
            gt_sem_seg = None
            sem_seg_results, sem_seg_losses, latent_features = self.dp_semseg_head(features, gt_sem_seg)

        return sem_seg_results, sem_seg_losses, latent_features

    def _forward_dp_keypoint(self, keypoint_logits, instances):

        if not self.dp_keypoint_on:
            return {} if self.training else instances
        num_images = len(instances)
        if self.training:
            # The loss is defined on positive proposals with at >=1 visible keypoints.
            normalizer = (
                num_images
                * self.batch_size_per_image
                * self.positive_sample_fraction
                * keypoint_logits.shape[1]
            )
            loss = dp_keypoint_rcnn_loss(
                keypoint_logits,
                instances,
                normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer,
            )
            return {"loss_keypoint": loss * self.keypoint_loss_weight}
        else:
            keypoint_rcnn_inference(keypoint_logits, instances)
            return instances

    def forward_dp_keypoint(self, features, instances):

        if not self.dp_keypoint_on:
            return {} if self.training else instances
        num_images = len(instances)
        if self.training:
            proposals, _ = select_foreground_proposals(instances, self.num_classes)
            proposals = select_proposals_with_visible_keypoints(proposals)
            # proposals = self.keypoint_data_filter(proposals)
            proposal_boxes = [x.proposal_boxes for x in proposals]
            # print('after:',len(proposal_boxes))
            # if len(proposal_boxes) == 0:
            #     return {"loss_keypoint": 0.}
            # if len(proposal_boxes) > 0:
            if self.use_mid:
                features = [self.mid_decoder(features)]
            keypoint_features = self.densepose_pooler(features, proposal_boxes)
            keypoint_output = self.densepose_head(keypoint_features)
            keypoint_logits = self.keypoint_predictor(keypoint_output)
            # The loss is defined on positive proposals with at >=1 visible keypoints.
            normalizer = (
                num_images
                * self.batch_size_per_image
                * self.positive_sample_fraction
                * keypoint_logits.shape[1]
            )
            # loss = dp_keypoint_rcnn_loss(
            #     keypoint_logits,
            #     instances,
            #     normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer,
            # )
            loss = keypoint_rcnn_loss(
                keypoint_logits,
                proposals,
                normalizer=None if self.normalize_loss_by_visible_keypoints else normalizer,
            )
            return {"loss_keypoint": loss * self.keypoint_loss_weight}
        else:
            if self.use_mid:
                features = [self.mid_decoder(features)]
            pred_boxes = [x.pred_boxes for x in instances]
            keypoint_features = self.densepose_pooler(features, pred_boxes)
            keypoint_output = self.densepose_head(keypoint_features)
            keypoint_logits = self.keypoint_predictor(keypoint_output)
            keypoint_rcnn_inference(keypoint_logits, instances)
            return instances

    def _forward_densepose(self, features, instances, extra=None):
        """
        Forward logic of the densepose prediction branch.

        Args:
            features (list[Tensor]): #level input features for densepose prediction
            instances (list[Instances]): the per-image instances to train/predict densepose.
                In training, they can be the proposals.
                In inference, they can be the predicted boxes.

        Returns:
            In training, a dict of losses.
            In inference, update `instances` with new fields "densepose" and return it.
        """
        if not self.densepose_on:
            return {} if self.training else instances
        bbox_locs_params = self.box_predictor.bbox_pred.weight
        bbox_cls_params = self.box_predictor.cls_score.weight
        with torch.no_grad():
            bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0)
        if self.training:
            proposals, _ = select_foreground_proposals(instances, self.num_classes)
            proposals_dp = self.densepose_data_filter(proposals)
            if len(proposals_dp) > 0:
                proposal_boxes = [x.proposal_boxes for x in proposals_dp]
                if self.use_mid:
                    features = [self.mid_decoder(features)]
                    if self.dp_semseg_on:
                        segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra)
                else:
                    if self.dp_semseg_on:
                        segm_res, seg_losses, latent_features = self._forward_semsegs([features[-1]], instances, extra)
                features_dp = self.densepose_pooler(features, proposal_boxes)
                # print('roi pooler:',features_dp.size())
                if self.inter_super_on:
                    densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp)
                    densepose_outputs,  densepose_inter_outputs= self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs)
                    densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=True)
                    # inter_loss_dict = self.densepose_losses(proposals_dp, densepose_inter_outputs, 'inter_')
                    inter_loss_dict = self.densepose_inter_losses(proposals_dp, densepose_inter_outputs, 'inter_')
                    for _, k in enumerate(inter_loss_dict.keys()):
                        if k == 'inter_loss_densepose_M':
                            densepose_loss_dict[k] = inter_loss_dict[k] * 1.
                        else:
                            densepose_loss_dict[k] = inter_loss_dict[k] * 0.
                else:
                    densepose_head_outputs = self.densepose_head(features_dp)
                    # densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs, bbox_params)
                    densepose_outputs, dp_outputs_from_kpt, dp_outputs_from_bbox = self.densepose_predictor(densepose_head_outputs, bbox_params)
                    if self.dp_keypoint_on:
                        keypoints_output = densepose_outputs[-1]
                        densepose_outputs = densepose_outputs[:-1]
                    densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs)
                    if self.dp_keypoint_on:
                        kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp)
                        for _, k in enumerate(kpt_loss_dict.keys()):
                            densepose_loss_dict[k] = kpt_loss_dict[k]
                    if self.dp_semseg_on:
                        for _, k in enumerate(seg_losses.keys()):
                            densepose_loss_dict[k] = seg_losses[k]
                    if dp_outputs_from_kpt is not None:
                        dp_kpt_loss_dict = self.densepose_losses(proposals_dp, dp_outputs_from_kpt)
                        densepose_loss_dict['loss_densepose_I_from_kpt'] = dp_kpt_loss_dict['loss_densepose_I']
                    if dp_outputs_from_bbox is not None:
                        dp_box_loss_dict = self.densepose_losses(proposals_dp, dp_outputs_from_bbox)
                        densepose_loss_dict['loss_densepose_I_from_box'] = dp_box_loss_dict['loss_densepose_I']
                return densepose_loss_dict
        else:
            pred_boxes = [x.pred_boxes for x in instances]
            if self.use_mid:
                features = [self.mid_decoder(features)]
                # if self.dp_semseg_on:
                #     segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra)
            features_dp = self.densepose_pooler(features, pred_boxes)
            if len(features_dp) > 0:
                if self.inter_super_on:
                    densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp)
                    densepose_outputs, densepose_inter_outputs = self.densepose_predictor(densepose_head_outputs,
                                                                                densepose_inter_head_outputs)
                else:
                    densepose_head_outputs = self.densepose_head(features_dp)
                    densepose_outputs, _, _ = self.densepose_predictor(densepose_head_outputs,bbox_params)
                    if self.dp_keypoint_on:
                        keypoints_output = densepose_outputs[-1]
                        densepose_outputs = densepose_outputs[:-1]
                        instances = self._forward_dp_keypoint(keypoints_output, instances)
            else:
                # If no detection occured instances
                # set densepose_outputs to empty tensors
                empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp.device)
                densepose_outputs = tuple([empty_tensor] * 5)

            densepose_inference(densepose_outputs, instances)
            return instances

    def _forward_ama_densepose(self, features, instances, extra=None):
        """
        Forward logic of the densepose prediction branch.

        Args:
            features (list[Tensor]): #level input features for densepose prediction
            instances (list[Instances]): the per-image instances to train/predict densepose.
                In training, they can be the proposals.
                In inference, they can be the predicted boxes.

        Returns:
            In training, a dict of losses.
            In inference, update `instances` with new fields "densepose" and return it.
        """
        if not self.densepose_on:
            return {} if self.training else instances
        latent_features = None
        if self.dp_semseg_on:
            segm_res, seg_losses, latent_features = self._forward_semsegs(features, instances, extra)
            # features = latent_features

        if self.training:
            proposals, _ = select_foreground_proposals(instances, self.num_classes)
            proposals_dp = self.densepose_data_filter(proposals)
            proposal_boxes = [x.proposal_boxes for x in proposals_dp]
            features_dp = self.densepose_pooler(features, proposal_boxes)

            if len(proposals_dp) > 0:
                densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features)
                bbox_locs_params = self.box_predictor.bbox_pred.weight
                bbox_cls_params  = self.box_predictor.cls_score.weight
                bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0)
                densepose_outputs, inter_output_1, inter_output_2 = self.densepose_predictor(densepose_head_outputs,
                                                                                densepose_inter_head_outputs, bbox_params)
                if self.dp_keypoint_on:
                    keypoints_output = densepose_outputs[-1]
                    densepose_outputs = densepose_outputs[:-1]
                densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs, cls_emb_loss_on=self.cfg.MODEL.ROI_DENSEPOSE_HEAD.IA_LOSS)

                inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_1, prefix='inter_1_')

                for _, k in enumerate(inter_loss_dict.keys()):
                    if 'loss_densepose_I' in k:
                        densepose_loss_dict[k] = inter_loss_dict[k]
                    else:
                        densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight
                inter_loss_dict = self.densepose_losses(proposals_dp, inter_output_2, prefix='inter_2_')
                for _, k in enumerate(inter_loss_dict.keys()):
                    if 'loss_densepose_I' in k:
                        densepose_loss_dict[k] = inter_loss_dict[k]
                    else:
                        densepose_loss_dict[k] = inter_loss_dict[k] * self.inter_weight
                if self.dp_semseg_on:
                    for _, k in enumerate(seg_losses.keys()):
                        densepose_loss_dict[k] = seg_losses[k]
                if self.dp_keypoint_on:

                    kpt_loss_dict = self._forward_dp_keypoint(keypoints_output, proposals_dp)
                    for _, k in enumerate(kpt_loss_dict.keys()):
                        densepose_loss_dict[k] = kpt_loss_dict[k]
                # box_features = self.box_pooler(features, proposal_boxes)
                # box_features = self.box_head(box_features)
                return densepose_loss_dict
        else:
            pred_boxes = [x.pred_boxes for x in instances]
            features_dp = self.densepose_pooler(features, pred_boxes)
            # if len(features_dp[0]) == 0:
            #     print('no proposal:', len(pred_boxes),pred_boxes)
            if len(features_dp[0]) > 0 and len(pred_boxes) > 0:

                densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp, latent_features)
                # densepose_head_outputs, densepose_inter_head_outputs = self.densepose_head(features_dp)
                bbox_locs_params = self.box_predictor.bbox_pred.weight
                bbox_cls_params = self.box_predictor.cls_score.weight
                bbox_params = torch.cat([bbox_cls_params, bbox_locs_params], dim=0)
                densepose_outputs, _, _ = self.densepose_predictor(densepose_head_outputs, densepose_inter_head_outputs,
                                                                   bbox_params)
                if self.dp_keypoint_on:
                    keypoints_output = densepose_outputs[-1]
                    densepose_outputs = densepose_outputs[:-1]
                    instances=self._forward_dp_keypoint(keypoints_output, instances)
            else:
                # If no detection occured instances
                # set densepose_outputs to empty tensors
                empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp[0].device)
                densepose_outputs = tuple([empty_tensor] * 5)

            densepose_inference(densepose_outputs, instances)
            # if self.dp_keypoint_on:
            #     instances = self._forward_dp_keypoint(keypoints_output, instances)
            return instances

    def forward(self, images, features, proposals, targets=None, extra=None):
        features_list = [features[f] for f in self.in_features]

        instances, losses = super().forward(images, features, proposals, targets)

        del targets, images

        if self.training:
            if self.cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME == 'DensePoseAMAHead':
                losses.update(self._forward_ama_densepose(features_list, instances, extra))

            else:
                losses.update(self._forward_densepose(features_list, instances, extra))
                if self.dp_keypoint_on and not self.densepose_on:
                    losses.update(self.forward_dp_keypoint(features_list, instances))
        else:
            if self.cfg.MODEL.ROI_DENSEPOSE_HEAD.NAME == 'DensePoseAMAHead':
                instances = self._forward_ama_densepose(features_list, instances)
            else:
                instances = self._forward_densepose(features_list, instances)
                if self.dp_keypoint_on and not self.densepose_on:
                    instances = self.forward_dp_keypoint(features_list, instances)
        return instances, losses


# -*- encoding: utf-8 -*-
from typing import Optional, Dict, List, Tuple, Union, Iterable
from typing_extensions import Literal

from ConfigSpace.configuration_space import Configuration
import dask.distributed
import joblib
import numpy as np
import pandas as pd
from scipy.sparse import spmatrix
from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
from sklearn.utils.multiclass import type_of_target
from smac.runhistory.runhistory import RunInfo, RunValue

from autosklearn.data.validation import (
    convert_if_sparse,
    SUPPORTED_FEAT_TYPES,
    SUPPORTED_TARGET_TYPES,
)
from autosklearn.pipeline.base import BasePipeline
from autosklearn.automl import AutoMLClassifier, AutoMLRegressor, AutoML
from autosklearn.metrics import Scorer

class AutoSklearnEstimator(BaseEstimator):

    def __init__(
        self,
        time_left_for_this_task=3600,
        per_run_time_limit=None,
        initial_configurations_via_metalearning=25,
        ensemble_size: int = 50,
        ensemble_nbest=50,
        max_models_on_disc=50,
        seed=1,
        memory_limit=3072,
        include=None,
        exclude=None,
        resampling_strategy='holdout',
        resampling_strategy_arguments=None,
        tmp_folder=None,
        delete_tmp_folder_after_terminate=True,
        n_jobs: Optional[int] = None,
        dask_client: Optional[dask.distributed.Client] = None,
        disable_evaluator_output=False,
        get_smac_object_callback=None,
        smac_scenario_args=None,
        logging_config=None,
        metadata_directory=None,
        metric=None,
        scoring_functions: Optional[List[Scorer]] = None,
        load_models: bool = True,
        get_trials_callback=None
    ):
        """
        Parameters
        ----------
        time_left_for_this_task : int, optional (default=3600)
            Time limit in seconds for the search of appropriate
            models. By increasing this value, *auto-sklearn* has a higher
            chance of finding better models.

        per_run_time_limit : int, optional (default=1/10 of time_left_for_this_task)
            Time limit for a single call to the machine learning model.
            Model fitting will be terminated if the machine learning
            algorithm runs over the time limit. Set this value high enough so
            that typical machine learning algorithms can be fit on the
            training data.

        initial_configurations_via_metalearning : int, optional (default=25)
            Initialize the hyperparameter optimization algorithm with this
            many configurations which worked well on previously seen
            datasets. Disable if the hyperparameter optimization algorithm
            should start from scratch.

        ensemble_size : int, optional (default=50)
            Number of models added to the ensemble built by *Ensemble
            selection from libraries of models*. Models are drawn with
            replacement.

        ensemble_nbest : int, optional (default=50)
            Only consider the ``ensemble_nbest`` models when building an
            ensemble.

        max_models_on_disc: int, optional (default=50),
            Defines the maximum number of models that are kept in the disc.
            The additional number of models are permanently deleted. Due to the
            nature of this variable, it sets the upper limit on how many models
            can be used for an ensemble.
            It must be an integer greater or equal than 1.
            If set to None, all models are kept on the disc.

        seed : int, optional (default=1)
            Used to seed SMAC. Will determine the output file names.

        memory_limit : int, optional (3072)
            Memory limit in MB for the machine learning algorithm.
            `auto-sklearn` will stop fitting the machine learning algorithm if
            it tries to allocate more than `memory_limit` MB.
            If None is provided, no memory limit is set.
            In case of multi-processing, `memory_limit` will be per job.
            This memory limit also applies to the ensemble creation process.

        include : dict, optional (None)
            If None, all possible algorithms are used. Otherwise specifies
            set of algorithms for each added component is used. Include and
            exclude are incompatible if used together on the same component

        exclude : dict, optional (None)
            If None, all possible algorithms are used. Otherwise specifies
            set of algorithms for each added component is not used.
            Incompatible with include. Include and exclude are incompatible
            if used together on the same component

        resampling_strategy : string or object, optional ('holdout')
            how to to handle overfitting, might need 'resampling_strategy_arguments'

            * 'holdout': 67:33 (train:test) split
            * 'holdout-iterative-fit':  67:33 (train:test) split, calls iterative
              fit where possible
            * 'cv': crossvalidation, requires 'folds'
            * 'cv-iterative-fit': crossvalidation, calls iterative fit where possible
            * 'partial-cv': crossvalidation with intensification, requires
              'folds'
            * BaseCrossValidator object: any BaseCrossValidator class found
                                        in scikit-learn model_selection module
            * _RepeatedSplits object: any _RepeatedSplits class found
                                      in scikit-learn model_selection module
            * BaseShuffleSplit object: any BaseShuffleSplit class found
                                      in scikit-learn model_selection module

        resampling_strategy_arguments : dict, optional if 'holdout' (train_size default=0.67)
            Additional arguments for resampling_strategy:

            * ``train_size`` should be between 0.0 and 1.0 and represent the
              proportion of the dataset to include in the train split.
            * ``shuffle`` determines whether the data is shuffled prior to
              splitting it into train and validation.

            Available arguments:

            * 'holdout': {'train_size': float}
            * 'holdout-iterative-fit':  {'train_size': float}
            * 'cv': {'folds': int}
            * 'cv-iterative-fit': {'folds': int}
            * 'partial-cv': {'folds': int, 'shuffle': bool}
            * BaseCrossValidator or _RepeatedSplits or BaseShuffleSplit object: all arguments
                required by chosen class as specified in scikit-learn documentation.
                If arguments are not provided, scikit-learn defaults are used.
                If no defaults are available, an exception is raised.
                Refer to the 'n_splits' argument as 'folds'.

        tmp_folder : string, optional (None)
            folder to store configuration output and log files, if ``None``
            automatically use ``/tmp/autosklearn_tmp_$pid_$random_number``

        delete_tmp_folder_after_terminate: bool, optional (True)
            remove tmp_folder, when finished. If tmp_folder is None
            tmp_dir will always be deleted

        n_jobs : int, optional, experimental
            The number of jobs to run in parallel for ``fit()``. ``-1`` means
            using all processors. By default, Auto-sklearn uses a single core
            for fitting the machine learning model and a single core for fitting
            an ensemble. Ensemble building is not affected by ``n_jobs`` but
            can be controlled by the number of models in the ensemble. In
            contrast to most scikit-learn models, ``n_jobs`` given in the
            constructor is not applied to the ``predict()`` method. If
            ``dask_client`` is None, a new dask client is created.

        dask_client : dask.distributed.Client, optional
            User-created dask client, can be used to start a dask cluster and then
            attach auto-sklearn to it.

        disable_evaluator_output: bool or list, optional (False)
            If True, disable model and prediction output. Cannot be used
            together with ensemble building. ``predict()`` cannot be used when
            setting this True. Can also be used as a list to pass more
            fine-grained information on what to save. Allowed elements in the
            list are:

            * ``'y_optimization'`` : do not save the predictions for the
              optimization/validation set, which would later on be used to build
              an ensemble.
            * ``'model'`` : do not save any model files

        smac_scenario_args : dict, optional (None)
            Additional arguments inserted into the scenario of SMAC. See the
            `SMAC documentation <https://automl.github.io/SMAC3/master/options.html?highlight=scenario#scenario>`_
            for a list of available arguments.

        get_smac_object_callback : callable
            Callback function to create an object of class
            `smac.optimizer.smbo.SMBO <https://automl.github.io/SMAC3/master/apidoc/smac.optimizer.smbo.html>`_.
            The function must accept the arguments ``scenario_dict``,
            ``instances``, ``num_params``, ``runhistory``, ``seed`` and ``ta``.
            This is an advanced feature. Use only if you are familiar with
            `SMAC <https://automl.github.io/SMAC3/master/index.html>`_.

        logging_config : dict, optional (None)
            dictionary object specifying the logger configuration. If None,
            the default logging.yaml file is used, which can be found in
            the directory ``util/logging.yaml`` relative to the installation.

        metadata_directory : str, optional (None)
            path to the metadata directory. If None, the default directory
            (autosklearn.metalearning.files) is used.

        metric : Scorer, optional (None)
            An instance of :class:`autosklearn.metrics.Scorer` as created by
            :meth:`autosklearn.metrics.make_scorer`. These are the `Built-in
            Metrics`_.
            If None is provided, a default metric is selected depending on the task.

        scoring_functions : List[Scorer], optional (None)
            List of scorers which will be calculated for each pipeline and results will be
            available via ``cv_results``

        load_models : bool, optional (True)
            Whether to load the models after fitting Auto-sklearn.

        get_trials_callback: callable
            Callback function to create an object of subclass defined in module
            `smac.callbacks <https://automl.github.io/SMAC3/master/apidoc/smac.callbacks.html>`_.
            This is an advanced feature. Use only if you are familiar with
            `SMAC <https://automl.github.io/SMAC3/master/index.html>`_.

        Attributes
        ----------

        cv_results\_ : dict of numpy (masked) ndarrays
            A dict with keys as column headers and values as columns, that can be
            imported into a pandas ``DataFrame``.

            Not all keys returned by scikit-learn are supported yet.

        performance_over_time\_ : pandas.core.frame.DataFrame
            A ``DataFrame`` containing the models performance over time data. Can be
            used for plotting directly. Please refer to the example
            :ref:`Train and Test Inputs <sphx_glr_examples_40_advanced_example_pandas_train_test.py>`.

        """  # noqa (links are too long)
        # Raise error if the given total time budget is less than 30 seconds.
        if time_left_for_this_task < 30:
            raise ValueError("Time left for this task must be at least "
                             "30 seconds. ")
        self.time_left_for_this_task = time_left_for_this_task
        self.per_run_time_limit = per_run_time_limit
        self.initial_configurations_via_metalearning = initial_configurations_via_metalearning
        self.ensemble_size = ensemble_size
        self.ensemble_nbest = ensemble_nbest
        self.max_models_on_disc = max_models_on_disc
        self.seed = seed
        self.memory_limit = memory_limit
        self.include = include
        self.exclude = exclude
        self.resampling_strategy = resampling_strategy
        self.resampling_strategy_arguments = resampling_strategy_arguments
        self.tmp_folder = tmp_folder
        self.delete_tmp_folder_after_terminate = delete_tmp_folder_after_terminate
        self.n_jobs = n_jobs
        self.dask_client = dask_client
        self.disable_evaluator_output = disable_evaluator_output
        self.get_smac_object_callback = get_smac_object_callback
        self.smac_scenario_args = smac_scenario_args
        self.logging_config = logging_config
        self.metadata_directory = metadata_directory
        self.metric = metric
        self.scoring_functions = scoring_functions
        self.load_models = load_models
        self.get_trials_callback = get_trials_callback

        self.automl_ = None  # type: Optional[AutoML]

        # Handle the number of jobs and the time for them
        self._n_jobs = None
        if self.n_jobs is None or self.n_jobs == 1:
            self._n_jobs = 1
        elif self.n_jobs == -1:
            self._n_jobs = joblib.cpu_count()
        else:
            self._n_jobs = self.n_jobs

        super().__init__()

    def __getstate__(self):
        # Cannot serialize a client!
        self.dask_client = None
        return self.__dict__

    def build_automl(self):

        automl = self._get_automl_class()(
            temporary_directory=self.tmp_folder,
            delete_tmp_folder_after_terminate=self.delete_tmp_folder_after_terminate,
            time_left_for_this_task=self.time_left_for_this_task,
            per_run_time_limit=self.per_run_time_limit,
            initial_configurations_via_metalearning=self.initial_configurations_via_metalearning,
            ensemble_size=self.ensemble_size,
            ensemble_nbest=self.ensemble_nbest,
            max_models_on_disc=self.max_models_on_disc,
            seed=self.seed,
            memory_limit=self.memory_limit,
            include=self.include,
            exclude=self.exclude,
            resampling_strategy=self.resampling_strategy,
            resampling_strategy_arguments=self.resampling_strategy_arguments,
            n_jobs=self._n_jobs,
            dask_client=self.dask_client,
            get_smac_object_callback=self.get_smac_object_callback,
            disable_evaluator_output=self.disable_evaluator_output,
            smac_scenario_args=self.smac_scenario_args,
            logging_config=self.logging_config,
            metadata_directory=self.metadata_directory,
            metric=self.metric,
            scoring_functions=self.scoring_functions,
            get_trials_callback=self.get_trials_callback
        )

        return automl

    def fit(self, **kwargs):

        # Automatically set the cutoff time per task
        if self.per_run_time_limit is None:
            self.per_run_time_limit = self._n_jobs * self.time_left_for_this_task // 10

        if self.automl_ is None:
            self.automl_ = self.build_automl()
        self.automl_.fit(load_models=self.load_models, **kwargs)

        return self

    def fit_pipeline(
        self,
        X: SUPPORTED_FEAT_TYPES,
        y: Union[SUPPORTED_TARGET_TYPES, spmatrix],
        config: Union[Configuration,  Dict[str, Union[str, float, int]]],
        dataset_name: Optional[str] = None,
        X_test: Optional[SUPPORTED_FEAT_TYPES] = None,
        y_test: Optional[Union[SUPPORTED_TARGET_TYPES, spmatrix]] = None,
        feat_type: Optional[List[str]] = None,
        *args,
        **kwargs: Dict,
    ) -> Tuple[Optional[BasePipeline], RunInfo, RunValue]:
        """ Fits and individual pipeline configuration and returns
        the result to the user.

        The Estimator constraints are honored, for example the resampling
        strategy, or memory constraints, unless directly provided to the method.
        By default, this method supports the same signature as fit(), and any extra
        arguments are redirected to the TAE evaluation function, which allows for
        further customization while building a pipeline.

        Any additional argument provided is directly passed to the worker exercising the run.

        Parameters
        ----------
        X: array-like, shape = (n_samples, n_features)
            The features used for training
        y: array-like
            The labels used for training
        X_test: Optionalarray-like, shape = (n_samples, n_features)
            If provided, the testing performance will be tracked on this features.
        y_test: array-like
            If provided, the testing performance will be tracked on this labels
        config: Union[Configuration,  Dict[str, Union[str, float, int]]]
            A configuration object used to define the pipeline steps.
            If a dictionary is passed, a configuration is created based on this dictionary.
        dataset_name: Optional[str]
            Name that will be used to tag the Auto-Sklearn run and identify the
            Auto-Sklearn run
        feat_type : list, optional (default=None)
            List of str of `len(X.shape[1])` describing the attribute type.
            Possible types are `Categorical` and `Numerical`. `Categorical`
            attributes will be automatically One-Hot encoded. The values
            used for a categorical attribute must be integers, obtained for
            example by `sklearn.preprocessing.LabelEncoder
            <https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.LabelEncoder.html>`_.

        Returns
        -------
        pipeline: Optional[BasePipeline]
            The fitted pipeline. In case of failure while fitting the pipeline,
            a None is returned.
        run_info: RunInFo
            A named tuple that contains the configuration launched
        run_value: RunValue
            A named tuple that contains the result of the run
        """
        if self.automl_ is None:
            self.automl_ = self.build_automl()
        return self.automl_.fit_pipeline(X=X, y=y,
                                         dataset_name=dataset_name,
                                         config=config,
                                         feat_type=feat_type,
                                         X_test=X_test, y_test=y_test,
                                         *args, **kwargs)

    def fit_ensemble(self, y, task=None, precision=32,
                     dataset_name=None, ensemble_nbest=None,
                     ensemble_size=None):
        """Fit an ensemble to models trained during an optimization process.

        All parameters are ``None`` by default. If no other value is given,
        the default values which were set in a call to ``fit()`` are used.

        Calling this function is only necessary if ``ensemble_size==0``, for
        example when executing *auto-sklearn* in parallel.

        Parameters
        ----------
        y : array-like
            Target values.

        task : int
            A constant from the module ``autosklearn.constants``. Determines
            the task type (binary classification, multiclass classification,
            multilabel classification or regression).

        precision : str
            Numeric precision used when loading ensemble data. Can be either
            ``'16'``, ``'32'`` or ``'64'``.

        dataset_name : str
            Name of the current data set.

        ensemble_nbest : int
            Determines how many models should be considered from the ensemble
            building. This is inspired by a concept called library pruning
            introduced in `Getting Most out of Ensemble Selection`.

        ensemble_size : int
            Size of the ensemble built by `Ensemble Selection`.

        Returns
        -------
        self

        """
        if self.automl_ is None:
            # Build a dummy automl object to call fit_ensemble
            # The ensemble size is honored in the .automl_.fit_ensemble
            # call
            self.automl_ = self.build_automl()
        self.automl_.fit_ensemble(
            y=y,
            task=task,
            precision=precision,
            dataset_name=dataset_name,
            ensemble_nbest=ensemble_nbest,
            ensemble_size=ensemble_size,
        )
        return self

    def refit(self, X, y):
        """Refit all models found with fit to new data.

        Necessary when using cross-validation. During training, auto-sklearn
        fits each model k times on the dataset, but does not keep any trained
        model and can therefore not be used to predict for new data points.
        This methods fits all models found during a call to fit on the data
        given. This method may also be used together with holdout to avoid
        only using 66% of the training data to fit the final model.

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The targets.

        Returns
        -------

        self

        """
        self.automl_.refit(X, y)
        return self

    def predict(self, X, batch_size=None, n_jobs=1):
        return self.automl_.predict(X, batch_size=batch_size, n_jobs=n_jobs)

    def predict_proba(self, X, batch_size=None, n_jobs=1):
        return self.automl_.predict_proba(
             X, batch_size=batch_size, n_jobs=n_jobs)

    def score(self, X, y):
        return self.automl_.score(X, y)

    def show_models(self):
        """Return a representation of the final ensemble found by auto-sklearn.

        Returns
        -------
        str

        """
        return self.automl_.show_models()

    def get_models_with_weights(self):
        """Return a list of the final ensemble found by auto-sklearn.

        Returns
        -------
        [(weight_1, model_1), ..., (weight_n, model_n)]

        """
        return self.automl_.get_models_with_weights()

    @property
    def performance_over_time_(self):
        return self.automl_.performance_over_time_

    @property
    def cv_results_(self):
        return self.automl_.cv_results_

    @property
    def trajectory_(self):
        return self.automl_.trajectory_

    @property
    def fANOVA_input_(self):
        return self.automl_.fANOVA_input_

    def sprint_statistics(self):
        """Return the following statistics of the training result:

        - dataset name
        - metric used
        - best validation score
        - number of target algorithm runs
        - number of successful target algorithm runs
        - number of crashed target algorithm runs
        - number of target algorithm runs that exceeded the memory limit
        - number of target algorithm runs that exceeded the time limit

        Returns
        -------
        str
        """
        return self.automl_.sprint_statistics()

    def leaderboard(
        self,
        detailed: bool = False,
        ensemble_only: bool = True,
        top_k: Union[int, Literal['all']] = 'all',
        sort_by: str = 'cost',
        sort_order: Literal['auto', 'ascending', 'descending'] = 'auto',
        include: Optional[Union[str, Iterable[str]]] = None
    ) -> pd.DataFrame:
        """ Returns a pandas table of results for all evaluated models.

        Gives an overview of all models trained during the search process along
        with various statistics about their training.

        The availble statistics are:

        **Simple**:

        * ``"model_id"`` - The id given to a model by ``autosklearn``.
        * ``"rank"`` - The rank of the model based on it's ``"cost"``.
        * ``"ensemble_weight"`` - The weight given to the model in the ensemble.
        * ``"type"`` - The type of classifier/regressor used.
        * ``"cost"`` - The loss of the model on the validation set.
        * ``"duration"`` - Length of time the model was optimized for.

        **Detailed**:
        The detailed view includes all of the simple statistics along with the
        following.

        * ``"config_id"`` - The id used by SMAC for optimization.
        * ``"budget"`` - How much budget was allocated to this model.
        * ``"status"`` - The return status of training the model with SMAC.
        * ``"train_loss"`` - The loss of the model on the training set.
        * ``"balancing_strategy"`` - The balancing strategy used for data preprocessing.
        * ``"start_time"`` - Time the model began being optimized
        * ``"end_time"`` - Time the model ended being optimized
        * ``"data_preprocessors"`` - The preprocessors used on the data
        * ``"feature_preprocessors"`` - The preprocessors for features types

        Parameters
        ----------
        detailed: bool = False
            Whether to give detailed information or just a simple overview.

        ensemble_only: bool = True
            Whether to view only models included in the ensemble or all models
            trained.

        top_k: int or "all" = "all"
            How many models to display.

        sort_by: str = 'cost'
            What column to sort by. If that column is not present, the
            sorting defaults to the ``"model_id"`` index column.

        sort_order: "auto" or "ascending" or "descending" = "auto"
            Which sort order to apply to the ``sort_by`` column. If left
            as ``"auto"``, it will sort by a sensible default where "better" is
            on top, otherwise defaulting to the pandas default for
            `DataFrame.sort_values`_ if there is no obvious "better".

            .. _DataFrame.sort_values: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.sort_values.html

        include: Optional[str or Iterable[str]]
            Items to include, other items not specified will be excluded.
            The exception is the ``"model_id"`` index column which is always included.

            If left as ``None``, it will resort back to using the ``detailed``
            param to decide the columns to include.

        Returns
        -------
        pd.DataFrame
            A dataframe of statistics for the models, ordered by ``sort_by``.

        """  # noqa (links are too long)
        # TODO validate that `self` is fitted. This is required for
        #      self.ensemble_ to get the identifiers of models it will generate
        #      weights for.
        column_types = AutoSklearnEstimator._leaderboard_columns()

        # Validation of top_k
        if (
            not (isinstance(top_k, str) or isinstance(top_k, int))
            or (isinstance(top_k, str) and top_k != 'all')
            or (isinstance(top_k, int) and top_k <= 0)
        ):
            raise ValueError(f"top_k={top_k} must be a positive integer or pass"
                             " `top_k`='all' to view results for all models")

        # Validate columns to include
        if isinstance(include, str):
            include = [include]

        if include == ['model_id']:
            raise ValueError('Must provide more than just `model_id`')

        if include is not None:
            columns = [*include]

            # 'model_id' should always be present as it is the unique index
            # used for pandas
            if 'model_id' not in columns:
                columns.append('model_id')

            invalid_include_items = set(columns) - set(column_types['all'])
            if len(invalid_include_items) != 0:
                raise ValueError(f"Values {invalid_include_items} are not known"
                                 f" columns to include, must be contained in "
                                 f"{column_types['all']}")
        elif detailed:
            columns = column_types['all']
        else:
            columns = column_types['simple']

        # Validation of sorting
        if sort_by not in column_types['all']:
            raise ValueError(f"sort_by='{sort_by}' must be one of included "
                             f"columns {set(column_types['all'])}")

        valid_sort_orders = ['auto', 'ascending', 'descending']
        if not (isinstance(sort_order, str) and sort_order in valid_sort_orders):
            raise ValueError(f"`sort_order` = {sort_order} must be a str in "
                             f"{valid_sort_orders}")

        # To get all the models that were optmized, we collect what we can from
        # runhistory first.
        def has_key(rv, key):
            return rv.additional_info and key in rv.additional_info

        model_runs = {
            rval.additional_info['num_run']: {
                'model_id': rval.additional_info['num_run'],
                'seed': rkey.seed,
                'budget': rkey.budget,
                'duration': rval.time,
                'config_id': rkey.config_id,
                'start_time': rval.starttime,
                'end_time': rval.endtime,
                'status': str(rval.status),
                'cost': rval.cost,
                'train_loss': rval.additional_info['train_loss']
                if has_key(rval, 'train_loss') else None,
                'config_origin': rval.additional_info['configuration_origin']
                if has_key(rval, 'configuration_origin') else None
            }
            for rkey, rval in self.automl_.runhistory_.data.items()
            if has_key(rval, 'num_run')
        }

        # Next we get some info about the model itself
        model_class_strings = {
            AutoMLClassifier: 'classifier',
            AutoMLRegressor: 'regressor'
        }
        model_type = model_class_strings.get(self._get_automl_class(), None)
        if model_type is None:
            raise RuntimeError(f"Unknown `automl_class` {self._get_automl_class()}")

        # A dict mapping model ids to their configurations
        configurations = self.automl_.runhistory_.ids_config

        for model_id, run_info in model_runs.items():
            config_id = run_info['config_id']
            run_config = configurations[config_id]._values

            run_info.update({
                'balancing_strategy': run_config.get('balancing:strategy', None),
                'type': run_config[f'{model_type}:__choice__'],
                'data_preprocessors': [
                    value for key, value in run_config.items()
                    if 'data_preprocessing' in key and '__choice__' in key
                ],
                'feature_preprocessors': [
                    value for key, value in run_config.items()
                    if 'feature_preprocessor' in key and '__choice__' in key
                ]
            })

        # Get the models ensemble weight if it has one
        # TODO both implementing classes of AbstractEnsemble have a property
        #      `identifiers_` and `weights_`, might be good to put it as an
        #       abstract property
        # TODO `ensemble_.identifiers_` and `ensemble_.weights_` are loosely
        #      tied together by ordering, might be better to store as tuple
        for i, weight in enumerate(self.automl_.ensemble_.weights_):
            (_, model_id, _) = self.automl_.ensemble_.identifiers_[i]
            model_runs[model_id]['ensemble_weight'] = weight

        # Filter out non-ensemble members if needed, else fill in a default
        # value of 0 if it's missing
        if ensemble_only:
            model_runs = {
                model_id: info
                for model_id, info in model_runs.items()
                if ('ensemble_weight' in info and info['ensemble_weight'] > 0)
            }
        else:
            for model_id, info in model_runs.items():
                if 'ensemble_weight' not in info:
                    info['ensemble_weight'] = 0

        # `rank` relies on `cost` so we include `cost`
        # We drop it later if it's not requested
        if 'rank' in columns and 'cost' not in columns:
            columns = [*columns, 'cost']

        # Finally, convert into a tabular format by converting the dict into
        # column wise orientation.
        dataframe = pd.DataFrame({
            col: [run_info[col] for run_info in model_runs.values()]
            for col in columns if col != 'rank'
        })

        # Give it an index, even if not in the `include`
        dataframe.set_index('model_id', inplace=True)

        # Add the `rank` column if needed, dropping `cost` if it's not
        # requested by the user
        if 'rank' in columns:
            dataframe.sort_values(by='cost', ascending=True, inplace=True)
            dataframe.insert(column='rank',
                             value=range(1, len(dataframe) + 1),
                             loc=list(columns).index('rank') - 1)  # account for `model_id`

            if 'cost' not in columns:
                dataframe.drop('cost', inplace=True)

        # Decide on the sort order depending on what it gets sorted by
        descending_columns = ['ensemble_weight', 'duration']
        if sort_order == 'auto':
            ascending_param = False if sort_by in descending_columns else True
        else:
            ascending_param = False if sort_order == 'descending' else True

        # Sort by the given column name, defaulting to 'model_id' if not present
        if sort_by not in dataframe.columns:
            self.automl_._logger.warning(f"sort_by = '{sort_by}' was not present"
                                         ", defaulting to sort on the index "
                                         "'model_id'")
            sort_by = 'model_id'

        # Cost can be the same but leave rank all over the place
        if 'rank' in columns and sort_by == 'cost':
            dataframe.sort_values(by=[sort_by, 'rank'],
                                  ascending=[ascending_param, True],
                                  inplace=True)
        else:
            dataframe.sort_values(by=sort_by,
                                  ascending=ascending_param,
                                  inplace=True)

        # Lastly, just grab the top_k
        if top_k == 'all' or top_k >= len(dataframe):
            top_k = len(dataframe)

        dataframe = dataframe.head(top_k)

        return dataframe

    @staticmethod
    def _leaderboard_columns() -> Dict[Literal['all', 'simple', 'detailed'], List[str]]:
        all = [
            "model_id", "rank", "ensemble_weight", "type", "cost", "duration",
            "config_id", "train_loss", "seed", "start_time", "end_time",
            "budget", "status", "data_preprocessors", "feature_preprocessors",
            "balancing_strategy", "config_origin"
        ]
        simple = [
            "model_id", "rank", "ensemble_weight", "type", "cost", "duration"
        ]
        detailed = all
        return {'all': all, 'detailed': detailed, 'simple': simple}

    def _get_automl_class(self):
        raise NotImplementedError()

    def get_configuration_space(
        self,
        X: SUPPORTED_FEAT_TYPES,
        y: Union[SUPPORTED_TARGET_TYPES, spmatrix],
        X_test: Optional[SUPPORTED_FEAT_TYPES] = None,
        y_test: Optional[Union[SUPPORTED_TARGET_TYPES, spmatrix]] = None,
        dataset_name: Optional[str] = None,
        feat_type: Optional[List[str]] = None,
    ):
        """
        Returns the Configuration Space object, from which Auto-Sklearn
        will sample configurations and build pipelines.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]
            Array with the training features, used to get characteristics like
            data sparsity
        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            Array with the problem labels
        X_test : array-like or sparse matrix of shape = [n_samples, n_features]
            Array with features used for performance estimation
        y_test : array-like, shape = [n_samples] or [n_samples, n_outputs]
            Array with the problem labels for the testing split
        dataset_name: Optional[str]
            A string to tag the Auto-Sklearn run
        """
        if self.automl_ is None:
            self.automl_ = self.build_automl()
        return self.automl_.fit(
            X, y,
            X_test=X_test, y_test=y_test,
            dataset_name=dataset_name,
            feat_type=feat_type,
            only_return_configuration_space=True,
        ) if self.automl_.configuration_space is None else self.automl_.configuration_space

class AutoSklearnClassifier(AutoSklearnEstimator, ClassifierMixin):
    """
    This class implements the classification task.

    """

    def fit(self, X, y,
            X_test=None,
            y_test=None,
            feat_type=None,
            dataset_name=None):
        """Fit *auto-sklearn* to given training set (X, y).

        Fit both optimizes the machine learning models and builds an ensemble
        out of them. To disable ensembling, set ``ensemble_size==0``.

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The target classes.

        X_test : array-like or sparse matrix of shape = [n_samples, n_features]
            Test data input samples. Will be used to save test predictions for
            all models. This allows to evaluate the performance of Auto-sklearn
            over time.

        y_test : array-like, shape = [n_samples] or [n_samples, n_outputs]
            Test data target classes. Will be used to calculate the test error
            of all models. This allows to evaluate the performance of
            Auto-sklearn over time.

        feat_type : list, optional (default=None)
            List of str of `len(X.shape[1])` describing the attribute type.
            Possible types are `Categorical` and `Numerical`. `Categorical`
            attributes will be automatically One-Hot encoded. The values
            used for a categorical attribute must be integers, obtained for
            example by `sklearn.preprocessing.LabelEncoder
            <https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.LabelEncoder.html>`_.

        dataset_name : str, optional (default=None)
            Create nicer output. If None, a string will be determined by the
            md5 hash of the dataset.

        Returns
        -------
        self

        """
        # AutoSklearn does not handle sparse y for now
        y = convert_if_sparse(y)

        # Before running anything else, first check that the
        # type of data is compatible with auto-sklearn. Legal target
        # types are: binary, multiclass, multilabel-indicator.
        target_type = type_of_target(y)
        supported_types = ['binary', 'multiclass', 'multilabel-indicator']
        if target_type not in supported_types:
            raise ValueError("Classification with data of type {} is "
                             "not supported. Supported types are {}. "
                             "You can find more information about scikit-learn "
                             "data types in: "
                             "https://scikit-learn.org/stable/modules/multiclass.html"
                             "".format(
                                    target_type,
                                    supported_types
                                )
                             )

        # remember target type for using in predict_proba later.
        self.target_type = target_type

        super().fit(
            X=X,
            y=y,
            X_test=X_test,
            y_test=y_test,
            feat_type=feat_type,
            dataset_name=dataset_name,
        )

        # After fit, a classifier is expected to define classes_
        # A list of class labels known to the classifier, mapping each label
        # to a numerical index used in the model representation our output.
        self.classes_ = self.automl_.InputValidator.target_validator.classes_

        return self

    def predict(self, X, batch_size=None, n_jobs=1):
        """Predict classes for X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_labels]
            The predicted classes.

        """
        return super().predict(X, batch_size=batch_size, n_jobs=n_jobs)

    def predict_proba(self, X, batch_size=None, n_jobs=1):

        """Predict probabilities of classes for all samples X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        batch_size : int (optional)
            Number of data points to predict for (predicts all points at once
            if ``None``.
        n_jobs : int

        Returns
        -------
        y : array of shape = [n_samples, n_classes] or [n_samples, n_labels]
            The predicted class probabilities.

        """
        pred_proba = super().predict_proba(
            X, batch_size=batch_size, n_jobs=n_jobs)

        # Check if all probabilities sum up to 1.
        # Assert only if target type is not multilabel-indicator.
        if self.target_type not in ['multilabel-indicator']:
            assert(
                np.allclose(
                    np.sum(pred_proba, axis=1),
                    np.ones_like(pred_proba[:, 0]))
            ), "prediction probability does not sum up to 1!"

        # Check that all probability values lie between 0 and 1.
        assert(
            (pred_proba >= 0).all() and (pred_proba <= 1).all()
        ), "found prediction probability value outside of [0, 1]!"

        return pred_proba

    def _get_automl_class(self):
        return AutoMLClassifier

class AutoSklearnRegressor(AutoSklearnEstimator, RegressorMixin):
    """
    This class implements the regression task.

    """

    def fit(self, X, y,
            X_test=None,
            y_test=None,
            feat_type=None,
            dataset_name=None):
        """Fit *Auto-sklearn* to given training set (X, y).

        Fit both optimizes the machine learning models and builds an ensemble
        out of them. To disable ensembling, set ``ensemble_size==0``.

        Parameters
        ----------

        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The training input samples.

        y : array-like, shape = [n_samples] or [n_samples, n_targets]
            The regression target.

        X_test : array-like or sparse matrix of shape = [n_samples, n_features]
            Test data input samples. Will be used to save test predictions for
            all models. This allows to evaluate the performance of Auto-sklearn
            over time.

        y_test : array-like, shape = [n_samples] or [n_samples, n_targets]
            The regression target. Will be used to calculate the test error
            of all models. This allows to evaluate the performance of
            Auto-sklearn over time.

        feat_type : list, optional (default=None)
            List of str of `len(X.shape[1])` describing the attribute type.
            Possible types are `Categorical` and `Numerical`. `Categorical`
            attributes will be automatically One-Hot encoded.

        dataset_name : str, optional (default=None)
            Create nicer output. If None, a string will be determined by the
            md5 hash of the dataset.

        Returns
        -------
        self

        """
        # Before running anything else, first check that the
        # type of data is compatible with auto-sklearn. Legal target
        # types are: continuous, continuous-multioutput, and the special cases:
        # multiclass : because [3.0, 1.0, 5.0] is considered as multiclass
        # binary: because [1.0, 0.0] is considered multiclass
        # AutoSklearn does not handle sparse y for now
        y = convert_if_sparse(y)

        target_type = type_of_target(y)
        supported_types = ['continuous', 'binary', 'multiclass', 'continuous-multioutput']
        if target_type not in supported_types:
            raise ValueError("Regression with data of type {} is "
                             "not supported. Supported types are {}. "
                             "You can find more information about scikit-learn "
                             "data types in: "
                             "https://scikit-learn.org/stable/modules/multiclass.html"
                             "".format(
                                    target_type,
                                    supported_types
                                )
                             )

        # Fit is supposed to be idempotent!
        # But not if we use share_mode.
        super().fit(
            X=X,
            y=y,
            X_test=X_test,
            y_test=y_test,
            feat_type=feat_type,
            dataset_name=dataset_name,
        )

        return self

    def predict(self, X, batch_size=None, n_jobs=1):
        """Predict regression target for X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_outputs]
            The predicted values.

        """
        return super().predict(X, batch_size=batch_size, n_jobs=n_jobs)

    def _get_automl_class(self):
        return AutoMLRegressor