from abc import ABC, abstractmethod
import collections
import pandas as pd

from autoscalingsim.utils.error_check import ErrorChecker

class Correlator(ABC):

    _Registry = {}

    @abstractmethod
    def _compute_correlation(self, metrics_vals_1 : pd.Series, metrics_vals_2 : pd.Series, lag : int):

        pass

    def __init__(self, config : dict):

        history_buffer_size_raw = ErrorChecker.key_check_and_load('history_buffer_size', config, self.__class__.__name__)
        history_buffer_size_value = ErrorChecker.key_check_and_load('value', history_buffer_size_raw, self.__class__.__name__)
        history_buffer_size_unit = ErrorChecker.key_check_and_load('unit', history_buffer_size_raw, self.__class__.__name__)
        self.history_buffer_size = pd.Timedelta(history_buffer_size_value, unit = history_buffer_size_unit)

        max_time_lag_raw = ErrorChecker.key_check_and_load('max_time_lag', config, self.__class__.__name__)
        max_time_lag_value = ErrorChecker.key_check_and_load('value', max_time_lag_raw, self.__class__.__name__)
        max_time_lag_unit = ErrorChecker.key_check_and_load('unit', max_time_lag_raw, self.__class__.__name__)
        self.max_time_lag = pd.Timedelta(max_time_lag_value, unit = max_time_lag_unit)

        self.associated_service_metric_vals = pd.DataFrame()
        self.other_service_metric_vals = collections.defaultdict(pd.DataFrame)

    def _update_data(self, associated_service_metric_vals : pd.DataFrame, other_service_metric_vals : pd.DataFrame):

        if len(self.associated_service_metric_vals.index) > 0:
            self.associated_service_metric_vals = self.associated_service_metric_vals.append(associated_service_metric_vals[associated_service_metric_vals.index > max(self.associated_service_metric_vals.index)])
        else:
            self.associated_service_metric_vals = self.associated_service_metric_vals.append(associated_service_metric_vals)
        if self.associated_service_metric_vals.shape[0] > 0:
            self.associated_service_metric_vals = self.associated_service_metric_vals[self.associated_service_metric_vals.index >= max(self.associated_service_metric_vals.index) - self.history_buffer_size]

        for service_name, metric_vals in other_service_metric_vals.items():
            if len(self.other_service_metric_vals[service_name].index) > 0:
                self.other_service_metric_vals[service_name] = self.other_service_metric_vals[service_name].append(metric_vals[metric_vals.index > max(self.other_service_metric_vals[service_name].index)])
            else:
                self.other_service_metric_vals[service_name] = self.other_service_metric_vals[service_name].append(metric_vals)
            if self.other_service_metric_vals[service_name].shape[0] > 0:
                self.other_service_metric_vals[service_name] = self.other_service_metric_vals[service_name][self.other_service_metric_vals[service_name].index >= max(self.other_service_metric_vals[service_name].index) - self.history_buffer_size]

    def get_lagged_correlation(self, associated_service_metric_vals : pd.DataFrame, other_service_metric_vals : pd.DataFrame) -> dict:

        self._update_data(associated_service_metric_vals, other_service_metric_vals)
        min_resolution = self._get_minimal_resolution()
        max_lag = self.max_time_lag // min_resolution
        lags_range = range(-max_lag, max_lag)

        lags_per_service = dict()
        for service_name, metric_vals in self.other_service_metric_vals.items():
            other_service_metric_vals_resampled = metric_vals.resample(min_resolution).mean()
            associated_service_metric_vals_resampled = self.associated_service_metric_vals.resample(min_resolution).mean()

            common_len = min(associated_service_metric_vals_resampled.shape[0], other_service_metric_vals_resampled.shape[0])
            associated_service_metric_vals_inp = associated_service_metric_vals_resampled['value'][-common_len:]
            other_service_metric_vals_inp = other_service_metric_vals_resampled['value'][-common_len:]
            if associated_service_metric_vals_inp.shape == other_service_metric_vals_inp.shape:
                corr_raw = { lag : self._compute_correlation(associated_service_metric_vals_inp, other_service_metric_vals_inp, lag) for lag in lags_range }
                corr_pruned = { lag : corr for lag, corr in corr_raw.items() if not corr is None}

                if len(corr_pruned) > 0:
                    linear_correlation_df = pd.DataFrame({'lags': list(corr_pruned.keys()), 'correlation': list(corr_pruned.values())}).set_index('lags')
                    lags_per_service[service_name] = { 'lag': int(linear_correlation_df.correlation.idxmax()) * min_resolution, 'correlation': linear_correlation_df.correlation.max() }

        return lags_per_service

    def _get_minimal_resolution(self):

        minimas_to_consider = [pd.Timedelta(1, unit = 's')]

        for service_name, metric_vals in self.other_service_metric_vals.items():
            if metric_vals.shape[0] > 0:
                other_service_metric_vals_min_resolution = min(metric_vals.index.to_series().diff()[1:])
                if not other_service_metric_vals_min_resolution is pd.NaT: minimas_to_consider.append(other_service_metric_vals_min_resolution)

        associated_service_metric_vals_min_resolution = min(self.associated_service_metric_vals.index.to_series().diff()[1:])
        if not associated_service_metric_vals_min_resolution is pd.NaT: minimas_to_consider.append(associated_service_metric_vals_min_resolution)

        return min(minimas_to_consider)

    @classmethod
    def register(cls, name : str):

        def decorator(correlator_class):
            cls._Registry[name] = correlator_class
            return correlator_class

        return decorator

    @classmethod
    def get(cls, name : str):

        if not name in cls._Registry:
            raise ValueError(f'An attempt to use a non-existent {cls.__name__} {name}')

        return cls._Registry[name]

