<commit_before># -*- coding: utf-8 -*-

import six
from .exceptions import InvalidQuery, InvalidParameterQuery, QueryError, \
    ScriptFieldsError
from .sort import SortFactory
from .facets import FacetFactory
from .aggs import AggFactory
from .filters import ANDFilter, Filter
from .highlight import HighLighter
from .scriptfields import ScriptFields
from .utils import clean_string, ESRange, EqualityComparableUsingAttributeDictionary

class Suggest(EqualityComparableUsingAttributeDictionary):
    def __init__(self, fields=None):
        self.fields = fields or {}

    def add(self, text, name, field, type='term', size=None, params=None):
        """
        Set the suggester of given type.

        :param text: text
        :param name: name of suggest
        :param field: field to be used
        :param type: type of suggester to add, available types are: completion,
                     phrase, term
        :param size: number of phrases
        :param params: dict of additional parameters to pass to the suggester
        :return: None
        """
        func = None

        if type == 'completion':
            func = self.add_completion
        elif type == 'phrase':
            func = self.add_phrase
        elif type == 'term':
            func = self.add_term
        else:
            raise InvalidQuery('Invalid type')

        func(text=text, name=name, field=field, size=size, params=params)

    def add_completion(self, text, name, field, size=None, params=None):
        """
        Add completion-type suggestion:

        http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html

        :param text: text for searching using autocompletion
        :param name: name for the suggester
        :param field: document's field to be autocompleted
        :param size: (optional) size of the autocomplete list
        :return:
        """
        data = {
            'field': field
        }

        if size:
            data['size'] = size

        if params:
            data.update(params)

        self.fields[name] = {
            'text': text,
            'completion': data,
        }


    def add_term(self, text, name, field, size=None, params=None):
        data = {"field": field}

        if size:
            data["size"] = size

        if params:
            data.update(params)

        self.fields[name] = {"text": text, "term": data}

    def add_phrase(self, text, name, field, size=None, params=None):
        tokens = text.split()
        gram = field + ".bigram"
        # if len(tokens) > 3:
        #     gram = field + ".trigram"

        data = {
            "analyzer": "standard_lower",
            "field": gram,
            "size": 4,
            "real_word_error_likelihood": 0.95,
            "confidence": 2.0,
            "gram_size": 2,
            "direct_generator": [{
                                     "field": field + ".tkl",
                                     "suggest_mode": "always",
                                     "min_word_len": 1
                                 }, {
                                     "field": field + ".reverse",
                                     "suggest_mode": "always",
                                     "min_word_len": 1,
                                     "pre_filter": "reverse",
                                     "post_filter": "reverse"
                                 }]
        }

        if size:
            data["size"] = size

        if params:
            data.update(params)

        self.fields[name] = {"text": text, "phrase": data}

    def is_valid(self):
        return len(self.fields) > 0

    def serialize(self):
        return self.fields

class FieldParameter(EqualityComparableUsingAttributeDictionary):

    def __init__(self, field, query, default_operator="OR", analyzer=None,
                 allow_leading_wildcard=True, lowercase_expanded_terms=True,
                 enable_position_increments=True, fuzzy_prefix_length=0,
                 fuzzy_min_sim=0.5, phrase_slop=0, boost=1.0):
        self.query = query
        self.field = field
        self.default_operator = default_operator
        self.analyzer = analyzer
        self.allow_leading_wildcard = allow_leading_wildcard
        self.lowercase_expanded_terms = lowercase_expanded_terms
        self.enable_position_increments = enable_position_increments
        self.fuzzy_prefix_length = fuzzy_prefix_length
        self.fuzzy_min_sim = fuzzy_min_sim
        self.phrase_slop = phrase_slop
        self.boost = boost

    def serialize(self):
        filters = {}

        if self.default_operator != "OR":
            filters["default_operator"] = self.default_operator
        if self.analyzer:
            filters["analyzer"] = self.analyzer
        if not self.allow_leading_wildcard:
            filters["allow_leading_wildcard"] = self.allow_leading_wildcard
        if not self.lowercase_expanded_terms:
            filters["lowercase_expanded_terms"] = self.lowercase_expanded_terms
        if not self.enable_position_increments:
            filters["enable_position_increments"] = self.enable_position_increments
        if self.fuzzy_prefix_length:
            filters["fuzzy_prefix_length"] = self.fuzzy_prefix_length
        if self.fuzzy_min_sim != 0.5:
            filters["fuzzy_min_sim"] = self.fuzzy_min_sim
        if self.phrase_slop:
            filters["phrase_slop"] = self.phrase_slop

        if self.boost != 1.0:
            filters["boost"] = self.boost
        if filters:
            filters["query"] = self.query
        else:
            filters = self.query
        return self.field, filters


