from numpy import where, isnan, nan
from ziplime.lib.quantiles import quantiles
from ziplime.pipeline.terms.classifiers.classifier import Classifier
from ziplime.utils.numpy_utils import (
    int64_dtype,
)

from ziplime.pipeline.mixins import (
    SingleInputMixin,
)


class Quantiles(SingleInputMixin, Classifier):
    """
    A classifier computing quantiles over an input.
    """

    params = ("bins",)
    dtype = int64_dtype
    window_length = 0
    missing_value = -1

    def _compute(self, arrays, dates, assets, mask):
        data = arrays[0]
        bins = self.params["bins"]
        to_bin = where(mask, data, nan)
        result = quantiles(to_bin, bins)
        # Write self.missing_value into nan locations, whether they were
        # generated by our input mask or not.
        result[isnan(result)] = self.missing_value
        return result.astype(int64_dtype)

    def graph_repr(self):
        """Short repr to use when rendering Pipeline graphs."""
        return type(self).__name__ + "(%d)" % self.params["bins"]