from .correlators import *

[((18, 34, 18, 121), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(18, 66, 18, 87): '"""history_buffer_size"""', (18, 89, 18, 95): 'config', (18, 97, 18, 120): 'self.__class__.__name__'}, {}), "('history_buffer_size', config, self.\n    __class__.__name__)", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((19, 36, 19, 126), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(19, 68, 19, 75): '"""value"""', (19, 77, 19, 100): 'history_buffer_size_raw', (19, 102, 19, 125): 'self.__class__.__name__'}, {}), "('value', history_buffer_size_raw, self.\n    __class__.__name__)", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((20, 35, 20, 124), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(20, 67, 20, 73): '"""unit"""', (20, 75, 20, 98): 'history_buffer_size_raw', (20, 100, 20, 123): 'self.__class__.__name__'}, {}), "('unit', history_buffer_size_raw, self.\n    __class__.__name__)", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((21, 35, 21, 107), 'pandas.Timedelta', 'pd.Timedelta', (), '', True, 'import pandas as pd\n'), ((23, 27, 23, 107), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(23, 59, 23, 73): '"""max_time_lag"""', (23, 75, 23, 81): 'config', (23, 83, 23, 106): 'self.__class__.__name__'}, {}), "('max_time_lag', config, self.__class__.__name__\n    )", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((24, 29, 24, 112), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(24, 61, 24, 68): '"""value"""', (24, 70, 24, 86): 'max_time_lag_raw', (24, 88, 24, 111): 'self.__class__.__name__'}, {}), "('value', max_time_lag_raw, self.__class__.\n    __name__)", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((25, 28, 25, 110), 'autoscalingsim.utils.error_check.ErrorChecker.key_check_and_load', 'ErrorChecker.key_check_and_load', ({(25, 60, 25, 66): '"""unit"""', (25, 68, 25, 84): 'max_time_lag_raw', (25, 86, 25, 109): 'self.__class__.__name__'}, {}), "('unit', max_time_lag_raw, self.__class__.\n    __name__)", False, 'from autoscalingsim.utils.error_check import ErrorChecker\n'), ((26, 28, 26, 86), 'pandas.Timedelta', 'pd.Timedelta', (), '', True, 'import pandas as pd\n'), ((28, 46, 28, 60), 'pandas.DataFrame', 'pd.DataFrame', ({}, {}), '()', True, 'import pandas as pd\n'), ((29, 41, 29, 78), 'collections.defaultdict', 'collections.defaultdict', ({(29, 65, 29, 77): 'pd.DataFrame'}, {}), '(pd.DataFrame)', False, 'import collections\n'), ((75, 31, 75, 58), 'pandas.Timedelta', 'pd.Timedelta', (), '', True, 'import pandas as pd\n')]