class Search(EqualityComparableUsingAttributeDictionary):
    """A search to be performed.

    This contains a query, and has additional parameters which are used to
    control how the search works, what it should return, etc.

    The rescore parameter takes a Search object created from a RescoreQuery.

    Example:

    q = QueryStringQuery('elasticsearch')
    s = Search(q, fields=['title', 'author'], start=100, size=50)
    results = conn.search(s)
    """

    def __init__(self, query=None, filter=None, fields=None, start=None,
                 size=None, highlight=None, sort=None, explain=False, facet=None, agg=None, rescore=None,
                 window_size=None, version=None, track_scores=None, script_fields=None, index_boost=None,
                 min_score=None, stats=None, bulk_read=None, partial_fields=None, _source=None, timeout=None):
        """
        fields: if is [], the _source is not returned
        """
        if not index_boost: index_boost = {}
        self.query = query
        self.filter = filter
        self.fields = fields
        self.start = start
        self.size = size
        self._highlight = highlight
        self.sort = sort or SortFactory()
        self.explain = explain
        self.facet = facet or FacetFactory()
        self.agg = agg or AggFactory()
        self.rescore = rescore
        self.window_size = window_size
        self.version = version
        self.track_scores = track_scores
        self._script_fields = script_fields
        self.index_boost = index_boost
        self.min_score = min_score
        self.stats = stats
        self.bulk_read = bulk_read
        self.partial_fields = partial_fields
        self._source = _source
        self.timeout = timeout

    def get_facet_factory(self):
        """
        Returns the facet factory
        """
        return self.facet

    def get_agg_factory(self):
        """
        Returns the agg factory
        """
        return self.agg

    def serialize(self):
        """Serialize the search to a structure as passed for a search body."""
        res = {}
        if self.query:
            if isinstance(self.query, dict):
                res["query"] = self.query
            elif isinstance(self.query, Query):
                res["query"] = self.query.serialize()
            else:
                raise InvalidQuery("Invalid query")
        if self.filter:
            res['filter'] = self.filter.serialize()
        if self.facet.facets:
            res['facets'] = self.facet.serialize()
        if self.agg.aggs:
            res['aggs'] = self.agg.serialize()
        if self.rescore:
            res['rescore'] = self.rescore.serialize()
        if self.window_size:
            res['window_size'] = self.window_size
        if self.fields is not None: #Deal properly with self.fields = []
            res['fields'] = self.fields
        if self.size is not None:
            res['size'] = self.size
        if self.start is not None:
            res['from'] = self.start
        if self._highlight:
            res['highlight'] = self._highlight.serialize()
        if self.sort:
            if isinstance(self.sort, SortFactory):
                sort = self.sort.serialize()
                if sort:
                    res['sort'] = sort
            else:
                res['sort'] = self.sort
        if self.explain:
            res['explain'] = self.explain
        if self.version:
            res['version'] = self.version
        if self.track_scores:
            res['track_scores'] = self.track_scores
        if self._script_fields:
            if isinstance(self.script_fields, ScriptFields):
                res['script_fields'] = self.script_fields.serialize()
            elif isinstance(self.script_fields, dict):
                res['script_fields'] = self.script_fields
            else:
                raise ScriptFieldsError("Parameter script_fields should of type ScriptFields or dict")
        if self.index_boost:
            res['indices_boost'] = self.index_boost
        if self.min_score:
            res['min_score'] = self.min_score
        if self.stats:
            res['stats'] = self.stats
        if self.partial_fields:
            res['partial_fields'] = self.partial_fields
        if self._source:
            res['_source'] = self._source
        if self.timeout:
            res['timeout'] = self.timeout
        return res

    @property
    def highlight(self):
        if self._highlight is None:
            self._highlight = HighLighter("<b>", "</b>")
        return self._highlight

    @property
    def script_fields(self):
        if self._script_fields is None:
            self._script_fields = ScriptFields()
        return self._script_fields

    def add_highlight(self, field, fragment_size=None,
                      number_of_fragments=None, fragment_offset=None, type=None):
        """Add a highlight field.

        The Search object will be returned, so calls to this can be chained.

        """
        if self._highlight is None:
            self._highlight = HighLighter("<b>", "</b>")
        self._highlight.add_field(field, fragment_size, number_of_fragments, fragment_offset, type=type)
        return self

    def add_index_boost(self, index, boost):
        """Add a boost on an index.

        The Search object will be returned, so calls to this can be chained.

        """
        if boost is None:
            if index in self.index_boost:
                del(self.index_boost[index])
        else:
            self.index_boost[index] = boost
        return self

    def __repr__(self):
        return str(self.serialize())


