import numpy as np


def mean_agg(x):
    return x.mean()


def sparse_sum_agg(x):
    return x.sparse.fill_value if x.sparse.npoints == 0 else x.sum()


def sparse_max_agg(x):
    return x.sparse.fill_value if x.sparse.npoints == 0 else np.max(x.values.sp_values)


def purity_agg(x):
    value_counts = x.value_counts(sort=False)
    largest = value_counts.nlargest(1)
    return largest[0] / value_counts.sum()


def mode_agg(x):
    return x.mode()[0]


def get_basis(basis, nbins=None, agg=None, dimensions=2, precomputed=False, coords=True):
    if isinstance(dimensions, str):
        dimensions = int(dimensions)
    coordinate_columns = []
    for i in range(dimensions):
        coordinate_columns.append(basis + '_' + str(i + 1))
    full_name = basis + '_' + str(dimensions)
    if nbins is not None:
        full_name = full_name + '_' + str(nbins) + '_' + str(agg)
    return {'name': basis, 'dimensions': dimensions, 'coordinate_columns': coordinate_columns, 'nbins': nbins,
            'agg': agg, 'full_name': full_name, 'precomputed': precomputed, 'coords': coords}


class EmbeddingAggregator:

    def __init__(self, measures, dimensions, nbins, basis, agg_function='max', quick=True,
                 coords=True):
        self.nbins = nbins
        self.agg_function = agg_function
        self.measures = measures
        self.dimensions = dimensions
        self.add_count = nbins is not None and len(measures) == 0 and len(dimensions) == 0
        self.basis = basis
        self.quick = quick  # do not compute purity.
        self.coords = coords  # return coordinates

    @staticmethod
    def convert_coords_to_bin(df, nbins, coordinate_columns, bin_name, coordinate_column_to_range=None):
        # replace coordinates with bin, set df[name] to bin
        for name in coordinate_columns:
            values = df[name].values
            if coordinate_column_to_range is not None:
                view_column_range = coordinate_column_to_range[name]
                column_min = view_column_range[0]
                column_max = view_column_range[1]
            else:
                column_min = values.min()
                column_max = values.max()
            df[name] = np.floor(np.interp(values, [column_min, column_max], [0, nbins - 1])).astype(int)
        if len(coordinate_columns) == 2:
            df[bin_name] = df[coordinate_columns[0]] * nbins + df[coordinate_columns[1]]
        else:
            df[bin_name] = df[coordinate_columns[2]] + nbins * (
                    df[coordinate_columns[1]] + nbins * df[coordinate_columns[0]])

    def execute(self, df):
        result = {'coordinates': {}, 'values': {}}
        measures = self.measures
        dimensions = self.dimensions
        add_count = self.add_count
        basis = self.basis
        if basis is not None:
            result['name'] = basis['full_name']
        nbins = self.nbins
        agg_function = self.agg_function
        compute_purity = not self.quick and nbins is not None
        if nbins is not None:
            if basis['full_name'] not in df:
                EmbeddingAggregator.convert_coords_to_bin(df=df, nbins=nbins,
                    coordinate_columns=basis['coordinate_columns'], bin_name=basis['full_name'])

            if add_count:
                df['__count'] = 1.0

            # bin level summary, coordinates have already been converted
            agg_dict = {}
            full_basis_name = basis['full_name']
            for column in basis['coordinate_columns']:
                agg_dict[column] = 'min'
            for column in measures:
                if hasattr(df[column], 'sparse'):
                    if agg_function == 'max':
                        sparse_agg = sparse_max_agg
                    elif agg_function == 'mean':
                        sparse_agg = mean_agg
                    elif agg_function == 'sum':
                        sparse_agg = sparse_sum_agg
                    agg_dict[column] = sparse_agg
                else:
                    agg_dict[column] = agg_function

            for column in dimensions:
                agg_dict[column] = (mode_agg, purity_agg) if compute_purity else mode_agg
            if add_count:
                agg_dict['__count'] = 'sum'

            agg_df = df.groupby(full_basis_name).agg(agg_dict)

            if add_count:
                series = agg_df['__count']
                if compute_purity:
                    series = series[series.columns[0]]
                result['values']['__count'] = series
            if self.coords:
                result['coordinates']['bins'] = agg_df.index

        else:  # no binning
            if add_count:
                result['values']['__count'] = np.ones(len(df))
            agg_df = df
        for column in dimensions:
            series = agg_df[column]
            if nbins is not None:
                if compute_purity:
                    result['values'][column] = dict(value=series['mode_agg'], purity=series['purity_agg'])
                else:
                    result['values'][column] = dict(value=series)
            else:
                result['values'][column] = series
        for column in measures:
            series = agg_df[column]
            if compute_purity:
                series = series[series.columns[0]]
            is_sparse = hasattr(series, 'sparse')
            if is_sparse:
                result['values'][column] = series.sparse.to_dense()
                # result['values'][column] = dict(index=series.values.sp_index, values=series.values.sp_values)
            else:
                result['values'][column] = series

        if self.coords:
            for column in basis['coordinate_columns']:
                series = agg_df[column]
                if compute_purity:
                    series = series[series.columns[0]]
                result['coordinates'][column] = series
        return result
