"""
The standard way of using djongo is to import models.py
in place of Django's standard models module.

Djongo Fields is where custom fields for working with
MongoDB is defined.

 - EmbeddedModelField
 - ArrayModelField
 - ArrayReferenceField
 - GenericReferenceField

These are the main fields for working with MongoDB.
"""

import functools
import typing

from bson import ObjectId
from django import forms
from django.core.exceptions import ValidationError
from django.db import connections as pymongo_connections
from django.db import router, connections, transaction
from django.db.models import (
    Manager, Model, Field, AutoField,
    ForeignKey, BigAutoField
)
from django.forms import modelform_factory
from django.utils.functional import cached_property
from django.utils.html import format_html_join, format_html
from django.utils.safestring import mark_safe
from django.utils.translation import gettext_lazy as _

def make_mdl(model, model_dict):
    """
    Builds an instance of model from the model_dict.
    """
    for field_name in model_dict:
        field = model._meta.get_field(field_name)
        model_dict[field_name] = field.to_python(model_dict[field_name])

    return model(**model_dict)

def useful_field(field):
    return field.concrete and not isinstance(field, (AutoField, BigAutoField))

class ModelSubterfuge:

    def __init__(self, embedded_model):
        self.subterfuge = embedded_model

class DjongoManager(Manager):
    """
    This modified manager allows to issue Mongo functions by prefixing
    them with 'mongo_'.

    This module allows methods to be passed directly to pymongo.
    """

    def __getattr__(self, name):
        if name.startswith('mongo'):
            name = name[6:]
            return getattr(self._client, name)
        else:
            return super().__getattr__(name)

    @property
    def _client(self):
        return (
            pymongo_connections[self.db]
                .cursor()
                .db_conn[self.model._meta.db_table]
        )

class FormlessField(Field):
    empty_strings_allowed = False

    def formfield(self, **kwargs):
        raise TypeError(
            'A Formless Field cannot be modified from Django Admin.'
        )

class ListField(FormlessField):
    """
    MongoDB allows the saving of python lists as BSON Array type data. The `ListField` is useful in such cases.
    """
    empty_strings_allowed = False

    def get_db_prep_value(self, value, connection, prepared=False):
        if not isinstance(value, list):
            raise ValueError(
                f'Value: {value} must be of type list'
            )
        return value

    def to_python(self, value):
        if not isinstance(value, list):
            raise ValueError(
                f'Value: {value} stored in DB must be of type list'
                'Did you miss any Migrations?'
            )
        return value

class DictField(FormlessField):
    """
    MongoDB allows the saving of python dicts as BSON object type data. The `DictField` is useful in such cases.
    """
    empty_strings_allowed = False

    def get_db_prep_value(self, value, connection, prepared=False):
        if not isinstance(value, dict):
            raise ValueError(
                f'Value: {value} must be of type dict'
            )
        return value

    def to_python(self, value):
        if not isinstance(value, dict):
            raise ValueError(
                f'Value: {value} stored in DB must be of type dict'
                'Did you miss any Migrations?'
            )
        return value