class Query(EqualityComparableUsingAttributeDictionary):
    """Base class for all queries."""

    def __init__(self, *args, **kwargs):
        if len(args) > 0 or len(kwargs) > 0:
            raise RuntimeWarning("No all parameters are processed by derivated query object")

    def search(self, **kwargs):
        """Return this query wrapped in a Search object.

        Any keyword arguments supplied to this call will be passed to the
        Search object.
        """
        return Search(query=self, **kwargs)

    def serialize(self):
        """Serialize the query to a structure using the query DSL."""
        return {self._internal_name: self._serialize()}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError


class BoolQuery(Query):
    """A boolean combination of other queries.

    BoolQuery maps to Lucene **BooleanQuery**. It is built using one or more
    boolean clauses, each clause with a typed occurrence.  The occurrence types
    are:

    ================  ========================================================
     Occur             Description
    ================  ========================================================
    **must**          The clause (query) must appear in matching documents.
    **should**        The clause (query) should appear in the matching
                      document. A boolean query with no **must** clauses, one
                      or more **should** clauses must match a document. The
                      minimum number of should clauses to match can be set
                      using **minimum_number_should_match** parameter.
    **must_not**      The clause (query) must not appear in the matching
                      documents. Note that it is not possible to search on
                      documents that only consists of a **must_not** clauses.
    ================  ========================================================

    The bool query also supports **disable_coord** parameter (defaults to
    **false**).
    """

    _internal_name = "bool"

    def __init__(self, must=None, must_not=None, should=None, boost=None,
                 minimum_number_should_match=1, disable_coord=None, **kwargs):
        super(BoolQuery, self).__init__(**kwargs)
        self._must = []
        self._must_not = []
        self._should = []
        self.boost = boost
        self.minimum_number_should_match = minimum_number_should_match
        self.disable_coord = disable_coord
        if must:
            self.add_must(must)
        if must_not:
            self.add_must_not(must_not)
        if should:
            self.add_should(should)

    def add_must(self, queries):
        """Add a query to the "must" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._must.extend(queries)
        else:
            self._must.append(queries)
        return self

    def add_should(self, queries):
        """Add a query to the "should" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._should.extend(queries)
        else:
            self._should.append(queries)
        return self

    def add_must_not(self, queries):
        """Add a query to the "must_not" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._must_not.extend(queries)
        else:
            self._must_not.append(queries)
        return self

    def is_empty(self):
        if self._must:
            return False
        if self._must_not:
            return False
        if self._should:
            return False
        return True

    def _serialize(self):
        filters = {}
        if self._must:
            filters['must'] = [f.serialize() for f in self._must]
        if self._must_not:
            filters['must_not'] = [f.serialize() for f in self._must_not]
        if self._should:
            filters['should'] = [f.serialize() for f in self._should]
            filters['minimum_number_should_match'] = self.minimum_number_should_match
        if self.boost:
            filters['boost'] = self.boost
        if self.disable_coord is not None:
            filters['disable_coord'] = self.disable_coord
        if not filters:
            raise RuntimeError("A least a filter must be declared")
        return filters


class ConstantScoreQuery(Query):
    """Returns a constant score for all documents matching a filter.

    Multiple filters may be supplied by passing a sequence or iterator as the
    filter parameter.  If multiple filters are supplied, documents must match
    all of them to be matched by this query.
    """

    _internal_name = "constant_score"

    def __init__(self, filter=None, boost=1.0, **kwargs):
        super(ConstantScoreQuery, self).__init__(**kwargs)
        self.filters = []
        self.queries = []
        self.boost = boost
        if filter:
            self.add(filter)

    def add(self, filter_or_query):
        """Add a filter, or a list of filters, to the query.

        If a sequence of filters is supplied, they are all added, and will be
        combined with an ANDFilter.

        If a sequence of queries is supplied, they are all added, and will be
        combined with an BooleanQuery(must) .

        """
        if isinstance(filter_or_query, Filter):
            if self.queries:
                raise QueryError("A Query is required")
            self.filters.append(filter_or_query)
        elif isinstance(filter_or_query, Query):
            if self.filters:
                raise QueryError("A Filter is required")
            self.queries.append(filter_or_query)
        else:
            for item in filter_or_query:
                self.add(item)
        return self

    def is_empty(self):
        """Returns True if the query is empty.

        """
        if self.filters:
            return False
        if self.queries:
            return False
        return True

    def _serialize(self):
        data = {}
        if self.boost != 1.0:
            data["boost"] = self.boost

        if self.filters:
            filters = {}
            if len(self.filters) == 1:
                filters.update(self.filters[0].serialize())
            else:
                filters.update(ANDFilter(self.filters).serialize())
            if not filters:
                raise QueryError("A filter is required")
            data['filter'] = filters
        else:
            queries = {}
            if len(self.queries) == 1:
                queries.update(self.queries[0].serialize())
            else:
                queries.update(BoolQuery(must=self.queries).serialize())
            data['query'] = queries

        return data


class HasQuery(Query):

    def __init__(self, type, query, _scope=None, score_mode=None, **kwargs):
        super(HasQuery, self).__init__(**kwargs)
        self.type = type
        self._scope = _scope
        self.score_mode = score_mode
        self.query = query

    def _serialize(self):
        data = {'type': self.type, 'query': self.query.serialize()}
        if self._scope is not None:
            data['_scope'] = self._scope
        if self.score_mode is not None:
            data['score_mode'] = self.score_mode
        return data


class HasChildQuery(HasQuery):

    _internal_name = "has_child"

    def __init__(self, *args, **kwargs):
        super(HasChildQuery, self).__init__(*args, **kwargs)
        if self.score_mode and self.score_mode not in ["max", "sum", "avg", "none"]:
            raise InvalidParameterQuery("Invalid value '%s' for score_mode" % self.score_mode)


class HasParentQuery(HasQuery):

    _internal_name = "has_parent"

    def __init__(self, *args, **kwargs):
        super(HasParentQuery, self).__init__(*args, **kwargs)
        if self.score_mode and self.score_mode not in ["score", "none"]:
            raise InvalidParameterQuery("Invalid value '%s' for score_mode" % self.score_mode)


class TopChildrenQuery(ConstantScoreQuery):

    _internal_name = "top_children"

    def __init__(self, type, score="max", factor=5, incremental_factor=2, **kwargs):
        super(TopChildrenQuery, self).__init__(**kwargs)
        self.type = type
        self.score = score
        self.factor = factor
        self.incremental_factor = incremental_factor

    def _serialize(self):
        if self.score not in ["max", "min", "avg", "sum"]:
            raise InvalidParameterQuery("Invalid value '%s' for score" % self.score)

        queries = {}
        if self.boost != 1.0:
            queries["boost"] = self.boost
        for f in self.queries:
            queries.update(f.serialize())
        return {
            'type': self.type,
            'query':queries ,
            'score': self.score,
            'factor': self.factor,
            'incremental_factor': self.incremental_factor,
        }


class NestedQuery(Query):
    """
    Nested query allows to query nested objects / docs (see nested mapping).
    The query is executed against the nested objects / docs as if they were
    indexed as separate docs (they are, internally) and resulting in the root
    parent doc (or parent nested mapping).<commit_msg>Space moved :S.<commit_after># -*- coding: utf-8 -*-

