# -*- coding: utf-8 -*-

from cf_units import Unit
import dask.array as da
import numpy as np

from ..util import change_units


SUPPORTED_OPERATORS = [
    '<',
    '>',
    '<=',
    '>=',
]
SUPPORTED_REDUCERS = [
    'min',
    'max',
    'sum',
    'mean',
]


NUMPY_OPERATORS = {
    '<': np.less,
    '>': np.greater,
    '<=': np.less_equal,
    '>=': np.greater_equal,
}
NUMPY_REDUCERS = {
    'min': np.min,
    'max': np.max,
    'sum': np.sum,
    'mean': np.mean,
}


DASK_OPERATORS = {
    '<': da.less,
    '>': da.greater,
    '<=': da.less_equal,
    '>=': da.greater_equal,
}
DASK_REDUCERS = {
    'min': da.min,
    'max': da.max,
    'sum': da.sum,
    'mean': da.mean,
}


def normalize_axis(axis, ndim):
    if isinstance(axis, list) and len(axis) == 1:
        axis = axis[0]
    if axis < 0:
        # just cope with negative axis numbers
        axis += ndim
    return axis


class IndexFunction:
    def __init__(self, standard_name=None, units=Unit('no_unit')):
        super().__init__()
        self.standard_name = standard_name
        self.units = units
        self.extra_coords = []

    def prepare(self, input_cube):
        pass


class ThresholdMixin:
    def __init__(self, threshold, condition, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.threshold = threshold
        self.condition = NUMPY_OPERATORS[condition]
        self.lazy_condition = DASK_OPERATORS[condition]
        self.extra_coords.append(threshold.copy())

    def prepare(self, input_cube):
        threshold = self.threshold
        change_units(threshold,
                     input_cube.units,
                     input_cube.standard_name)
        super().prepare(input_cube)


class ReducerMixin:
    def __init__(self, reducer, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.reducer = NUMPY_REDUCERS[reducer]
        self.lazy_reducer = DASK_REDUCERS[reducer]