class ArrayModelField(Field):
    """
    Implements an array of objects inside the document.

    The allowed object type is defined on model declaration. The
    declared instance will accept a python list of instances of the
    given model as its contents.

    The model of the container must be declared as abstract, thus should
    not be treated as a collection of its own.

    Example:

    class Author(models.Model):
        name = models.CharField(max_length=100)
        email = models.CharField(max_length=100)

        class Meta:
            abstract = True

    class AuthorForm(forms.ModelForm):
        class Meta:
            model = Author
            fields = (
                'name', 'email'
            )

    class MultipleBlogPosts(models.Model):
        h1 = models.CharField(max_length=100)
        content = models.ArrayModelField(
            model_container=BlogContent,
            model_form_class=BlogContentForm
        )

    """

    empty_strings_allowed = False

    def __init__(self,
                 model_container: typing.Type[Model],
                 model_form_class: typing.Type[forms.ModelForm] = None,
                 model_form_kwargs_l: dict = None,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_container = model_container
        self.model_form_class = model_form_class

        if model_form_kwargs_l is None:
            model_form_kwargs_l = {}
        self.model_form_kwargs_l = model_form_kwargs_l

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        kwargs['model_container'] = self.model_container

        if self.model_form_class is not None:
            kwargs['model_form_class'] = self.model_form_class

        if self.model_form_kwargs_l:
            kwargs['model_form_kwargs_l'] = self.model_form_kwargs_l

        return name, path, args, kwargs

    def get_db_prep_value(self, value, connection, prepared=False):
        if prepared:
            return value

        if not isinstance(value, list):
            raise ValueError(
                'Expected value to be type list,'
                f'Got type {type(value)} instead'
            )

        ret = []
        for a_mdl in value:
            mdl_ob = {}
            if not isinstance(a_mdl, Model):
                raise ValueError('Array items must be Model instances')
            for fld in a_mdl._meta.get_fields():
                if not useful_field(fld):
                    continue
                fld_value = getattr(a_mdl, fld.attname)
                mdl_ob[fld.attname] = fld.get_db_prep_value(fld_value, connection, prepared)
            ret.append(mdl_ob)

        return ret

    def from_db_value(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        """
        Overrides standard to_python method from django models to allow
        correct translation of Mongo array to a python list.
        """
        if value is None:
            return value

        assert isinstance(value, list)
        ret = []
        for mdl_dict in value:
            if isinstance(mdl_dict, self.model_container):
                ret.append(mdl_dict)
                continue
            mdl = make_mdl(self.model_container, mdl_dict)
            ret.append(mdl)

        return ret

    def formfield(self, **kwargs):
        """
        Returns the formfield for the array.
        """
        defaults = {
            'form_class': ArrayFormField,
            'model_container': self.model_container,
            'model_form_class': self.model_form_class,
            'name': self.attname,
            'mdl_form_kw_l': self.model_form_kwargs_l

        }
        defaults.update(kwargs)
        return super().formfield(**defaults)

    def validate(self, value, model_instance):
        super().validate(value, model_instance)
        errors = []
        for mdl in value:
            try:
                mdl.full_clean()
            except ValidationError as e:
                errors.append(e)
        if errors:
            raise ValidationError(errors)

def _get_model_form_class(model_form_class, model_container, admin, request):
    if not model_form_class:
        form_kwargs = dict(
            form=forms.ModelForm,
            fields=forms.ALL_FIELDS,
        )

        if admin and request:
            form_kwargs['formfield_callback'] = functools.partial(
                admin.formfield_for_dbfield, request=request)

        model_form_class = modelform_factory(model_container, **form_kwargs)

    return model_form_class

class ArrayFormField(forms.Field):
    def __init__(self, name, model_form_class, model_container, mdl_form_kw_l,
                 widget=None, admin=None, request=None, *args, **kwargs):

        self.name = name
        self.model_container = model_container
        self.model_form_class = _get_model_form_class(
            model_form_class, model_container, admin, request)
        self.mdl_form_kw_l = mdl_form_kw_l
        self.admin = admin
        self.request = request

        if not widget:
            widget = ArrayFormWidget(self.model_form_class.__name__)

        error_messages = {
            'incomplete': 'Enter all required fields.',
        }

        self.ArrayFormSet = forms.formset_factory(
            self.model_form_class, can_delete=True)
        super().__init__(error_messages=error_messages,
                         widget=widget, *args, **kwargs)

    def clean(self, value):
        if not value:
            return []

        form_set = self.ArrayFormSet(value, prefix=self.name)
        if form_set.is_valid():
            ret = []
            for itm in form_set.cleaned_data:
                if itm.get('DELETE', True):
                    continue
                itm.pop('DELETE')
                ret.append(self.model_form_class._meta.model(**itm))
            return ret

        else:
            raise ValidationError(form_set.errors + form_set.non_form_errors())

    def has_changed(self, initial, data):
        form_set_initial = []
        for init in initial:
            form_set_initial.append(
                forms.model_to_dict(
                    init,
                    fields=self.model_form_class._meta.fields,
                    exclude=self.model_form_class._meta.exclude
                )
            )
        form_set = self.ArrayFormSet(data, initial=form_set_initial, prefix=self.name)
        return form_set.has_changed()

    def get_bound_field(self, form, field_name):
        return ArrayFormBoundField(form, self, field_name)

class ArrayFormBoundField(forms.BoundField):
    def __init__(self, form, field, name):
        super().__init__(form, field, name)

        data = self.data if form.is_bound else None
        initial = []
        if self.initial is not None:
            for ini in self.initial:
                if isinstance(ini, Model):
                    initial.append(
                        forms.model_to_dict(
                            ini,
                            fields=field.model_form_class._meta.fields,
                            exclude=field.model_form_class._meta.exclude
                        ))

        self.form_set = field.ArrayFormSet(data, initial=initial, prefix=name)

    def __getitem__(self, idx):
        if not isinstance(idx, (int, slice)):
            raise TypeError
        return self.form_set[idx]

    def __iter__(self):
        for form in self.form_set:
            yield form

    def __str__(self):
        table = format_html_join(
            '\n', '
{}
',
            ((form.as_table(),)
             for form in self.form_set))
        table = format_html(
            '\n
'
            '\n{}'
            '\n
',
            self.name,
            table)
        return format_html('{}\n{}', table, self.form_set.management_form)

    def __len__(self):
        return len(self.form_set)

class ArrayFormWidget(forms.Widget):
    def __init__(self, first_field_id, attrs=None):
        self.first_field_id = first_field_id
        super().__init__(attrs)

    def render(self, name, value, attrs=None, renderer=None):
        assert False

    def id_for_label(self, id_):
        return '{}-0-{}'.format(id_, self.first_field_id)

    def value_from_datadict(self, data, files, name):
        ret = {key: data[key] for key in data if key.startswith(name)}
        return ret

    def value_omitted_from_data(self, data, files, name):
        for key in data:
            if key.startswith(name):
                return False
        return True

class EmbeddedModelField(Field):
    """
    Allows for the inclusion of an instance of an abstract model as
    a field inside a document.

    Example:

    class Blog(models.Model):
        name = models.CharField(max_length=100)
        tagline = models.TextField()

        class Meta:
            abstract = True

    class BlogForm(forms.ModelForm):
        class Meta:
            model = Blog
            fields = (
                'comment', 'author'
            )

    class Entry(models.Model):
        blog = models.EmbeddedModelField(
            model_container=Blog,
            model_form_class=BlogForm
        )

        headline = models.CharField(max_length=255)
        objects = models.DjongoManager()

    """
    empty_strings_allowed = False

    def __init__(self,
                 model_container: typing.Type[Model],
                 model_form_class: typing.Type[forms.ModelForm] = None,
                 model_form_kwargs: dict = None,
                 admin=None,
                 request=None,
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_container = model_container
        self.model_form_class = model_form_class
        self.null = True
        self.instance = None
        self.admin = admin
        self.request = request

        if model_form_kwargs is None:
            model_form_kwargs = {}
        self.model_form_kwargs = model_form_kwargs

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        kwargs['model_container'] = self.model_container
        if self.model_form_class is not None:
            kwargs['model_form_class'] = self.model_form_class
        return name, path, args, kwargs

    def pre_save(self, model_instance, add):
        value = getattr(model_instance, self.attname)
        if isinstance(value, ModelSubterfuge):
            return value

        subterfuge = ModelSubterfuge(value)
        # setattr(model_instance, self.attname, subterfuge)
        return subterfuge

    def get_db_prep_value(self, value, connection=None, prepared=False):
        if isinstance(value, dict):
            return value
        if isinstance(value, ModelSubterfuge):
            value = value.subterfuge
        if value is None and self.blank:
            return None
        if not isinstance(value, Model):
            raise ValueError(
                'Value: {value} must be instance of Model: {model}'.format(
                    value=value,
                    model=Model
                )
            )

        mdl_ob = {}
        for fld in value._meta.get_fields():
            if not useful_field(fld):
                continue
            fld_value = getattr(value, fld.attname)
            mdl_ob[fld.attname] = fld.get_db_prep_value(fld_value, connection, prepared)

        return mdl_ob

    def from_db_value(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        """
        Overrides Django's default to_python to allow correct
        translation to instance.
        """
        if value is None or isinstance(value, self.model_container):
            return value
        assert isinstance(value, dict)

        instance = make_mdl(self.model_container, value)
        return instance

    def formfield(self, **kwargs):
        defaults = {
            'form_class': EmbeddedFormField,
            'model_container': self.model_container,
            'model_form_class': self.model_form_class,
            'model_form_kw': self.model_form_kwargs,
            'name': self.attname
        }

        defaults.update(kwargs)
        return super().formfield(**defaults)

class EmbeddedFormField(forms.MultiValueField):
    def __init__(self, name, model_form_class, model_form_kw, model_container,
                 admin=None, request=None, *args, **kwargs):
        form_fields = []
        mdl_form_field_names = []
        widgets = []
        model_form_kwargs = model_form_kw.copy()

        try:
            model_form_kwargs['prefix'] = model_form_kwargs['prefix'] + '-' + name
        except KeyError:
            model_form_kwargs['prefix'] = name

        initial = kwargs.pop('initial', None)
        if initial:
            model_form_kwargs['initial'] = initial

        self.model_form_class = _get_model_form_class(
            model_form_class, model_container, admin, request)
        self.model_form_kwargs = model_form_kwargs
        self.admin = admin
        self.request = request

        error_messages = {
            'incomplete': 'Enter all required fields.',
        }

        self.model_form = self.model_form_class(**model_form_kwargs)
        for field_name, field in self.model_form.fields.items():
            if isinstance(field, (forms.ModelChoiceField, forms.ModelMultipleChoiceField)):
                continue
            form_fields.append(field)
            mdl_form_field_names.append(field_name)
            widgets.append(field.widget)

        widget = EmbeddedFormWidget(mdl_form_field_names, widgets)
        super().__init__(error_messages=error_messages, fields=form_fields,
                         widget=widget, require_all_fields=False, *args, **kwargs)

    def compress(self, clean_data_vals):
        model_field = dict(zip(self.model_form.fields.keys(), clean_data_vals))
        return self.model_form._meta.model(**model_field)

    def get_bound_field(self, form, field_name):
        if form.prefix:
            self.model_form.prefix = '{}-{}'.format(form.prefix, self.model_form.prefix)
        return EmbeddedFormBoundField(form, self, field_name)

    def bound_data(self, data, initial):
        if self.disabled:
            return initial
        return self.compress(data)

class EmbeddedFormBoundField(forms.BoundField):
    # def __getitem__(self, name):
    #     return self.field.model_form[name]
    #
    # def __getattr__(self, name):
    #     return getattr(self.field.model_form, name)

    def __str__(self):
        instance = self.value()
        model_form = self.field.model_form_class(instance=instance, **self.field.model_form_kwargs)

        return mark_safe(f'
\n{ model_form.as_table() }\n
')

class EmbeddedFormWidget(forms.MultiWidget):
    def __init__(self, field_names, *args, **kwargs):
        self.field_names = field_names
        super().__init__(*args, **kwargs)

    def decompress(self, value):
        if value is None:
            return []
        elif isinstance(value, list):
            return value
        elif isinstance(value, Model):
            return [getattr(value, f_n) for f_n in self.field_names]
        else:
            raise forms.ValidationError('Expected model-form')

    def value_from_datadict(self, data, files, name):
        ret = []
        for i, wid in enumerate(self.widgets):
            f_n = '{}-{}'.format(name, self.field_names[i])
            ret.append(wid.value_from_datadict(
                data, files, f_n
            ))
        return ret

    def value_omitted_from_data(self, data, files, name):
        return all(
            widget.value_omitted_from_data(
                data, files, '{}-{}'.format(name, self.field_names[i])
            )
            for i, widget in enumerate(self.widgets)
        )

class ObjectIdFieldMixin:
    description = _("ObjectId")

    def get_db_prep_value(self, value, connection, prepared=False):
        return self.to_python(value)

    def to_python(self, value):
        if isinstance(value, str):
            return ObjectId(value)
        return value

    def get_internal_type(self):
        return "ObjectIdField"

class GenericObjectIdField(ObjectIdFieldMixin, Field):
    empty_strings_allowed = False

class ObjectIdField(ObjectIdFieldMixin, AutoField):
    """
    For every document inserted into a collection MongoDB internally creates an field.
    The field can be referenced from within the Model as _id.
    """

    def __init__(self, *args, **kwargs):
        id_field_args = {
            'primary_key': True,
            'auto_created': True
        }
        id_field_args.update(kwargs)
        super().__init__(*args, **id_field_args)

    def get_prep_value(self, value):
        value = super(AutoField, self).get_prep_value(value)
        if value is None:
            return None
        return value

class ArrayReferenceManagerMixin:

    def _apply_rel_filters(self, queryset):
        """
        Filter the queryset for the instance this manager is bound to.
        """
        queryset._add_hints(instance=self.instance)
        if self._db:
            queryset = queryset.using(self._db)
        queryset = queryset.filter(**self.core_filters)

        return queryset

    def get_queryset(self):
        try:
            return self.instance._prefetched_objects_cache[self.field.related_query_name()]
        except (AttributeError, KeyError):
            queryset = super().get_queryset()
            return self._apply_rel_filters(queryset)

    def update_or_create(self, **kwargs):
        db = router.db_for_write(self.instance.__class__, instance=self.instance)
        obj, created = super(ArrayReferenceManagerMixin, self.db_manager(db)).update_or_create(**kwargs)
        # We only need to add() if created because if we got an object back
        # from get() then the relationship already exists.
        if created:
            self.add(obj)
        return obj, created

    update_or_create.alters_data = True

    def get_or_create(self, **kwargs):
        db = router.db_for_write(self.instance.__class__, instance=self.instance)
        obj, created = super(ArrayReferenceManagerMixin, self.db_manager(db)).get_or_create(**kwargs)
        # We only need to add() if created because if we got an object back
        # from get() then the relationship already exists.
        if created:
            self.add(obj)
        return obj, created

    get_or_create.alters_data = True

    def create(self, **kwargs):
        db = router.db_for_write(self.instance.__class__, instance=self.instance)
        new_obj = super(ArrayReferenceManagerMixin, self.db_manager(db)).create(**kwargs)
        self.add(new_obj)
        return new_obj

    create.alters_data = True

def create_reverse_array_reference_manager(superclass, rel):
    if issubclass(superclass, DjongoManager):
        baseclass = superclass
    else:
        baseclass = type('baseclass', (DjongoManager, superclass), {})

    class ReverseArrayReferenceManager(ArrayReferenceManagerMixin, baseclass):
        def __init__(self, instance):
            super().__init__()
            self.instance = instance
            self.model = rel.related_model
            self.field = rel.remote_field

            name = rel.remote_field.name
            self.core_filters = {name: instance}

        def __call__(self, *, manager):
            manager = getattr(self.model, manager)
            manager_class = create_reverse_array_reference_manager(manager.__class__, rel)
            return manager_class(instance=self.instance)

        do_not_call_in_templates = True

        def _apply_rel_filters(self, queryset):
            queryset = super()._apply_rel_filters(queryset)
            db = self._db or router.db_for_read(self.model, instance=self.instance)
            empty_strings_as_null = connections[db].features.interprets_empty_strings_as_nulls

            for field in self.field.foreign_related_fields:
                val = getattr(self.instance, field.attname)
                if val is None or (val == '' and empty_strings_as_null):
                    return queryset.none()
            return queryset

        def _make_filter(self, *objs):
            ids = set(obj.pk for obj in objs)
            return {
                self.model._meta.pk.name: {
                    '$in': list(ids)
                }
            }

        def add(self, *objs):
            _filter = self._make_filter(*objs)
            lh_field, rh_field = self.field.related_fields[0]
            self.mongo_update(
                _filter,
                {
                    '$addToSet': {
                        lh_field.get_attname():
                            getattr(self.instance, rh_field.get_attname())
                    }
                }
            )
            for obj in objs:
                fk_field = getattr(obj, lh_field.get_attname())
                fk_field.add(getattr(self.instance, rh_field.get_attname()))

        add.alters_data = True

        def remove(self, *objs):
            pass

        remove.alters_data = True

        def clear(self):
            pass

        clear.alters_data = True

        def set(self, objs, *, clear=False):
            pass

        set.alters_data = True

        def create(self, **kwargs):
            pass

        create.alters_data = True

    return ReverseArrayReferenceManager

def create_forward_array_reference_manager(superclass, rel):
    if issubclass(superclass, DjongoManager):
        baseclass = superclass
    else:
        baseclass = type('baseclass', (DjongoManager, superclass), {})

    class ArrayReferenceManager(ArrayReferenceManagerMixin, baseclass):

        def __init__(self, instance):
            super().__init__()

            self.instance = instance
            self.model = rel.model
            self.field = rel.remote_field
            self.instance_manager = DjongoManager()
            self.instance_manager.model = instance
            name = self.field.related_fields[0][1].attname
            ids = getattr(instance, self.field.attname) or []

            self.core_filters = {f'{name}__in': ids}

        def __call__(self, *, manager):
            manager = getattr(self.model, manager)
            manager_class = create_forward_array_reference_manager(manager.__class__, rel)
            return manager_class(instance=self.instance)

        do_not_call_in_templates = True

        def _apply_rel_filters(self, queryset):
            queryset = super()._apply_rel_filters(queryset)
            if not getattr(self.instance, self.field.attname):
                return queryset.none()
            return queryset

        def get_queryset(self):
            try:
                return self.instance._prefetched_objects_cache[self.field.related_query_name()]
            except (AttributeError, KeyError):
                queryset = super().get_queryset()
                return self._apply_rel_filters(queryset)

        def _make_filter(self):
            return {self.instance._meta.pk.name: self.instance.pk}

        def add(self, *objs):
            fks = getattr(self.instance, self.field.get_attname())
            if fks is None:
                fks = set()
                setattr(self.instance, self.field.get_attname(), fks)

            new_fks = set()
            rh_field = self.field.foreign_related_fields[0]
            for obj in objs:
                new_fks.add(getattr(obj, rh_field.get_attname()))
            fks.update(new_fks)

            db = router.db_for_write(self.instance.__class__, instance=self.instance)
            self.instance_manager.db_manager(db).mongo_update(
                self._make_filter(),
                {
                    '$addToSet': {
                        self.field.attname: {
                            '$each': list(new_fks)
                        }
                    }
                }
            )

        add.alters_data = True

        def remove(self, *objs):
            to_del = set(
                getattr(obj, self.field.foreign_related_fields[0].attname)
                for obj in objs
            )
            self._remove(to_del)

        remove.alters_data = True

        def _remove(self, to_del):
            fks = getattr(self.instance, self.field.attname)
            fks.difference_update(to_del)
            db = self._db or router.db_for_write(self.instance.__class__, instance=self.instance)
            self.instance_manager.db_manager(db).mongo_update(
                self._make_filter(),
                {
                    '$pull': {
                        self.field.attname: {
                            '$in': list(to_del)
                        }
                    }
                }
            )

        def clear(self):
            db = router.db_for_write(self.instance.__class__, instance=self.instance)
            self.instance_manager.db_manager(db).mongo_update(
                self._make_filter(),
                {
                    '$set': {
                        self.field.attname: []
                    }
                }
            )
            setattr(self.instance, self.field.attname, set())

        clear.alters_data = True

        def set(self, objs, *, clear=False):
            objs = tuple(objs)

            db = router.db_for_write(self.instance.__class__, instance=self.instance)
            with transaction.atomic(using=db, savepoint=False):
                if clear:
                    self.clear()
                    self.add(*objs)
                else:
                    fks = getattr(self.instance, self.field.attname)
                    rh_field = self.field.foreign_related_fields[0]
                    new_fks = set(getattr(obj, rh_field.get_attname()) for obj in objs)
                    to_del = fks - new_fks
                    self._remove(to_del)
                    fks = getattr(self.instance, self.field.attname)
                    to_add = []
                    for obj in objs:
                        if getattr(obj, rh_field.get_attname()) not in fks:
                            to_add.append(obj)
                    self.add(to_add)

        set.alters_data = True

    return ArrayReferenceManager

class ArrayReferenceDescriptor:

    def __init__(self, field_with_rel):
        self.field = field_with_rel

    @cached_property
    def related_manager_cls(self):
        related_model = self.field.related_model

        return create_forward_array_reference_manager(
            related_model._default_manager.__class__,
            self.field.remote_field,
        )

    def __get__(self, instance, cls=None):
        """
        Get the related objects through the reverse relation.

        With the example above, when getting ``parent.children``:

        - ``self`` is the descriptor managing the ``children`` attribute
        - ``instance`` is the ``parent`` instance
        - ``cls`` is the ``Parent`` class (unused)
        """
        if instance is None:
            return self

        return self.related_manager_cls(instance)

class ReverseArrayReferenceDescriptor:

    def __init__(self, related):
        self.rel = related

    @cached_property
    def related_manager_cls(self):
        related_model = self.rel.related_model

        return create_reverse_array_reference_manager(
            related_model._default_manager.__class__,
            self.rel,
        )

    def __get__(self, instance, cls=None):
        """
        Get the related objects through the relation.

        With the example above, when getting ``parent.children``:

        - ``self`` is the descriptor managing the ``children`` attribute
        - ``instance`` is the ``parent`` instance
        - ``cls`` is the ``Parent`` class (unused)
        """
        if instance is None:
            return self

        return self.related_manager_cls(instance)

# class ArrayManyToManyField(ManyToManyField):
#
#     def contribute_to_class(self, cls, name, **kwargs):
#         super().contribute_to_class(cls, name, **kwargs)
#         setattr(cls, self.name, ArrayManyToManyDescriptor(self.remote_field, reverse=False))

class ArrayReferenceField(ForeignKey):
    """
    When the entry gets saved, only a reference to the primary_key of the model is saved in the array.
    For all practical aspects, the ArrayReferenceField behaves like a Django ManyToManyField.
    """

    many_to_many = False
    many_to_one = True
    one_to_many = False
    one_to_one = False
    # rel_class = ManyToManyRel
    related_accessor_class = ReverseArrayReferenceDescriptor
    forward_related_accessor_class = ArrayReferenceDescriptor

    def __init__(self, to, on_delete=None, related_name=None, related_query_name=None,
                 limit_choices_to=None, parent_link=False, to_field=None,
                 db_constraint=True, **kwargs):

        on_delete = on_delete or self._on_delete
        super().__init__(to,
                         on_delete=on_delete,
                         related_name=related_name,
                         related_query_name=related_query_name,
                         limit_choices_to=limit_choices_to,
                         parent_link=parent_link,
                         to_field=to_field,
                         db_constraint=db_constraint,
                         **kwargs)

        self.concrete = False

    @staticmethod
    def _on_delete(collector, field, sub_objs, using):
        for model, instances in collector.data.items():
            for obj in sub_objs:
                getattr(obj, field.name).db_manager(using).remove(*instances)

    def from_db_value(self, value, expression, connection, context):
        return self.to_python(value)

    def to_python(self, value):
        if value is None:
            return set()
        return set(value)

    def get_db_prep_value(self, value, connection, prepared=False):
        if value is None:
            return []
        elif isinstance(value, set):
            return list(value)
        return value
        # return super().get_db_prep_value(value, connection, prepared)

    def get_db_prep_save(self, value, connection):
        return self.get_db_prep_value(value, connection)
        # if value is None:
        #     return []
        # return list(value)

    def validate(self, value, model_instance):
        pass
        # super().validate(list(value), model_instance)

    def save_form_data(self, instance, data):
        getattr(instance, self.name).set(list(data), clear=True)

    def formfield(self, *, using=None, **kwargs):
        defaults = {
            'form_class': forms.ModelMultipleChoiceField,
            'queryset': self.remote_field.model._default_manager.using(using),
            **kwargs,
        }
        # If initial is passed in, it's a list of related objects, but the
        # MultipleChoiceField takes a list of IDs.
        if defaults.get('initial') is not None:
            initial = defaults['initial']
            if callable(initial):
                initial = initial()
            defaults['initial'] = [i.pk for i in initial]
        return super().formfield(**defaults)


"""
This module is for creating optimal partition intervals for LSH Ensemble:

    [l1, u1], [l2, u2], ...

with all boundaries inclusive, using set size distribution.
"""

from collections import Counter
import numpy as np

def _compute_nfp_uniform(l, u, cum_counts, sizes):
    """Computes the expected number of false positives caused by using
    u to approximate set sizes in the interval [l, u], assuming uniform
    distribution of set sizes within the interval.

    Args:
        l: the lower bound on set sizes.
        u: the upper bound on set sizes.
        cum_counts: the complete cummulative distribution of set sizes.
        sizes: the complete domain of set sizes.

    Return (float): the expected number of false positives.
    """
    if l > u:
        raise ValueError("l must be less or equal to u")
    if l == 0:
        n = cum_counts[u]
    else:
        n = cum_counts[u]-cum_counts[l-1]
    return n * float(sizes[u] - sizes[l]) / float(2*sizes[u])

def _compute_nfps_uniform(cum_counts, sizes):
    """Computes the matrix of expected false positives for all possible
    sub-intervals of the complete domain of set sizes, assuming uniform
    distribution of set_sizes within each sub-intervals.

    Args:
        cum_counts: the complete cummulative distribution of set sizes.
        sizes: the complete domain of set sizes.

    Return (np.array): the 2-D array of expected number of false positives
        for every pair of [l, u] interval, where l is axis-0 and u is
        axis-1.
    """
    nfps = np.zeros((len(sizes), len(sizes)))
    # All u an l are inclusive bounds for intervals.
    # Compute p = 1, the NFPs
    for l in range(len(sizes)):
        for u in range(l, len(sizes)):
            nfps[l, u] = _compute_nfp_uniform(l, u, cum_counts, sizes)
    return nfps

def _compute_nfp_real(l, u, counts, sizes):
    """Computes the expected number of false positives caused by using
    u to approximate set sizes in the interval [l, u], using the real
    set size distribution.

    Args:
        l: the lower bound on set sizes.
        u: the upper bound on set sizes.
        counts: the complete distribution of set sizes.
        sizes: the complete domain of set sizes.

    Return (float): the expected number of false positives.
    """
    if l > u:
        raise ValueError("l must be less or equal to u")
    return np.sum((float(sizes[u])-sizes[l:u+1])/float(sizes[u])*counts[l:u+1])

def _compute_nfps_real(counts, sizes):
    """Computes the matrix of expected false positives for all possible
    sub-intervals of the complete domain of set sizes.

    Args:
        counts: the complete distribution of set sizes.
        sizes: the complete domain of set sizes.

    Return (np.array): the 2-D array of expected number of false positives
        for every pair of [l, u] interval, where l is axis-0 and u is
        axis-1.
    """
    nfps = np.zeros((len(sizes), len(sizes)))
    # All u an l are inclusive bounds for intervals.
    # Compute p = 1, the NFPs
    for l in range(len(sizes)):
        for u in range(l, len(sizes)):
            nfps[l, u] = _compute_nfp_real(l, u, counts, sizes)
    return nfps

def _compute_best_partitions(num_part, sizes, nfps):
    """Computes the optimal partitions given the size distributions
    and computed number of expected false positives for all sub-intervals.

    Args:
        num_part (int): The number of partitions to create.
        sizes (numpy.array): The complete domain of set sizes in sorted order.
        nfps (numpy.array): The computed number of expected false positives
            for all sub-intervals; axis-0 is for the indexes of lower bounds and
            axis-1 is for the indexes of upper bounds.

    Returns:
        partitions (list): list of lower and upper bounds of set sizes for
            all partitions.
        total_nfps (float): total number of expected false positives from all
            partitions.
        cost (numpy.array): a N x p-1 matrix of the computed optimal NFPs for
            all sub-problems given upper bound set size and number of partitions.
    """

    if num_part < 2:
        raise ValueError("num_part cannot be less than 2")
    if num_part > len(sizes):
        raise ValueError("num_part cannot be greater than the domain size of "
                "all set sizes")

    # If number of partitions is 2, then simply find the upper bound
    # of the first partition.
    if num_part == 2:
        total_nfps, u = min((nfps[0, u1]+nfps[u1+1, len(sizes)-1], u1)
            for u1 in range(0, len(sizes)-1))
        return [(sizes[0], sizes[u]), (sizes[u+1], sizes[-1]),], \
                total_nfps, None

    # Initialize subproblem total NFPs.
    cost = np.zeros((len(sizes), num_part-2))

    # Note: p is the number of partitions in the subproblem.
    # p2i translates the number of partition into the index in the matrix.
    p2i = lambda p : p - 2

    # Compute p >= 2 until before p = num_part.
    for p in range(2, num_part):
        # Compute best partition for subproblems with increasing
        # max index u, starting from the smallest possible u given the p.
        # The smallest possible u can be considered as the max index that
        # generates p partitions each with only one size.
        for u in range(p-1, len(sizes)):
            if p == 2:
                cost[u, p2i(p)] = min(nfps[0, u1]+nfps[u1+1,u]
                        for u1 in range(u))
            else:
                cost[u, p2i(p)] = min(cost[u1, p2i(p-1)] + nfps[u1+1, u]
                        for u1 in range((p-1)-1, u))
    p = num_part
    # Find the optimal upper bound index of the 2nd right-most partition given
    # the number of partitions (p).
    total_nfps, u = min((cost[u1, p2i(p-1)]+nfps[u1+1, len(sizes)-1], u1)
            for u1 in range((p-1)-1, len(sizes)-1))
    partitions = [(sizes[u+1], sizes[-1]),]
    p -= 1
    # Back track to find the best partitions.
    while p > 1:
        # Find the optimal upper bound index of the 2nd right-most partition
        # givne the number of partitions (p) and upper bound index (u) in this
        # sub-problem.
        _, u1_best = min((cost[u1, p2i(p)]+nfps[u1+1, u], u1)
                for u1 in range((p-1)-1, u))
        partitions.insert(0, (sizes[u1_best+1], sizes[u]))
        u = u1_best
        p -= 1
    partitions.insert(0, (sizes[0], sizes[u]))
    return [partitions, total_nfps, cost]

def optimal_partitions(sizes, counts, num_part):
    """Compute the optimal partitions given a distribution of set sizes.

    Args:
        sizes (numpy.array): The complete domain of set sizes in ascending
            order.
        counts (numpy.array): The frequencies of all set sizes in the same
            order as `sizes`.
        num_part (int): The number of partitions to create.

    Returns:
        list: A list of partitions in the form of `(lower, upper)` tuples,
            where `lower` and `upper` are lower and upper bound (inclusive)
            set sizes of each partition.
    """
    if num_part < 2:
        return [(sizes[0], sizes[-1])]
    if num_part >= len(sizes):
        partitions = [(x, x) for x in sizes]
        return partitions
    nfps = _compute_nfps_real(counts, sizes)
    partitions, _, _ = _compute_best_partitions(num_part, sizes, nfps)
    return partitions


'''
This module implements the b-bit MinHash.
http://research.microsoft.com/pubs/120078/wfc0398-liPS.pdf

b-bit MinHash reduces storage space by storing only the b lowest bits
of each minimum hashed value, without significant loss of accuracy.
'''

import struct
import numpy as np

class bBitMinHash(object):
    '''
    The b-bit MinHash object
    '''

    __slots__ = ('seed', 'b', 'r', 'hashvalues')

    # seed as int64
    # b as uint8
    # r as float64
    # num_perm as int32
    _serial_fmt_params = '<qBdi'
    # each block as uint64
    _serial_fmt_block = 'Q'

    def __init__(self, minhash, b=1, r=0.0):
        '''
        Initialize a b-bit MinHash given an existing full MinHash
        object and parameter b - the number of bits to store for
        each minimum hashed values in the MinHash object.
        '''
        b = int(b)
        r = float(r)
        if b > 32 or b < 0:
            raise ValueError("b must be an integer in [0, 32]")
        if r > 1.0:
            raise ValueError("r must be a float in [0.0, 1.0]")
        bmask = (1 << b) - 1
        self.hashvalues = np.bitwise_and(minhash.hashvalues, bmask)\
                .astype(np.uint32)
        self.seed = minhash.seed
        self.b = b
        self.r = r

    def __eq__(self, other):
        '''
        Check for full equality of two b-bit MinHash objects.
        '''
        return type(self) is type(other) and \
            self.seed == other.seed and \
            self.b == other.b and \
            self.r == other.r and \
            np.array_equal(self.hashvalues, other.hashvalues)

    def jaccard(self, other):
        '''
        Estimate the Jaccard similarity (resemblance) between this b-bit
        MinHash and the other.
        '''
        if self.b != other.b:
            raise ValueError("Cannot compare two b-bit MinHashes with different\
                    b values")
        if self.seed != other.seed:
            raise ValueError("Cannot compare two b-bit MinHashes with different\
                    set of permutations")
        intersection = np.count_nonzero(self.hashvalues==other.hashvalues)
        raw_est = float(intersection) / float(self.hashvalues.size)
        a1 = self._calc_a(self.r, self.b)
        a2 = self._calc_a(other.r, other.b)
        c1, c2 = self._calc_c(a1, a2, self.r, other.r)
        return (raw_est - c1) / (1 - c2)

    def bytesize(self):
        '''
        Get the serialized size of this b-bit MinHash in number of bytes.
        '''
        return self._bytesize()[-1]

    def __getstate__(self):
        '''
        This function is called when pickling the b-bit MinHash object.
        Returns a bytearray which will then be pickled.
        '''
        slot_size, n, num_blocks, total = self._bytesize()
        buffer = bytearray(total)
        blocks = np.zeros((num_blocks,), dtype=np.uint64)
        for i in range(num_blocks):
            # Obtain the current segment of n hashed values
            start = i * n
            hvs = self.hashvalues[start:start+n]
            # Store the n b-bit hashed values in the current block
            for j, hv in enumerate(hvs):
                blocks[i] |= np.uint64(hv << (n - 1 - j) * slot_size)
        fmt = self._serial_fmt_params + \
                "%d%s" % (num_blocks, self._serial_fmt_block)
        struct.pack_into(fmt, buffer, 0, self.seed, self.b, self.r, \
                self.hashvalues.size, *blocks)
        return buffer

    def __setstate__(self, buf):
        '''
        This function is called when unpickling the b-bit MinHash object.
        Initialize the object with data in the buffer.
        '''
        try:
            self.seed, self.b, self.r, num_perm = \
                    struct.unpack_from(self._serial_fmt_params, buf, 0)
        except TypeError:
            self.seed, self.b, self.r, num_perm = \
                    struct.unpack_from(self._serial_fmt_params, buffer(buf), 0)
        offset = struct.calcsize(self._serial_fmt_params)
        self.hashvalues = np.zeros((num_perm,), dtype=np.uint32)
        # Reconstruct the hash values
        slot_size, n, num_blocks, total = self._bytesize()
        fmt = "%d%s" % (num_blocks, self._serial_fmt_block)
        try:
            blocks = struct.unpack_from(fmt, buf, offset)
        except TypeError:
            blocks = struct.unpack_from(fmt, buffer(buf), offset)
        mask = (1 << slot_size) - 1
        for i in range(num_blocks):
            start = i * n
            for j, _ in enumerate(self.hashvalues[start:start+n]):
                hv = (blocks[i] >> (n - 1 - j) * slot_size) & mask
                self.hashvalues[start+j] = np.uint32(hv)

    def _calc_a(self, r, b):
        '''
        Compute the function A(r, b)
        '''
        if r == 0.0:
            # Find the limit of A(r, b) as r -> 0.
            return 1.0 / (1 << b)
        return r * (1 - r) ** (2 ** b - 1) / (1 - (1 - r) ** (2 * b))

    def _calc_c(self, a1, a2, r1, r2):
        '''
        Compute the functions C1 and C2
        '''
        if r1 == 0.0 and r2 == 0.0:
            # Find the limits of C1 and C2 as r1 -> 0 and r2 -> 0
            # Since the b-value must be the same and r1 = r2,
            # we have A1(r1, b1) = A2(r2, b2) = A,
            # then the limits for both C1 and C2 are A.
            return a1, a2
        div = 1 / (r1 + r2)
        c1 = (a1 * r2 + a2 * r1) * div
        c2 = (a1 * r1 + a2 * r2) * div
        return c1, c2

    def _find_slot_size(self, b):
        if b == 1:
            return 1
        if b == 2:
            return 2
        if b <= 4:
            return 4
        if b <= 8:
            return 8
        if b <= 16:
            return 16
        if b <= 32:
            return 32
        raise ValueError("Incorrect value of b")

    def _bytesize(self):
        block_size = struct.calcsize(self._serial_fmt_block)
        # Get the size of the slot for storing one b-bit hashed value
        slot_size = self._find_slot_size(self.b)
        # Get the number of slots to be stored in each block
        num_slots_per_block = int(block_size * 8 / slot_size)
        # Get the number of blocks required
        num_blocks = int(np.ceil(float(self.hashvalues.size) /\
                num_slots_per_block))
        # Get the total serialized size
        total = struct.calcsize(self._serial_fmt_params + \
                "%d%s" % (num_blocks, self._serial_fmt_block))
        return slot_size, num_slots_per_block, num_blocks, total


import struct
import numpy as np

from datasketch import MinHash

class LeanMinHash(MinHash):
    '''Lean MinHash is MinHash with a smaller memory footprint
    and faster deserialization, but with its internal state frozen
    -- no `update()`.

    Lean MinHash inherits all methods from :class:`datasketch.MinHash`.
    It does not store the `permutations` and the `hashfunc` needed for updating.
    If a MinHash does not need further updates, convert it into a lean MinHash
    to save memory.

    Example:
        To create a lean MinHash from an existing MinHash:

        .. code-block:: python

            lean_minhash = LeanMinHash(minhash)

            # You can compute the Jaccard similarity between two lean MinHash
            lean_minhash.jaccard(lean_minhash2)

            # Or between a lean MinHash and a MinHash
            lean_minhash.jaccard(minhash2)

        To create a MinHash from a lean MinHash:

        .. code-block:: python

            minhash = MinHash(seed=lean_minhash.seed,
                              hashvalues=lean_minhash.hashvalues)

            # Or if you want to prevent further updates on minhash
            # from affecting the state of lean_minhash
            minhash = MinHash(seed=lean_minhash.seed,
                              hashvalues=lean_minhash.digest())

    Note:
        Lean MinHash can also be used in :class:`datasketch.MinHashLSH`,
        :class:`datasketch.MinHashLSHForest`, and :class:`datasketch.MinHashLSHEnsemble`.

    Args:
        minhash: The :class:`datasketch.MinHash` object used to initialize the LeanMinHash.
    '''

    __slots__ = ('seed', 'hashvalues')

    def _initialize_slots(self, seed, hashvalues):
        '''Initialize the slots of the LeanMinHash.

        Args:
            seed (int): The random seed controls the set of random
                permutation functions generated for this LeanMinHash.
            hashvalues: The hash values is the internal state of the LeanMinHash.
        '''
        self.seed = seed
        self.hashvalues = self._parse_hashvalues(hashvalues)

    def __init__(self, minhash):
        self._initialize_slots(minhash.seed, minhash.hashvalues)

    def update(self, b):
        '''This method is not available on a LeanMinHash.
        Calling it raises a TypeError.
        '''
        raise TypeError("Cannot update a LeanMinHash")

    def copy(self):
        lmh = object.__new__(LeanMinHash)
        lmh._initialize_slots(*self.__slots__)
        return lmh

    def bytesize(self, byteorder='@'):
        '''Compute the byte size after serialization.

        Args:
            byteorder (str, optional): This is byte order of the serialized data. Use one
                of the `byte order characters
                <https://docs.python.org/3/library/struct.html#byte-order-size-and-alignment>`_:
                ``@``, ``=``, ``<``, ``>``, and ``!``.
                Default is ``@`` -- the native order.

        Returns:
            int: Size in number of bytes after serialization.
        '''
        # Use 8 bytes to store the seed integer
        seed_size = struct.calcsize(byteorder+'q')
        # Use 4 bytes to store the number of hash values
        length_size = struct.calcsize(byteorder+'i')
        # Use 4 bytes to store each hash value as we are using the lower 32 bit
        hashvalue_size = struct.calcsize(byteorder+'I')
        return seed_size + length_size + len(self) * hashvalue_size

    def serialize(self, buf, byteorder='@'):
        '''
        Serialize this lean MinHash and store the result in an allocated buffer.

        Args:
            buf (buffer): `buf` must implement the `buffer`_ interface.
                One such example is the built-in `bytearray`_ class.
            byteorder (str, optional): This is byte order of the serialized data. Use one
                of the `byte order characters
                <https://docs.python.org/3/library/struct.html#byte-order-size-and-alignment>`_:
                ``@``, ``=``, ``<``, ``>``, and ``!``.
                Default is ``@`` -- the native order.

        This is preferred over using `pickle`_ if the serialized lean MinHash needs
        to be used by another program in a different programming language.

        The serialization schema:
            1. The first 8 bytes is the seed integer
            2. The next 4 bytes is the number of hash values
            3. The rest is the serialized hash values, each uses 4 bytes

        Example:
            To serialize a single lean MinHash into a `bytearray`_ buffer.

            .. code-block:: python

                buf = bytearray(lean_minhash.bytesize())
                lean_minhash.serialize(buf)

            To serialize multiple lean MinHash into a `bytearray`_ buffer.

            .. code-block:: python

                # assuming lean_minhashs is a list of LeanMinHash with the same size
                size = lean_minhashs[0].bytesize()
                buf = bytearray(size*len(lean_minhashs))
                for i, lean_minhash in enumerate(lean_minhashs):
                    lean_minhash.serialize(buf[i*size:])

        .. _`buffer`: https://docs.python.org/3/c-api/buffer.html
        .. _`bytearray`: https://docs.python.org/3.6/library/functions.html#bytearray
        .. _`byteorder`: https://docs.python.org/3/library/struct.html
        '''
        if len(buf) < self.bytesize():
            raise ValueError("The buffer does not have enough space\
                    for holding this MinHash.")
        fmt = "%sqi%dI" % (byteorder, len(self))
        struct.pack_into(fmt, buf, 0,
                self.seed, len(self), *self.hashvalues)

    @classmethod
    def deserialize(cls, buf, byteorder='@'):
        '''
        Deserialize a lean MinHash from a buffer.

        Args:
            buf (buffer): `buf` must implement the `buffer`_ interface.
                One such example is the built-in `bytearray`_ class.
            byteorder (str. optional): This is byte order of the serialized data. Use one
                of the `byte order characters
                <https://docs.python.org/3/library/struct.html#byte-order-size-and-alignment>`_:
                ``@``, ``=``, ``<``, ``>``, and ``!``.
                Default is ``@`` -- the native order.

        Return:
            datasketch.LeanMinHash: The deserialized lean MinHash

        Example:
            To deserialize a lean MinHash from a buffer.

            .. code-block:: python

                lean_minhash = LeanMinHash.deserialize(buf)
        '''
        fmt_seed_size = "%sqi" % byteorder
        fmt_hash = byteorder + "%dI"
        try:
            seed, num_perm = struct.unpack_from(fmt_seed_size, buf, 0)
        except TypeError:
            seed, num_perm = struct.unpack_from(fmt_seed_size, buffer(buf), 0)
        offset = struct.calcsize(fmt_seed_size)
        try:
            hashvalues = struct.unpack_from(fmt_hash % num_perm, buf, offset)
        except TypeError:
            hashvalues = struct.unpack_from(fmt_hash % num_perm, buffer(buf), offset)
        lmh = object.__new__(LeanMinHash)
        lmh._initialize_slots(seed, hashvalues)
        return lmh

    def __getstate__(self):
        buf = bytearray(self.bytesize())
        fmt = "qi%dI" % len(self)
        struct.pack_into(fmt, buf, 0,
                self.seed, len(self), *self.hashvalues)
        return buf

    def __setstate__(self, buf):
        try:
            seed, num_perm = struct.unpack_from('qi', buf, 0)
        except TypeError:
            seed, num_perm = struct.unpack_from('qi', buffer(buf), 0)
        offset = struct.calcsize('qi')
        try:
            hashvalues = struct.unpack_from('%dI' % num_perm, buf, offset)
        except TypeError:
            hashvalues = struct.unpack_from('%dI' % num_perm, buffer(buf), offset)
        self._initialize_slots(seed, hashvalues)

    def __hash__(self):
        return hash((self.seed, tuple(self.hashvalues)))

    @classmethod
    def union(cls, *lmhs):
        if len(lmhs) < 2:
            raise ValueError("Cannot union less than 2 MinHash")
        num_perm = len(lmhs[0])
        seed = lmhs[0].seed
        if any((seed != m.seed or num_perm != len(m)) for m in lmhs):
            raise ValueError("The unioning MinHash must have the\
                    same seed, number of permutation functions.")
        hashvalues = np.minimum.reduce([m.hashvalues for m in lmhs])

        lmh = object.__new__(LeanMinHash)
        lmh._initialize_slots(seed, hashvalues)
        return lmh


import random, copy, struct
import warnings
import numpy as np

from datasketch.hashfunc import sha1_hash32

# The size of a hash value in number of bytes
hashvalue_byte_size = len(bytes(np.int64(42).data))

# http://en.wikipedia.org/wiki/Mersenne_prime
_mersenne_prime = (1 << 61) - 1
_max_hash = (1 << 32) - 1
_hash_range = (1 << 32)

class MinHash(object):
    '''MinHash is a probabilistic data structure for computing
    `Jaccard similarity`_ between sets.

    Args:
        num_perm (int, optional): Number of random permutation functions.
            It will be ignored if `hashvalues` is not None.
        seed (int, optional): The random seed controls the set of random
            permutation functions generated for this MinHash.
        hashfunc (optional): The hash function used by this MinHash.
            It takes the input passed to the `update` method and
            returns an integer that can be encoded with 32 bits.
            The default hash function is based on SHA1 from hashlib_.
        hashobj (**deprecated**): This argument is deprecated since version
            1.4.0. It is a no-op and has been replaced by `hashfunc`.
        hashvalues (`numpy.array` or `list`, optional): The hash values is
            the internal state of the MinHash. It can be specified for faster
            initialization using the existing state from another MinHash.
        permutations (optional): The permutation function parameters. This argument
            can be specified for faster initialization using the existing
            state from another MinHash.

    Note:
        To save memory usage, consider using :class:`datasketch.LeanMinHash`.

    Note:
        Since version 1.1.1, MinHash will only support serialization using
        `pickle`_. ``serialize`` and ``deserialize`` methods are removed,
        and are supported in :class:`datasketch.LeanMinHash` instead.
        MinHash serialized before version 1.1.1 cannot be deserialized properly
        in newer versions (`need to migrate? <https://github.com/ekzhu/datasketch/issues/18>`_).

    Note:
        Since version 1.1.3, MinHash uses Numpy's random number generator
        instead of Python's built-in random package. This change makes the
        hash values consistent across different Python versions.
        The side-effect is that now MinHash created before version 1.1.3 won't
        work (i.e., ``jaccard``, ``merge`` and ``union``)
        with those created after.

    .. _`Jaccard similarity`: https://en.wikipedia.org/wiki/Jaccard_index
    .. _hashlib: https://docs.python.org/3.5/library/hashlib.html
    .. _`pickle`: https://docs.python.org/3/library/pickle.html
    '''

    def __init__(self, num_perm=128, seed=1,
            hashfunc=sha1_hash32,
            hashobj=None, # Deprecated.
            hashvalues=None, permutations=None):
        if hashvalues is not None:
            num_perm = len(hashvalues)
        if num_perm > _hash_range:
            # Because 1) we don't want the size to be too large, and
            # 2) we are using 4 bytes to store the size value
            raise ValueError("Cannot have more than %d number of\
                    permutation functions" % _hash_range)
        self.seed = seed
        # Check the hash function.
        if not callable(hashfunc):
            raise ValueError("The hashfunc must be a callable.")
        self.hashfunc = hashfunc
        # Check for use of hashobj and issue warning.
        if hashobj is not None:
            warnings.warn("hashobj is deprecated, use hashfunc instead.",
                    DeprecationWarning)
        # Initialize hash values
        if hashvalues is not None:
            self.hashvalues = self._parse_hashvalues(hashvalues)
        else:
            self.hashvalues = self._init_hashvalues(num_perm)
        # Initalize permutation function parameters
        if permutations is not None:
            self.permutations = permutations
        else:
            generator = np.random.RandomState(self.seed)
            # Create parameters for a random bijective permutation function
            # that maps a 32-bit hash value to another 32-bit hash value.
            # http://en.wikipedia.org/wiki/Universal_hashing
            self.permutations = np.array([(generator.randint(1, _mersenne_prime, dtype=np.uint64),
                                           generator.randint(0, _mersenne_prime, dtype=np.uint64))
                                          for _ in range(num_perm)], dtype=np.uint64).T
        if len(self) != len(self.permutations[0]):
            raise ValueError("Numbers of hash values and permutations mismatch")

    def _init_hashvalues(self, num_perm):
        return np.ones(num_perm, dtype=np.uint64)*_max_hash

    def _parse_hashvalues(self, hashvalues):
        return np.array(hashvalues, dtype=np.uint64)

    def update(self, b):
        '''Update this MinHash with a new value.
        The value will be hashed using the hash function specified by
        the `hashfunc` argument in the constructor.

        Args:
            b: The value to be hashed using the hash function specified.

        Example:
            To update with a new string value (using the default SHA1 hash
            function, which requires bytes as input):

            .. code-block:: python

                minhash = Minhash()
                minhash.update("new value".encode('utf-8'))

            We can also use a different hash function, for example, `pyfarmhash`:

            .. code-block:: python

                import farmhash
                def _hash_32(b):
                    return farmhash.hash32(b)
                minhash = MinHash(hashfunc=_hash_32)
                minhash.update("new value")
        '''
        hv = self.hashfunc(b)
        a, b = self.permutations
        phv = np.bitwise_and((a * hv + b) % _mersenne_prime, np.uint64(_max_hash))
        self.hashvalues = np.minimum(phv, self.hashvalues)

    def jaccard(self, other):
        '''Estimate the `Jaccard similarity`_ (resemblance) between the sets
        represented by this MinHash and the other.

        Args:
            other (datasketch.MinHash): The other MinHash.

        Returns:
            float: The Jaccard similarity, which is between 0.0 and 1.0.
        '''
        if other.seed != self.seed:
            raise ValueError("Cannot compute Jaccard given MinHash with\
                    different seeds")
        if len(self) != len(other):
            raise ValueError("Cannot compute Jaccard given MinHash with\
                    different numbers of permutation functions")
        return np.float(np.count_nonzero(self.hashvalues==other.hashvalues)) /\
                np.float(len(self))

    def count(self):
        '''Estimate the cardinality count based on the technique described in
        `this paper <http://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=365694>`_.

        Returns:
            int: The estimated cardinality of the set represented by this MinHash.
        '''
        k = len(self)
        return np.float(k) / np.sum(self.hashvalues / np.float(_max_hash)) - 1.0

    def merge(self, other):
        '''Merge the other MinHash with this one, making this one the union
        of both.

        Args:
            other (datasketch.MinHash): The other MinHash.
        '''
        if other.seed != self.seed:
            raise ValueError("Cannot merge MinHash with\
                    different seeds")
        if len(self) != len(other):
            raise ValueError("Cannot merge MinHash with\
                    different numbers of permutation functions")
        self.hashvalues = np.minimum(other.hashvalues, self.hashvalues)

    def digest(self):
        '''Export the hash values, which is the internal state of the
        MinHash.

        Returns:
            numpy.array: The hash values which is a Numpy array.
        '''
        return copy.copy(self.hashvalues)

    def is_empty(self):
        '''
        Returns:
            bool: If the current MinHash is empty - at the state of just
                initialized.
        '''
        if np.any(self.hashvalues != _max_hash):
            return False
        return True

    def clear(self):
        '''
        Clear the current state of the MinHash.
        All hash values are reset.
        '''
        self.hashvalues = self._init_hashvalues(len(self))

    def copy(self):
        '''
        :returns: datasketch.MinHash -- A copy of this MinHash by exporting its state.
        '''
        return MinHash(seed=self.seed, hashfunc=self.hashfunc,
                hashvalues=self.digest(),
                permutations=self.permutations)

    def __len__(self):
        '''
        :returns: int -- The number of hash values.
        '''
        return len(self.hashvalues)

    def __eq__(self, other):
        '''
        :returns: bool -- If their seeds and hash values are both equal then two are equivalent.
        '''
        return type(self) is type(other) and \
            self.seed == other.seed and \
            np.array_equal(self.hashvalues, other.hashvalues)

    @classmethod
    def union(cls, *mhs):
        '''Create a MinHash which is the union of the MinHash objects passed as arguments.

        Args:
            *mhs: The MinHash objects to be united. The argument list length is variable,
                but must be at least 2.

        Returns:
            datasketch.MinHash: A new union MinHash.
        '''
        if len(mhs) < 2:
            raise ValueError("Cannot union less than 2 MinHash")
        num_perm = len(mhs[0])
        seed = mhs[0].seed
        if any((seed != m.seed or num_perm != len(m)) for m in mhs):
            raise ValueError("The unioning MinHash must have the\
                    same seed and number of permutation functions")
        hashvalues = np.minimum.reduce([m.hashvalues for m in mhs])
        permutations = mhs[0].permutations
        return cls(num_perm=num_perm, seed=seed, hashvalues=hashvalues,
                permutations=permutations)


from collections import deque, Counter
import numpy as np
from datasketch.lsh import integrate, MinHashLSH
from datasketch.lshensemble_partition import optimal_partitions

def _false_positive_probability(threshold, b, r, xq):
    '''
    Compute the false positive probability given the containment threshold.
    xq is the ratio of x/q.
    '''
    _probability = lambda t : 1 - (1 - (t/(1 + xq - t))**float(r))**float(b)
    if xq >= threshold:
        a, err = integrate(_probability, 0.0, threshold)
        return a
    a, err = integrate(_probability, 0.0, xq)
    return a

def _false_negative_probability(threshold, b, r, xq):
    '''
    Compute the false negative probability given the containment threshold
    '''
    _probability = lambda t : 1 - (1 - (1 - (t/(1 + xq - t))**float(r))**float(b))
    if xq >= 1.0:
        a, err = integrate(_probability, threshold, 1.0)
        return a
    if xq >= threshold:
        a, err = integrate(_probability, threshold, xq)
        return a
    return 0.0

def _optimal_param(threshold, num_perm, max_r, xq, false_positive_weight,
        false_negative_weight):
    '''
    Compute the optimal parameters that minimizes the weighted sum
    of probabilities of false positive and false negative.
    xq is the ratio of x/q.
    '''
    min_error = float("inf")
    opt = (0, 0)
    for b in range(1, num_perm+1):
        for r in range(1, max_r+1):
            if b*r > num_perm:
                continue
            fp = _false_positive_probability(threshold, b, r, xq)
            fn = _false_negative_probability(threshold, b, r, xq)
            error = fp*false_positive_weight + fn*false_negative_weight
            if error < min_error:
                min_error = error
                opt = (b, r)
    return opt

class MinHashLSHEnsemble(object):
    '''
    The :ref:`minhash_lsh_ensemble` index. It supports
    :ref:`containment` queries.
    The implementation is based on
    `E. Zhu et al. <http://www.vldb.org/pvldb/vol9/p1185-zhu.pdf>`_.

    Args:
        threshold (float): The Containment threshold between 0.0 and
            1.0. The initialized LSH Ensemble will be optimized for the threshold by
            minizing the false positive and false negative.
        num_perm (int, optional): The number of permutation functions used
            by the MinHash to be indexed. For weighted MinHash, this
            is the sample size (`sample_size`).
        num_part (int, optional): The number of partitions in LSH Ensemble.
        m (int, optional): The memory usage factor: an LSH Ensemble uses approximately
            `m` times more memory space than a MinHash LSH with the same number of
            sets indexed. The higher the `m` the better the accuracy.
        weights (tuple, optional): Used to adjust the relative importance of
            minizing false positive and false negative when optimizing
            for the Containment threshold. Similar to the `weights` parameter
            in :class:`datasketch.MinHashLSH`.

    Note:
        Using more partitions (`num_part`) leads to better accuracy, at the
        expense of slower query performance.
        This is different from `the paper`_ and the `Go implementation`_, in which
        more partitions leads to better accuracy AND faster query performance,
        due to parallelism.

    Note:
        More information about the parameter `m` can be found in the
        `Go implementation`_
        of LSH Ensemble, in which `m` is named `MaxK`.

    .. _`Go implementation`: https://github.com/ekzhu/lshensemble
    .. _`the paper`: http://www.vldb.org/pvldb/vol9/p1185-zhu.pdf
    '''

    def __init__(self, threshold=0.9, num_perm=128, num_part=16, m=8, weights=(0.5,0.5)):
        if threshold > 1.0 or threshold < 0.0:
            raise ValueError("threshold must be in [0.0, 1.0]")
        if num_perm < 2:
            raise ValueError("Too few permutation functions")
        if num_part < 1:
            raise ValueError("num_part must be at least 1")
        if m < 2 or m > num_perm:
            raise ValueError("m must be in the range of [2, num_perm]")
        if any(w < 0.0 or w > 1.0 for w in weights):
            raise ValueError("Weight must be in [0.0, 1.0]")
        if sum(weights) != 1.0:
            raise ValueError("Weights must sum to 1.0")
        self.threshold = threshold
        self.h = num_perm
        self.m = m
        rs = self._init_optimal_params(weights)
        # Initialize multiple LSH indexes for each partition
        self.indexes = [dict((r, MinHashLSH(num_perm=self.h,
            params=(int(self.h/r), r))) for r in rs)
            for _ in range(0, num_part)]
        self.lowers = [None for _ in self.indexes]
        self.uppers = [None for _ in self.indexes]

    def _init_optimal_params(self, weights):
        false_positive_weight, false_negative_weight = weights
        self.xqs = np.exp(np.linspace(-5, 5, 10))
        self.params = np.array([_optimal_param(self.threshold, self.h, self.m,
                                               xq,
                                               false_positive_weight,
                                               false_negative_weight)
                                for xq in self.xqs], dtype=np.int)
        # Find all unique r
        rs = set()
        for _, r in self.params:
            rs.add(r)
        return rs

    def _get_optimal_param(self, x, q):
        i = np.searchsorted(self.xqs, float(x)/float(q), side='left')
        if i == len(self.params):
            i = i - 1
        return self.params[i]

    def index(self, entries):
        '''
        Index all sets given their keys, MinHashes, and sizes.
        It can be called only once after the index is created.

        Args:
            entries (`iterable` of `tuple`): An iterable of tuples, each must be
                in the form of `(key, minhash, size)`, where `key` is the unique
                identifier of a set, `minhash` is the MinHash of the set,
                and `size` is the size or number of unique items in the set.

        Note:
            `size` must be positive.
        '''
        if not self.is_empty():
            raise ValueError("Cannot call index again on a non-empty index")
        if not isinstance(entries, list):
            queue = deque([])
            for key, minhash, size in entries:
                if size <= 0:
                    raise ValueError("Set size must be positive")
                queue.append((key, minhash, size))
            entries = list(queue)
        if len(entries) == 0:
            raise ValueError("entries is empty")
        # Create optimal partitions.
        sizes, counts = np.array(sorted(
            Counter(e[2] for e in entries).most_common())).T
        partitions = optimal_partitions(sizes, counts, len(self.indexes))
        for i, (lower, upper) in enumerate(partitions):
            self.lowers[i], self.uppers[i] = lower, upper
        # Insert into partitions.
        entries.sort(key=lambda e : e[2])
        curr_part = 0
        for key, minhash, size in entries:
            if size > self.uppers[curr_part]:
                curr_part += 1
            for r in self.indexes[curr_part]:
                self.indexes[curr_part][r].insert(key, minhash)

    def query(self, minhash, size):
        '''
        Giving the MinHash and size of the query set, retrieve
        keys that references sets with containment with respect to
        the query set greater than the threshold.

        Args:
            minhash (datasketch.MinHash): The MinHash of the query set.
            size (int): The size (number of unique items) of the query set.

        Returns:
            `iterator` of keys.
        '''
        for i, index in enumerate(self.indexes):
            u = self.uppers[i]
            if u is None:
                continue
            b, r = self._get_optimal_param(u, size)
            for key in index[r]._query_b(minhash, b):
                yield key

    def __contains__(self, key):
        '''
        Args:
            key (hashable): The unique identifier of a set.

        Returns:
            bool: True only if the key exists in the index.
        '''
        return any(any(key in index[r] for r in index)
                   for index in self.indexes)

    def is_empty(self):
        '''
        Returns:
            bool: Check if the index is empty.
        '''
        return all(all(index[r].is_empty() for r in index)
                   for index in self.indexes)

if __name__ == "__main__":
    import numpy as np
    xqs = np.exp(np.linspace(-5, 5, 10))
    threshold = 0.5
    max_r = 8
    num_perm = 256
    false_negative_weight, false_positive_weight = 0.5, 0.5
    for xq in xqs:
        b, r = _optimal_param(threshold, num_perm, max_r, xq,
                false_positive_weight, false_negative_weight)
        print("threshold: %.2f, xq: %.3f, b: %d, r: %d" % (threshold, xq, b, r))


import collections
import copy
import numpy as np

class WeightedMinHash(object):
    '''New weighted MinHash is generated by
    :class:`datasketch.WeightedMinHashGenerator`.
    You can also initialize weighted MinHash by using the state
    from an existing one.

    Args:
        seed (int): The random seed used to generate this weighted
            MinHash.
        hashvalues: The internal state of this weighted MinHash.
    '''

    def __init__(self, seed, hashvalues):
        self.seed = seed
        self.hashvalues = hashvalues

    def jaccard(self, other):
        '''Estimate the `weighted Jaccard similarity`_ between the
        multi-sets represented by this weighted MinHash and the other.

        Args:
            other (datasketch.WeightedMinHash): The other weighted MinHash.

        Returns:
            float: The weighted Jaccard similarity between 0.0 and 1.0.

        .. _`weighted Jaccard similarity`: http://mathoverflow.net/questions/123339/weighted-jaccard-similarity
        '''
        if other.seed != self.seed:
            raise ValueError("Cannot compute Jaccard given WeightedMinHash objects with\
                    different seeds")
        if len(self) != len(other):
            raise ValueError("Cannot compute Jaccard given WeightedMinHash objects with\
                    different numbers of hash values")
        # Check how many pairs of (k, t) hashvalues are equal
        intersection = 0
        for this, that in zip(self.hashvalues, other.hashvalues):
            if np.array_equal(this, that):
                intersection += 1
        return float(intersection) / float(len(self))

    def digest(self):
        '''Export the hash values, which is the internal state of the
        weighted MinHash.

        Returns:
            numpy.array: The hash values which is a Numpy array.
        '''
        return copy.copy(self.hashvalues)

    def copy(self):
        '''
        Returns:
            datasketch.WeightedMinHash: A copy of this weighted MinHash by exporting
            its state.
        '''
        return WeightedMinHash(self.seed, self.digest())

    def __len__(self):
        '''
        Returns:
            int: The number of hash values.
        '''
        return len(self.hashvalues)

    def __eq__(self, other):
        '''
        Returns:
            bool: If their seeds and hash values are both equal then two
            are equivalent.
        '''
        return type(self) is type(other) and \
            self.seed == other.seed and \
            np.array_equal(self.hashvalues, other.hashvalues)

class WeightedMinHashGenerator(object):
    '''The weighted MinHash generator is used for creating
    new :class:`datasketch.WeightedMinHash` objects.

    This weighted MinHash implementation is based on Sergey Ioffe's paper,
    `Improved Consistent Sampling, Weighted Minhash and L1 Sketching
    <http://static.googleusercontent.com/media/research.google.com/en//pubs/archive/36928.pdf>`_

    Args:
        dim (int): The number of dimensions of the input Jaccard vectors.
        sample_size (int, optional): The number of samples to use for creating
            weighted MinHash.
        seed (int): The random seed to use for generating permutation functions.

    '''

    def __init__(self, dim, sample_size=128, seed=1):
        self.dim = dim
        self.sample_size = sample_size
        self.seed = seed
        generator = np.random.RandomState(seed=seed)
        self.rs = generator.gamma(2, 1, (sample_size, dim)).astype(np.float32)
        self.ln_cs = np.log(generator.gamma(2, 1, (sample_size, dim))).astype(np.float32)
        self.betas = generator.uniform(0, 1, (sample_size, dim)).astype(np.float32)

    def minhash(self, v):
        '''Create a new weighted MinHash given a weighted Jaccard vector.
        Each dimension is an integer
        frequency of the corresponding element in the multi-set represented
        by the vector.

        Args:
            v (numpy.array): The Jaccard vector.
        '''
        if not isinstance(v, collections.Iterable):
            raise TypeError("Input vector must be an iterable")
        if not len(v) == self.dim:
            raise ValueError("Input dimension mismatch, expecting %d" % self.dim)
        if not isinstance(v, np.ndarray):
            v = np.array(v, dtype=np.float32)
        elif v.dtype != np.float32:
            v = v.astype(np.float32)
        hashvalues = np.zeros((self.sample_size, 2), dtype=np.int)
        vzeros = (v == 0)
        if vzeros.all():
            raise ValueError("Input is all zeros")
        v[vzeros] = np.nan
        vlog = np.log(v)
        for i in range(self.sample_size):
            t = np.floor((vlog / self.rs[i]) + self.betas[i])
            ln_y = (t - self.betas[i]) * self.rs[i]
            ln_a = self.ln_cs[i] - ln_y - self.rs[i]
            k = np.nanargmin(ln_a)
            hashvalues[i][0], hashvalues[i][1] = k, int(t[k])
        return WeightedMinHash(self.seed, hashvalues)


import pickle
from datasketch.storage import (
    ordered_storage, unordered_storage, _random_name)

_integration_precision = 0.001
def _integration(f, a, b):
    p = _integration_precision
    area = 0.0
    x = a
    while x < b:
        area += f(x+0.5*p)*p
        x += p
    return area, None

try:
    from scipy.integrate import quad as integrate
except ImportError:
    # For when no scipy installed
    integrate = _integration

def _false_positive_probability(threshold, b, r):
    _probability = lambda s : 1 - (1 - s**float(r))**float(b)
    a, err = integrate(_probability, 0.0, threshold)
    return a

def _false_negative_probability(threshold, b, r):
    _probability = lambda s : 1 - (1 - (1 - s**float(r))**float(b))
    a, err = integrate(_probability, threshold, 1.0)
    return a

def _optimal_param(threshold, num_perm, false_positive_weight,
        false_negative_weight):
    '''
    Compute the optimal `MinHashLSH` parameter that minimizes the weighted sum
    of probabilities of false positive and false negative.
    '''
    min_error = float("inf")
    opt = (0, 0)
    for b in range(1, num_perm+1):
        max_r = int(num_perm / b)
        for r in range(1, max_r+1):
            fp = _false_positive_probability(threshold, b, r)
            fn = _false_negative_probability(threshold, b, r)
            error = fp*false_positive_weight + fn*false_negative_weight
            if error < min_error:
                min_error = error
                opt = (b, r)
    return opt

class MinHashLSH(object):
    '''
    The :ref:`minhash_lsh` index.
    It supports query with `Jaccard similarity`_ threshold.
    Reference: `Chapter 3, Mining of Massive Datasets
    <http://www.mmds.org/>`_.

    Args:
        threshold (float): The Jaccard similarity threshold between 0.0 and
            1.0. The initialized MinHash LSH will be optimized for the threshold by
            minizing the false positive and false negative.
        num_perm (int, optional): The number of permutation functions used
            by the MinHash to be indexed. For weighted MinHash, this
            is the sample size (`sample_size`).
        weights (tuple, optional): Used to adjust the relative importance of
            minimizing false positive and false negative when optimizing
            for the Jaccard similarity threshold.
            `weights` is a tuple in the format of
            :code:`(false_positive_weight, false_negative_weight)`.
        params (tuple, optional): The LSH parameters (i.e., number of bands and size
            of each bands). This is used to bypass the parameter optimization
            step in the constructor. `threshold` and `weights` will be ignored
            if this is given.
        storage_config (dict, optional): Type of storage service to use for storing
            hashtables and keys.
            `basename` is an optional property whose value will be used as the prefix to
            stored keys. If this is not set, a random string will be generated instead. If you
            set this, you will be responsible for ensuring there are no key collisions.
        prepickle (bool, optional): If True, all keys are pickled to bytes before
            insertion. If None, a default value is chosen based on the
            `storage_config`.

    Note:
        `weights` must sum to 1.0, and the format is
        (false positive weight, false negative weight).
        For example, if minimizing false negative (or maintaining high recall) is more
        important, assign more weight toward false negative: weights=(0.4, 0.6).
        Try to live with a small difference between weights (i.e. < 0.5).
    '''

    def __init__(self, threshold=0.9, num_perm=128, weights=(0.5, 0.5),
                 params=None, storage_config=None, prepickle=None):
        storage_config = {'type': 'dict'} if not storage_config else storage_config
        self._buffer_size = 50000
        if threshold > 1.0 or threshold < 0.0:
            raise ValueError("threshold must be in [0.0, 1.0]")
        if num_perm < 2:
            raise ValueError("Too few permutation functions")
        if any(w < 0.0 or w > 1.0 for w in weights):
            raise ValueError("Weight must be in [0.0, 1.0]")
        if sum(weights) != 1.0:
            raise ValueError("Weights must sum to 1.0")
        self.h = num_perm
        if params is not None:
            self.b, self.r = params
            if self.b * self.r > num_perm:
                raise ValueError("The product of b and r must be less than num_perm")
        else:
            false_positive_weight, false_negative_weight = weights
            self.b, self.r = _optimal_param(threshold, num_perm,
                    false_positive_weight, false_negative_weight)

        self.prepickle = storage_config['type'] == 'redis' if prepickle is None else prepickle

        basename = storage_config.get('basename', _random_name(11))
        self.hashtables = [
            unordered_storage(storage_config, name=b''.join([basename, b'_bucket_', bytes([i])]))
            for i in range(self.b)]
        self.hashranges = [(i*self.r, (i+1)*self.r) for i in range(self.b)]
        self.keys = ordered_storage(storage_config, name=b''.join([basename, b'_keys']))

    @property
    def buffer_size(self):
        return self._buffer_size

    @buffer_size.setter
    def buffer_size(self, value):
        self.keys.buffer_size = value
        for t in self.hashtables:
            t.buffer_size = value
        self._buffer_size = value

    def insert(self, key, minhash, check_duplication=True):
        '''
        Insert a key to the index, together
        with a MinHash (or weighted MinHash) of the set referenced by
        the key.

        :param str key: The identifier of the set.
        :param datasketch.MinHash minhash: The MinHash of the set.
        :param bool check_duplication: To avoid duplicate keys in the storage (`default=True`).
                                       It's recommended to not change the default, but
                                       if you want to avoid the overhead during insert
                                       you can set `check_duplication = False`.
        '''
        self._insert(key, minhash, check_duplication=check_duplication, buffer=False)

    def insertion_session(self, buffer_size=50000):
        '''
        Create a context manager for fast insertion into this index.

        :param int buffer_size: The buffer size for insert_session mode (default=50000).

        Returns:
            datasketch.lsh.MinHashLSHInsertionSession
        '''
        return MinHashLSHInsertionSession(self, buffer_size=buffer_size)

    def _insert(self, key, minhash, check_duplication=True, buffer=False):
        if len(minhash) != self.h:
            raise ValueError("Expecting minhash with length %d, got %d"
                    % (self.h, len(minhash)))
        if check_duplication and key in self.keys:
            raise ValueError("The given key already exists")
        Hs = [self._H(minhash.hashvalues[start:end])
              for start, end in self.hashranges]
        if self.prepickle:
            key = pickle.dumps(key)
        self.keys.insert(key, *Hs, buffer=buffer)
        for H, hashtable in zip(Hs, self.hashtables):
            hashtable.insert(H, key, buffer=buffer)

    def query(self, minhash):
        '''
        Giving the MinHash of the query set, retrieve
        the keys that references sets with Jaccard
        similarities greater than the threshold.

        Args:
            minhash (datasketch.MinHash): The MinHash of the query set.

        Returns:
            `list` of unique keys.
        '''
        if len(minhash) != self.h:
            raise ValueError("Expecting minhash with length %d, got %d"
                    % (self.h, len(minhash)))
        candidates = set()
        for (start, end), hashtable in zip(self.hashranges, self.hashtables):
            H = self._H(minhash.hashvalues[start:end])
            for key in hashtable.get(H):
                candidates.add(key)
        if self.prepickle:
            return [pickle.loads(key) for key in candidates]
        else:
            return list(candidates)

    def __contains__(self, key):
        '''
        Args:
            key (hashable): The unique identifier of a set.

        Returns:
            bool: True only if the key exists in the index.
        '''
        if self.prepickle:
            key = pickle.dumps(key)
        return key in self.keys

    def remove(self, key):
        '''
        Remove the key from the index.

        Args:
            key (hashable): The unique identifier of a set.

        '''
        if self.prepickle:
            key = pickle.dumps(key)
        if key not in self.keys:
            raise ValueError("The given key does not exist")
        for H, hashtable in zip(self.keys[key], self.hashtables):
            hashtable.remove_val(H, key)
            if not hashtable.get(H):
                hashtable.remove(H)
        self.keys.remove(key)

    def is_empty(self):
        '''
        Returns:
            bool: Check if the index is empty.
        '''
        return any(t.size() == 0 for t in self.hashtables)

    @staticmethod
    def _H(hs):
        return bytes(hs.byteswap().data)

    def _query_b(self, minhash, b):
        if len(minhash) != self.h:
            raise ValueError("Expecting minhash with length %d, got %d"
                    % (self.h, len(minhash)))
        if b > len(self.hashtables):
            raise ValueError("b must be less or equal to the number of hash tables")
        candidates = set()
        for (start, end), hashtable in zip(self.hashranges[:b], self.hashtables[:b]):
            H = self._H(minhash.hashvalues[start:end])
            if H in hashtable:
                for key in hashtable[H]:
                    candidates.add(key)
        if self.prepickle:
            return {pickle.loads(key) for key in candidates}
        else:
            return candidates

    def get_counts(self):
        '''
        Returns a list of length ``self.b`` with elements representing the
        number of keys stored under each bucket for the given permutation.
        '''
        counts = [
            hashtable.itemcounts() for hashtable in self.hashtables]
        return counts

    def get_subset_counts(self, *keys):
        '''
        Returns the bucket allocation counts (see :func:`~datasketch.MinHashLSH.get_counts` above)
        restricted to the list of keys given.

        Args:
            keys (hashable) : the keys for which to get the bucket allocation
                counts
        '''
        if self.prepickle:
            key_set = [pickle.dumps(key) for key in set(keys)]
        else:
            key_set = list(set(keys))
        hashtables = [unordered_storage({'type': 'dict'}) for _ in
                      range(self.b)]
        Hss = self.keys.getmany(*key_set)
        for key, Hs in zip(key_set, Hss):
            for H, hashtable in zip(Hs, hashtables):
                hashtable.insert(H, key)
        return [hashtable.itemcounts() for hashtable in hashtables]

class MinHashLSHInsertionSession:
    '''Context manager for batch insertion of documents into a MinHashLSH.
    '''

    def __init__(self, lsh, buffer_size):
        self.lsh = lsh
        self.lsh.buffer_size = buffer_size

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self):
        self.lsh.keys.empty_buffer()
        for hashtable in self.lsh.hashtables:
            hashtable.empty_buffer()

    def insert(self, key, minhash, check_duplication=True):
        '''
        Insert a unique key to the index, together
        with a MinHash (or weighted MinHash) of the set referenced by
        the key.

        Args:
            key (hashable): The unique identifier of the set.
            minhash (datasketch.MinHash): The MinHash of the set.
        '''
        self.lsh._insert(key, minhash, check_duplication=check_duplication,
                         buffer=True)