import six
from .exceptions import InvalidQuery, InvalidParameterQuery, QueryError, \
    ScriptFieldsError
from .sort import SortFactory
from .facets import FacetFactory
from .aggs import AggFactory
from .filters import ANDFilter, Filter
from .highlight import HighLighter
from .scriptfields import ScriptFields
from .utils import clean_string, ESRange, EqualityComparableUsingAttributeDictionary

class Suggest(EqualityComparableUsingAttributeDictionary):
    def __init__(self, fields=None):
        self.fields = fields or {}

    def add(self, text, name, field, type='term', size=None, params=None):
        """
        Set the suggester of given type.

        :param text: text
        :param name: name of suggest
        :param field: field to be used
        :param type: type of suggester to add, available types are: completion,
                     phrase, term
        :param size: number of phrases
        :param params: dict of additional parameters to pass to the suggester
        :return: None
        """
        func = None

        if type == 'completion':
            func = self.add_completion
        elif type == 'phrase':
            func = self.add_phrase
        elif type == 'term':
            func = self.add_term
        else:
            raise InvalidQuery('Invalid type')

        func(text=text, name=name, field=field, size=size, params=params)

    def add_completion(self, text, name, field, size=None, params=None):
        """
        Add completion-type suggestion:

        http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html

        :param text: text for searching using autocompletion
        :param name: name for the suggester
        :param field: document's field to be autocompleted
        :param size: (optional) size of the autocomplete list
        :return:
        """
        data = {
            'field': field
        }

        if size:
            data['size'] = size

        if params:
            data.update(params)

        self.fields[name] = {
            'text': text,
            'completion': data,
        }


    def add_term(self, text, name, field, size=None, params=None):
        data = {"field": field}

        if size:
            data["size"] = size

        if params:
            data.update(params)

        self.fields[name] = {"text": text, "term": data}

    def add_phrase(self, text, name, field, size=None, params=None):
        tokens = text.split()
        gram = field + ".bigram"
        # if len(tokens) > 3:
        #     gram = field + ".trigram"

        data = {
            "analyzer": "standard_lower",
            "field": gram,
            "size": 4,
            "real_word_error_likelihood": 0.95,
            "confidence": 2.0,
            "gram_size": 2,
            "direct_generator": [{
                                     "field": field + ".tkl",
                                     "suggest_mode": "always",
                                     "min_word_len": 1
                                 }, {
                                     "field": field + ".reverse",
                                     "suggest_mode": "always",
                                     "min_word_len": 1,
                                     "pre_filter": "reverse",
                                     "post_filter": "reverse"
                                 }]
        }

        if size:
            data["size"] = size

        if params:
            data.update(params)

        self.fields[name] = {"text": text, "phrase": data}

    def is_valid(self):
        return len(self.fields) > 0

    def serialize(self):
        return self.fields

