## @package bucket_weighted
# Module caffe2.python.layers.bucket_weighted





import logging
import numpy as np

from caffe2.python import core, schema
from caffe2.python.layers.layers import (
    get_categorical_limit,
    ModelLayer,
)

from caffe2.python.layers.tags import Tags

logger = logging.getLogger(__name__)


class BucketWeighted(ModelLayer):
    def __init__(self, model, input_record, max_score=0, bucket_boundaries=None,
                 hash_buckets=True, weight_optim=None, name="bucket_weighted"):
        super(BucketWeighted, self).__init__(model, name, input_record)

        assert isinstance(input_record, schema.List), "Incorrect input type"
        self.bucket_boundaries = bucket_boundaries
        self.hash_buckets = hash_buckets
        if bucket_boundaries is not None:
            self.shape = len(bucket_boundaries) + 1
        elif max_score > 0:
            self.shape = max_score
        else:
            self.shape = get_categorical_limit(input_record)

        self.bucket_w = self.create_param(param_name='bucket_w',
                                       shape=[self.shape, ],
                                       initializer=('ConstantFill', {'value': 1.0}),
                                       optimizer=weight_optim)

        self.output_schema = schema.Struct(
            ('bucket_weights',
                schema.Scalar((np.float32, self.shape),
                              self.get_next_blob_reference("bucket_w_gather")))
        )

        self.tags.update({Tags.HANDLE_AS_SPARSE_LAYER})

    def get_memory_usage(self):
        return self.shape

    def add_ops(self, net):
        if self.bucket_boundaries is not None:
            buckets_int = net.Bucketize(
                self.input_record.values(),
                "buckets_int",
                boundaries=self.bucket_boundaries
            )
        else:
            buckets = self.input_record.values()
            buckets_int = net.Cast(
                buckets,
                "buckets_int",
                to=core.DataType.INT32
            )
        if self.hash_buckets:
            buckets_int = net.IndexHash(
                buckets_int, "hashed_buckets_int", seed=0, modulo=self.shape
            )
        net.Gather(
            [self.bucket_w, buckets_int],
            self.output_schema.bucket_weights.field_blobs())