# -*- coding: utf-8 -*-
"""FLAVORS.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/18WHD_qE3MQkJy51ZjQkq9qR7EhVyLlAT
"""

from collections import defaultdict
from heapq import heappop, heappush
from sklearn.utils import check_array
import numpy as np
import random
import datetime
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from joblib import Parallel, delayed

class RankOrderTree:
    def __init__(self):
        self.tree = []

    def insert(self, score):
        heappush(self.tree, -score)

    def remove_max(self):
        heappop(self.tree)

    def get_best_score (self):
        if self.tree:
            return -self.tree[0]
        return 0

    def get_second_best_score(self):
        if len(self.tree) >= 2:
            return -self.tree[1]
        return 0

class FLAVORS:
    def default_metric(X,y):
      return cross_val_score(LogisticRegression(),X,y,error_score='raise').mean()
    def __init__(self, budget, minimize=False,metrics=[], c_sub=4):
      self.leaderboard = []
      #self.subset_prob_dct = {}
      self.budget=budget
      self.c_sub = c_sub
      self.minimize = minimize
      self.fitted=False
      self.dup_dct=defaultdict(int)
      metrics=[FLAVORS.default_metric] if not metrics else metrics

      metrics=[FLAVORS.validate_input_arrays(metric) for metric in metrics]
      for metric in metrics:
        FLAVORS.test_custom_function(metric)
      self.metrics=metrics

    def update_pareto_history(self, new_tuple):
        self.pareto_history.append(new_tuple)

    def calculate_hypervolume(self):
        n = len(self.pareto_history)

        # Sort the history based on the first objective value in ascending order
        self.pareto_history.sort(key=lambda x: x[0])

        # Initialize the hypervolume to zero
        hypervolume = 0.0

        # Iterate through the history in reverse order
        for i in range(n - 1, -1, -1):
            # Calculate the volume of the current point relative to the origin (0, 0)
            volume = 1.0
            for j in range(len(self.pareto_history[i])):
                volume *= max(0.0, self.pareto_history[i][j])

            # Add the volume to the hypervolume
            hypervolume += volume

            # Calculate the volume of the current point relative to the previous point
            for j in range(i - 1, -1, -1):
                volume = 1.0
                for k in range(len(self.pareto_history[i])):
                    volume *= max(0.0, min(self.pareto_history[i][k], self.pareto_history[j][k]))

                # Subtract the volume from the hypervolume
                hypervolume -= volume

        return hypervolume

    @staticmethod
    def validate_input_arrays(func):

        def wrapper(X, y):
            # Check X and y using check_array
            X = check_array(X)
            y = check_array(y,ensure_2d=False)
            return func(X, y)
        return wrapper
    @staticmethod
    def test_custom_function(custom_metric):
        # Create example arrays
        X = np.random.rand(100,3)
        y = np.array([1]*50+[0]*50)

        try:
            # Call the custom metric function
            result = custom_metric(X, y)  # Pass `None` as the model argument for testing purposes
        except Exception as e:
            raise RuntimeError("Error occurred while testing the custom metric function") from e

        assert isinstance(result,(int,float)), 'must return number'

    def estimate_cost_feat(self, feat):

        if self.k_dct[0,feat] == 0:
            return self.initial_ECI(feat)

        K0, K1, K2 = self.k_dct[:, feat]

        delta= self.error_history[feat].get_best_score() - self.error_history[feat].get_second_best_score()
        if delta == 0:
            delta = self.current_error
            tau = K0
        else:
            tau = K0 - K2

        assert K1 >= K2, 'K1 has to be greater than or equal to K2'
        assert K0>= K1, 'K0 has to be greater than or equal to K1'

        return self.calculate_ECI(K0, K1, K2, delta, self.best_error, self.current_error, tau)

    def initial_ECI(self,feat):

        ECI1 = self.placeholder_coefficient* self.unique_counts[feat]/self.unique_counts[self.fastest_feat] if self.placeholder_coefficient!=np.inf else self.placeholder_coefficient
        ECI2 = self.c * self.kl
        return min(ECI1, ECI2)

    def calculate_ECI(self, K0, K1, K2, delta, best_error, feature_error, tau):
        ECI1 = max(K0 - K1, K1 - K2)
        ECI2 = self.c * self.kl
        if feature_error == best_error:
            ECI = min(ECI1, ECI2)
        else:
            ECI = max((feature_error - best_error) * (tau) / delta, min(ECI1, ECI2))

        return ECI

    def evaluate_subset(self, subset):
        start=datetime.datetime.now()
        sign = 1 if self.minimize else -1

        self.dup_dct[str(np.sort(subset))]+=1

        if len(self.metrics)>1:
          tuple_metric=tuple(metric(self.X[:,subset],self.y) for metric in self.metrics)
          self.update_pareto_history(tuple_metric)
          self.current_error= self.calculate_hypervolume()
          self.error_marker=tuple_metric

        else:
          self.current_error= sign* self.metrics[0](self.X[:,subset],self.y)
          self.error_marker= self.current_error

        return (datetime.datetime.now()-start).total_seconds()

    def update_v1(self, feat,time_cost):
        if time_cost<self.placeholder_coefficient:
          self.placeholder_coefficient=time_cost
          self.fastest_feat=feat

        self.k_dct[0,feat] += time_cost

        if self.current_error < self.best_error:
            self.k_dct[2,feat] = self.k_dct[1,feat]
            self.k_dct[1,feat] = self.k_dct[0,feat]

            self.error_history[feat].insert(self.current_error)

    def update_search(self, subset,time_cost):
        n_feats = len(subset)

        # Update k0, k1, k2 dictionaries in parallel
        lst=[feat for feat in subset if self.k_dct[0,feat]==0]
        Parallel(n_jobs=-1,require='sharedmem')(delayed(self.update_v1)(feat, time_cost) for feat in subset)

        self.iters += 1
        if self.current_error < self.best_error:
          self.iters_best = self.iters
          self.leaderboard.append((self.error_marker, subset))
          self.best_error = self.current_error

        self.c = min(self.c_sub, self.n_feats / n_feats)
        self.kl = time_cost

        for feat in subset:
          if feat not in lst:
            self.feat_eci_dct[feat]=self.estimate_cost_feat(feat)
          else:
            self.feat_eci_dct[feat]=self.initial_ECI(feat)

    def adjust_step_size(self, d):
        self.step_size = np.round(np.sqrt(d))
        no_improvement_count = self.iters - self.iters_best
        if no_improvement_count > (2 ** (d - 1)):
            reduction_ratio = self.iters / self.iters_best if self.iters_best != 0 else 1.1
            self.step_size /= reduction_ratio

        return np.round(self.step_size)

    def search_strategy(self, n_feats):

      unscaled_w = np.array([self.feat_eci_dct[feat] for feat in range(self.n_feats)])

      scaled_w = unscaled_w / unscaled_w.sum()
      remain_w= scaled_w.copy()
      selected_feats = []
      remaining_feats = set(range(self.n_feats))

      for _ in range(n_feats):
        feat = random.choices(list(remaining_feats), weights=remain_w, k=1)[-1]
        selected_feats.append(feat)
        remaining_feats.remove(feat)
        remain_w = scaled_w[list(remaining_feats)]
        remain_w=remain_w/remain_w.sum() if remain_w.sum()!=0 else np.array([1/len(remaining_feats) for x in remaining_feats])

      return selected_feats

    def select_best_feature_subset(self, budget):
        marker = datetime.datetime.now() + datetime.timedelta(seconds=budget)
        self.num_features = 1 if not self.fitted else self.num_features
        self.num_features+=int(self.adjust_step_size(self.num_features))

        subset = [np.random.choice(range(self.n_feats)) for _ in range(self.num_features)]
        time_cost=self.evaluate_subset(subset)

        if not self.fitted:
          for feat in subset:
            self.error_history[feat].insert(self.current_error)

          self.best_error = self.current_error
        self.update_search(subset,time_cost)

        while datetime.datetime.now() < marker:

            self.num_features += int(self.adjust_step_size(self.num_features))

            if self.num_features > self.n_feats:
                self.num_features = np.random.choice(range(2, self.n_feats))

            subset = self.search_strategy(self.num_features)
            time_cost=self.evaluate_subset(subset)
            self.update_search(subset,time_cost)

        if not len(self.leaderboard) > 0:
          self.leaderboard.append((self.error_marker, subset))

    def fit(self,X,y=None,budget=None):
      self.budget= budget if budget else self.budget
      self.X = check_array(X)
      self.y = check_array(y,ensure_2d=False)
      self.n_feats = self.X.shape[-1]

      if not len(self.leaderboard) > 0:
        self.feat_eci_dct= np.full(self.n_feats,0,dtype=np.float)
        print(self.feat_eci_dct)
        self.error_history=defaultdict(RankOrderTree)

        for feat in range(self.n_feats):
          self.error_history[feat].insert(0)

        self.k_dct= np.zeros((3, self.n_feats))

        self.current_error = np.nan
        self.best_error = np.nan
        self.kl = np.nan
        self.c = np.nan
        self.placeholder_coefficient = 2 * np.inf
        self.iters = 0
        self.iters_best = 0

        self.fastest_feat=np.nan
        self.unique_counts = np.apply_along_axis(lambda x: len(np.unique(x)), axis=0, arr=self.X)

        if len(self.metrics)>1:
          self.pareto_history = []
          self.hypervolume_values = []

      else:
        assert X.shape[-1]==self.feat_eci_dct.shape[-1], 'new input shape does not match warm shape'

      self.select_best_feature_subset(self.budget)
      self.fitted=True

      return self

    def transform(self,X):
      assert self.fitted, 'not fitted'
      X=check_array(X)
      assert X.shape[-1]== self.X.shape[-1], 'transform shape does not match fit shape'
      lowest_score_tuple = min(self.leaderboard, key=lambda x: x[0])
      lowest_score_indices = [index for score, index in self.leaderboard if score == lowest_score_tuple[0]]
      return X[:,lowest_score_indices[0]]