class FieldParameter(EqualityComparableUsingAttributeDictionary):

    def __init__(self, field, query, default_operator="OR", analyzer=None,
                 allow_leading_wildcard=True, lowercase_expanded_terms=True,
                 enable_position_increments=True, fuzzy_prefix_length=0,
                 fuzzy_min_sim=0.5, phrase_slop=0, boost=1.0):
        self.query = query
        self.field = field
        self.default_operator = default_operator
        self.analyzer = analyzer
        self.allow_leading_wildcard = allow_leading_wildcard
        self.lowercase_expanded_terms = lowercase_expanded_terms
        self.enable_position_increments = enable_position_increments
        self.fuzzy_prefix_length = fuzzy_prefix_length
        self.fuzzy_min_sim = fuzzy_min_sim
        self.phrase_slop = phrase_slop
        self.boost = boost

    def serialize(self):
        filters = {}

        if self.default_operator != "OR":
            filters["default_operator"] = self.default_operator
        if self.analyzer:
            filters["analyzer"] = self.analyzer
        if not self.allow_leading_wildcard:
            filters["allow_leading_wildcard"] = self.allow_leading_wildcard
        if not self.lowercase_expanded_terms:
            filters["lowercase_expanded_terms"] = self.lowercase_expanded_terms
        if not self.enable_position_increments:
            filters["enable_position_increments"] = self.enable_position_increments
        if self.fuzzy_prefix_length:
            filters["fuzzy_prefix_length"] = self.fuzzy_prefix_length
        if self.fuzzy_min_sim != 0.5:
            filters["fuzzy_min_sim"] = self.fuzzy_min_sim
        if self.phrase_slop:
            filters["phrase_slop"] = self.phrase_slop

        if self.boost != 1.0:
            filters["boost"] = self.boost
        if filters:
            filters["query"] = self.query
        else:
            filters = self.query
        return self.field, filters


