# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=no-else-return,invalid-name,consider-using-enumerate,abstract-method
"""Base class for model-based tuner
This type of tuner will fit a cost model and use some optimization methods to
find optimums points of cost model in space.
"""
import gc

import numpy as np

from .tuner import Tuner
from ..env import GLOBAL_SCOPE


class FeatureCache(object):
    """Feature cache manager for cache sharing between different cost models"""

    def __init__(self):
        self.feature_cache = {}

    def get(self, key):
        """Get feature cache dictionary for a key

        Parameters
        ----------
        key: str
            The key of a feature type

        Returns
        -------
        fea_cache: dict
            cache dictionary
        """
        if key not in self.feature_cache:
            self.feature_cache[key] = {}

        return self.feature_cache[key]

    def size(self, key):
        """ " Get the size of a feature cache dictionary

        Parameters
        ----------
        key: str
            The key of a feature type

        Returns
        -------
        n: int
        """
        return len(self.feature_cache.get(key, tuple()))

    def clear(self, key):
        """Clear feature cache for a key

        Parameters
        ----------
        key: str
            The key of a feature type
        """
        del self.feature_cache[key]
        self.feature_cache[key] = {}
        gc.collect()


class CostModel(object):
    """Cost model to predict the speed of a config"""

    def __init__(self):
        pass

    def fit(self, xs, ys, plan_size):
        """Fit to training data

        Parameters
        ----------
        xs: Array of int
            indexes of configs in the config space
        ys: Array of float
            The speed (flop, float number operations per second)
        plan_size: int
            The plan size of tuner
        """
        raise NotImplementedError()

    def fit_log(self, records, plan_size):
        """Fit training data from log.

        Parameters
        ----------
        records: Array of Tuple(MeasureInput, MeasureResult)
            The tuning records
        plan_size: int
            The plan size of tuner
        """
        raise NotImplementedError()

    def predict(self, xs, output_margin=False):
        """Predict the speed of configs

        Parameters
        ----------
        xs: Array of int
            The indexes of configs to predict
        output_margin: bool, optional
            Whether output the untransformed margin.
            When a model is used as base model, it should output untransformed margin

        Returns
        -------
        preds: Array of float
            The prediction
        """
        raise NotImplementedError()

    def load_basemodel(self, base_model):
        """Load base model for transfer learning

        Parameters
        ----------
        base_model: CostModel
                base model
        """
        raise NotImplementedError()

    def spawn_base_model(self):
        """Clone a base model with the same parameters.
        The base model is used to fit history data in transfer learning.

        Returns
        -------
        model: CostModel
            A model with the same hyperparameter (argument)
        """
        raise NotImplementedError()


class ModelOptimizer(object):
    """Optimizer used to find optimal points of cost model"""

    def __init__(self):
        pass