class Search(EqualityComparableUsingAttributeDictionary):
    """A search to be performed.

    This contains a query, and has additional parameters which are used to
    control how the search works, what it should return, etc.

    The rescore parameter takes a Search object created from a RescoreQuery.

    Example:

    q = QueryStringQuery('elasticsearch')
    s = Search(q, fields=['title', 'author'], start=100, size=50)
    results = conn.search(s)
    """

    def __init__(self, query=None, filter=None, fields=None, start=None,
                 size=None, highlight=None, sort=None, explain=False, facet=None, agg=None, rescore=None,
                 window_size=None, version=None, track_scores=None, script_fields=None, index_boost=None,
                 min_score=None, stats=None, bulk_read=None, partial_fields=None, _source=None, timeout=None):
        """
        fields: if is [], the _source is not returned
        """
        if not index_boost: index_boost = {}
        self.query = query
        self.filter = filter
        self.fields = fields
        self.start = start
        self.size = size
        self._highlight = highlight
        self.sort = sort or SortFactory()
        self.explain = explain
        self.facet = facet or FacetFactory()
        self.agg = agg or AggFactory()
        self.rescore = rescore
        self.window_size = window_size
        self.version = version
        self.track_scores = track_scores
        self._script_fields = script_fields
        self.index_boost = index_boost
        self.min_score = min_score
        self.stats = stats
        self.bulk_read = bulk_read
        self.partial_fields = partial_fields
        self._source = _source
        self.timeout = timeout

    def get_facet_factory(self):
        """
        Returns the facet factory
        """
        return self.facet

    def get_agg_factory(self):
        """
        Returns the agg factory
        """
        return self.agg

    def serialize(self):
        """Serialize the search to a structure as passed for a search body."""
        res = {}
        if self.query:
            if isinstance(self.query, dict):
                res["query"] = self.query
            elif isinstance(self.query, Query):
                res["query"] = self.query.serialize()
            else:
                raise InvalidQuery("Invalid query")
        if self.filter:
            res['filter'] = self.filter.serialize()
        if self.facet.facets:
            res['facets'] = self.facet.serialize()
        if self.agg.aggs:
            res['aggs'] = self.agg.serialize()
        if self.rescore:
            res['rescore'] = self.rescore.serialize()
        if self.window_size:
            res['window_size'] = self.window_size
        if self.fields is not None: #Deal properly with self.fields = []
            res['fields'] = self.fields
        if self.size is not None:
            res['size'] = self.size
        if self.start is not None:
            res['from'] = self.start
        if self._highlight:
            res['highlight'] = self._highlight.serialize()
        if self.sort:
            if isinstance(self.sort, SortFactory):
                sort = self.sort.serialize()
                if sort:
                    res['sort'] = sort
            else:
                res['sort'] = self.sort
        if self.explain:
            res['explain'] = self.explain
        if self.version:
            res['version'] = self.version
        if self.track_scores:
            res['track_scores'] = self.track_scores
        if self._script_fields:
            if isinstance(self.script_fields, ScriptFields):
                res['script_fields'] = self.script_fields.serialize()
            elif isinstance(self.script_fields, dict):
                res['script_fields'] = self.script_fields
            else:
                raise ScriptFieldsError("Parameter script_fields should of type ScriptFields or dict")
        if self.index_boost:
            res['indices_boost'] = self.index_boost
        if self.min_score:
            res['min_score'] = self.min_score
        if self.stats:
            res['stats'] = self.stats
        if self.partial_fields:
            res['partial_fields'] = self.partial_fields
        if self._source:
            res['_source'] = self._source
        if self.timeout:
            res['timeout'] = self.timeout
        return res

    @property
    def highlight(self):
        if self._highlight is None:
            self._highlight = HighLighter("<b>", "</b>")
        return self._highlight

    @property
    def script_fields(self):
        if self._script_fields is None:
            self._script_fields = ScriptFields()
        return self._script_fields

    def add_highlight(self, field, fragment_size=None,
                      number_of_fragments=None, fragment_offset=None, type=None):
        """Add a highlight field.

        The Search object will be returned, so calls to this can be chained.

        """
        if self._highlight is None:
            self._highlight = HighLighter("<b>", "</b>")
        self._highlight.add_field(field, fragment_size, number_of_fragments, fragment_offset, type=type)
        return self

    def add_index_boost(self, index, boost):
        """Add a boost on an index.

        The Search object will be returned, so calls to this can be chained.

        """
        if boost is None:
            if index in self.index_boost:
                del(self.index_boost[index])
        else:
            self.index_boost[index] = boost
        return self

    def __repr__(self):
        return str(self.serialize())


class Query(EqualityComparableUsingAttributeDictionary):
    """Base class for all queries."""

    def __init__(self, *args, **kwargs):
        if len(args) > 0 or len(kwargs) > 0:
            raise RuntimeWarning("No all parameters are processed by derivated query object")

    def search(self, **kwargs):
        """Return this query wrapped in a Search object.

        Any keyword arguments supplied to this call will be passed to the
        Search object.
        """
        return Search(query=self, **kwargs)

    def serialize(self):
        """Serialize the query to a structure using the query DSL."""
        return {self._internal_name: self._serialize()}

    def _serialize(self):
        raise NotImplementedError

    @property
    def _internal_name(self):
        raise NotImplementedError


class BoolQuery(Query):
    """A boolean combination of other queries.

    BoolQuery maps to Lucene **BooleanQuery**. It is built using one or more
    boolean clauses, each clause with a typed occurrence.  The occurrence types
    are:

    ================  ========================================================
     Occur             Description
    ================  ========================================================
    **must**          The clause (query) must appear in matching documents.
    **should**        The clause (query) should appear in the matching
                      document. A boolean query with no **must** clauses, one
                      or more **should** clauses must match a document. The
                      minimum number of should clauses to match can be set
                      using **minimum_number_should_match** parameter.
    **must_not**      The clause (query) must not appear in the matching
                      documents. Note that it is not possible to search on
                      documents that only consists of a **must_not** clauses.
    ================  ========================================================

    The bool query also supports **disable_coord** parameter (defaults to
    **false**).
    """

    _internal_name = "bool"

    def __init__(self, must=None, must_not=None, should=None, boost=None,
                 minimum_number_should_match=1, disable_coord=None, **kwargs):
        super(BoolQuery, self).__init__(**kwargs)
        self._must = []
        self._must_not = []
        self._should = []
        self.boost = boost
        self.minimum_number_should_match = minimum_number_should_match
        self.disable_coord = disable_coord
        if must:
            self.add_must(must)
        if must_not:
            self.add_must_not(must_not)
        if should:
            self.add_should(should)

    def add_must(self, queries):
        """Add a query to the "must" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._must.extend(queries)
        else:
            self._must.append(queries)
        return self

    def add_should(self, queries):
        """Add a query to the "should" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._should.extend(queries)
        else:
            self._should.append(queries)
        return self

    def add_must_not(self, queries):
        """Add a query to the "must_not" clause of the query.

        The Query object will be returned, so calls to this can be chained.
        """
        if isinstance(queries, list):
            self._must_not.extend(queries)
        else:
            self._must_not.append(queries)
        return self

    def is_empty(self):
        if self._must:
            return False
        if self._must_not:
            return False
        if self._should:
            return False
        return True

    def _serialize(self):
        filters = {}
        if self._must:
            filters['must'] = [f.serialize() for f in self._must]
        if self._must_not:
            filters['must_not'] = [f.serialize() for f in self._must_not]
        if self._should:
            filters['should'] = [f.serialize() for f in self._should]
            filters['minimum_number_should_match'] = self.minimum_number_should_match
        if self.boost:
            filters['boost'] = self.boost
        if self.disable_coord is not None:
            filters['disable_coord'] = self.disable_coord
        if not filters:
            raise RuntimeError("A least a filter must be declared")
        return filters


class ConstantScoreQuery(Query):
    """Returns a constant score for all documents matching a filter.

    Multiple filters may be supplied by passing a sequence or iterator as the
    filter parameter.  If multiple filters are supplied, documents must match
    all of them to be matched by this query.
    """

    _internal_name = "constant_score"

    def __init__(self, filter=None, boost=1.0, **kwargs):
        super(ConstantScoreQuery, self).__init__(**kwargs)
        self.filters = []
        self.queries = []
        self.boost = boost
        if filter:
            self.add(filter)

    def add(self, filter_or_query):
        """Add a filter, or a list of filters, to the query.

        If a sequence of filters is supplied, they are all added, and will be
        combined with an ANDFilter.

        If a sequence of queries is supplied, they are all added, and will be
        combined with an BooleanQuery(must) .

        """
        if isinstance(filter_or_query, Filter):
            if self.queries:
                raise QueryError("A Query is required")
            self.filters.append(filter_or_query)
        elif isinstance(filter_or_query, Query):
            if self.filters:
                raise QueryError("A Filter is required")
            self.queries.append(filter_or_query)
        else:
            for item in filter_or_query:
                self.add(item)
        return self

    def is_empty(self):
        """Returns True if the query is empty.

        """
        if self.filters:
            return False
        if self.queries:
            return False
        return True

    def _serialize(self):
        data = {}
        if self.boost != 1.0:
            data["boost"] = self.boost

        if self.filters:
            filters = {}
            if len(self.filters) == 1:
                filters.update(self.filters[0].serialize())
            else:
                filters.update(ANDFilter(self.filters).serialize())
            if not filters:
                raise QueryError("A filter is required")
            data['filter'] = filters
        else:
            queries = {}
            if len(self.queries) == 1:
                queries.update(self.queries[0].serialize())
            else:
                queries.update(BoolQuery(must=self.queries).serialize())
            data['query'] = queries

        return data


class HasQuery(Query):

    def __init__(self, type, query, _scope=None, score_mode=None, **kwargs):
        super(HasQuery, self).__init__(**kwargs)
        self.type = type
        self._scope = _scope
        self.score_mode = score_mode
        self.query = query

    def _serialize(self):
        data = {'type': self.type, 'query': self.query.serialize()}
        if self._scope is not None:
            data['_scope'] = self._scope
        if self.score_mode is not None:
            data['score_mode'] = self.score_mode
        return data


class HasChildQuery(HasQuery):

    _internal_name = "has_child"

    def __init__(self, *args, **kwargs):
        super(HasChildQuery, self).__init__(*args, **kwargs)
        if self.score_mode and self.score_mode not in ["max", "sum", "avg", "none"]:
            raise InvalidParameterQuery("Invalid value '%s' for score_mode" % self.score_mode)


class HasParentQuery(HasQuery):

    _internal_name = "has_parent"

    def __init__(self, *args, **kwargs):
        super(HasParentQuery, self).__init__(*args, **kwargs)
        if self.score_mode and self.score_mode not in ["score", "none"]:
            raise InvalidParameterQuery("Invalid value '%s' for score_mode" % self.score_mode)


class TopChildrenQuery(ConstantScoreQuery):

    _internal_name = "top_children"

    def __init__(self, type, score="max", factor=5, incremental_factor=2, **kwargs):
        super(TopChildrenQuery, self).__init__(**kwargs)
        self.type = type
        self.score = score
        self.factor = factor
        self.incremental_factor = incremental_factor

    def _serialize(self):
        if self.score not in ["max", "min", "avg", "sum"]:
            raise InvalidParameterQuery("Invalid value '%s' for score" % self.score)

        queries = {}
        if self.boost != 1.0:
            queries["boost"] = self.boost
        for f in self.queries:
            queries.update(f.serialize())
        return {
            'type': self.type,
            'query': queries,
            'score': self.score,
            'factor': self.factor,
            'incremental_factor': self.incremental_factor,
        }


class NestedQuery(Query):
    """
    Nested query allows to query nested objects / docs (see nested mapping).
    The query is executed against the nested objects / docs as if they were
    indexed as separate docs (they are, internally) and resulting in the root
    parent doc (or parent nested mapping).
{'fraction_non_alphanumeric': 0.10384689178818111, 'fraction_numerical': 0.002014581734458941, 'mean_word_length': 2.0267857142857144, 'pattern_counts': {'":': 42, '<': 11, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3655521', 'n_tokens_mistral': 11333, 'n_tokens_neox': 10836, 'n_words': 3170}