<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct  4 18:43:40 2017

@author: m75380

Directions:
    - Parents do not include filters

Filter by core when should is almost never used

TODO:

- ref_gen
- buffer_for_future_matches in same row when previous is called:
    overwrite ref_gen
- query_filter   


for query in labeller.current_queries:
    print('\n*******')
    print(query._as_tuple())
    print('precision:', query.precision)
    print('recall', query.recall)
    print('score', query.score)
    print('thresh', query.thresh)


TODO: check why precision is 1 even after no/yes

TODO: try multiple training orders on data

"""
import copy
import itertools
import pickle
import random

from elasticsearch import Elasticsearch, client
import numpy as np

from es_helpers import _gen_body, _bulk_search
from my_json_encoder import MyEncoder

# =============================================================================
# Dev decorators
# =============================================================================

INDENT_LVL = 0

def print_name(function):
    '''Print the function being run with indentation for inclusion'''
    def wrapper(*args, **kwargs):
        global INDENT_LVL
        my_id = random.randint(0, 100000)
        ind = 2 * INDENT_LVL * ' '
        INDENT_LVL += 1
        print('{0}<< Starting: {1} ({2})'.format(ind, function.__name__, my_id))
        res = function(*args, **kwargs)
        INDENT_LVL -= 1
        print('{0}>> Ending: {1} ({2})'.format(ind, function.__name__, my_id))
        return res
    return wrapper

# =============================================================================
# 
# =============================================================================

class SingleQueryTemplate():
    '''Class representing an Elasticsearch query template'''
    def __init__(self, bool_lvl, source_col, ref_col, analyzer_suffix, boost):
        self.bool_lvl = bool_lvl
        
        if isinstance(source_col, str):
            source_col = [source_col]
            
        if isinstance(ref_col, str):
            ref_col = [ref_col]
            
        self.source_col = tuple(sorted(source_col))
        self.ref_col = tuple(sorted(ref_col))
        self.analyzer_suffix = analyzer_suffix
        self.boost = boost
        
        self.core = self._core()
        

    def to_dict(self):
        '''Returns a dict representation of the class'''
        return {'bool_lvl': self.bool_lvl,
                'source_col': self.source_col,
                'ref_col': self.ref_col,
                'analyzer_suffix': self.analyzer_suffix,
                'boost': self.boost}
        
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        query_template = cls(dict_['bool_lvl'], 
                                   dict_['source_col'], 
                                   dict_['ref_col'], 
                                   dict_['analyzer_suffix'], 
                                   dict_['boost'])
        query_template.core = query_template._core()
        return query_template

    def _core(self):
        '''
        Return the minimal compononent that guarantees equivalence of the claim:
        "this query has results"    
        '''
        return (self.source_col, self.ref_col, self.analyzer_suffix)
        
    def __hash__(self):
        try:
            return self.__hash
        except:
            self.__hash = hash((self.bool_lvl, tuple(sorted(self.source_col)), tuple(sorted(self.source_col)), 
                     self.analyzer_suffix, self.boost))
            return self.__hash
    
    def __eq__(self, other):
        return self._as_tuple() == other._as_tuple()
    

    def __gt__(self, other):
        return self._as_tuple() > other._as_tuple()

#    def __str__(self):
#        return self._as_tuple().__str__()
    
    def _as_tuple(self):
        return (self.bool_lvl, self.source_col, self.ref_col, self.analyzer_suffix, self.boost)


class CompoundQueryTemplate():
    '''Information regarding a query to be used in the labeller'''
    def __init__(self, single_query_templates_tuple):        
        self.musts = tuple(sorted([SingleQueryTemplate(*x) for x in single_query_templates_tuple if x[0] == 'must']))
        self.shoulds = tuple(sorted([SingleQueryTemplate(*x) for x in single_query_templates_tuple if x[0] == 'should']))
        
        assert self.musts
        
        self.core = self._core()
        self.parent_cores = self._parent_cores()

    def add_must(self, query):
        self.musts = tuple(sorted(list(self.musts) + [query]))
        self.core = self._core()
        
    def add_should(self, query):
        self.shoulds = tuple(sorted(list(self.shoulds) + [query]))

    def to_dict(self):
        '''Returns a dict representation of the instance'''
        return [x.to_dict() for x in self.musts + self.shoulds]
        
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        single_query_templates_tuple = [SingleQueryTemplate.from_dict(d)._as_tuple() \
                                        for d in dict_]
        return cls(single_query_templates_tuple)

    def _core(self):
        '''Same as in SingleQueryTemplate'''
        cores = []
        for must in self.musts:
            cores.append(must._core())
        return tuple(sorted(cores))

    def __hash__(self):
        try:
            return self.__hash
        except:
            self.__hash = hash((q.__hash__() for q in self.musts + self.shoulds))    
            return self.__hash

    def __eq__(self, other):        
        bool_1 = self.musts == other.musts
        bool_2 = self.shoulds == other.shoulds
        
        return bool_1 and bool_2


    def __len__(self):
        return len(self.musts) + len(self.shoulds)

#    def __str__(self):
#        return [x.__str__() for x in self.musts + self.shoulds].__str__()

    def _parent_cores(self):
        '''Returns the core of all possible parents'''
        if len(self.core) == 1:
            return []
        else:
            to_return = []
            for i in range(1, len(self.core)-1):
                to_return.extend(list(itertools.combinations(self.core, i)))
            return to_return        
    
    
    def _gen_body(self, row, must_filter={}, must_not_filter={}, num_results=3):
        '''
        Generate the string to pass to Elastic search for it to execute query
        
        INPUT:
            - row: pandas.Series from the source object
            - must: terms to filter by field (AND: will include ONLY IF ALL are in text)
            - must_not: terms to exclude by field from search (OR: will exclude if ANY is found)
            - num_results: Max number of results for the query
        
        OUTPUT:
            - body: the query as string
        '''
        
        # TODO: do we use this ?
        body = _gen_body(self.musts + self.shoulds, row, must_filter, must_not_filter, num_results)
        return body
        
    def _as_tuple(self):
        return tuple(x._as_tuple() for x in self.musts + self.shoulds)
     
        
class LabellerQueryTemplate(CompoundQueryTemplate):
    '''
    Extension of CompoundQueryTemplate which can also store data relative
    to the performance of the query in the context of labelling.
    '''
    
    def __init__(self, *argv, **kwargs):
        super().__init__(*argv, **kwargs)
        
        self.reset()
        
    def _sanity_check(self):
        assert len(self.history_pairs) \
                == len(self.first_scores) \
                == len(self.has_results)
                
        assert len(self.first_is_match) == len(self.any_is_match)
        

    def reset(self):
        self.history_pairs = {} # {source_id': [(), (), ]}
        #self.steps = [] # Num rows labelled at which this pair was added (in case of later creation)
        self.first_scores = {} # ES scores of first result of queies
        self.has_results = {} # Whether the query returned any result
        self.first_is_match = {} # First hit is the actual match
        self.any_is_match = {} # Any of the hits is the actual match
        
        self.thresh = None
        self.precision = None
        self.recall = None
        self.inclusion_ratio = None
        self.score = 0.1 # General score of the template based on precision and recall        

    def to_dict(self):
        '''Returns a dict representation of the instance'''
        dict_ = super().to_dict()
        dict_['history_pairs'] = self.history_pairs
        dict_['first_scores'] = self.first_scores
        dict_['has_results'] = self.has_results
        dict_['first_is_match'] = self.first_is_match
        dict_['any_is_match'] = self.any_is_match
        
        dict_['thresh'] = self.thresh
        dict_['precision'] = self.precision
        dict_['recall'] = self.recall
        dict_['inclusion_ratio'] = self.inclusion_ratio
        dict_['score'] = self.score
        return dict_
    
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        labeller = super(LabellerQueryTemplate, cls).from_dict(cls, dict_)
        labeller['history_pairs'] = dict_['history_pairs']
        labeller['first_scores'] = dict_['first_scores']
        labeller['has_results'] = dict_['has_results']
        labeller['first_scores'] = dict_['first_scores']
        labeller['any_is_match'] = dict_['any_is_match']
        
        labeller['thresh'] = dict_['thresh']
        labeller['precision'] = dict_['precision']
        labeller['recall'] = dict_['recall']
        labeller['inclusion_ratio'] = dict_['inclusion_ratio']
        labeller['score'] = dict_['score']
        return labeller
        
    def add_pairs(self, source_idx, pairs, first_score):
        '''Add pairs to history_pairs (pairs that were generated by this query)'''
        self.history_pairs[source_idx] = pairs
        self.has_results[source_idx] = bool(pairs)
        self.first_scores[source_idx] = first_score
    
    def add_labelled_pair(self, source_idx, labelled_pair):
        ''' Update first_is_match, any_is_match with a labelled_pair (commes after add_pairs)'''
        
        assert len(self.first_scores) == len(self.history_pairs)
        
        # If row is forgotten or no matches were found in all queries, mark as unaplicable
        if labelled_pair[1] in ['__FORGET', '__NO_RESULT']: 
            self.first_is_match[source_idx] = None # None if does not apply
            self.any_is_match[source_idx] = None # None if does not apply
        else:
            if self.history_pairs[source_idx]:        
                self.first_is_match[source_idx] = labelled_pair==self.history_pairs[source_idx][0]
                self.any_is_match[source_idx] = labelled_pair in self.history_pairs[source_idx]
            else:
                self.first_is_match[source_idx] = False
                self.any_is_match[source_idx] = False
            
    def unlabel(self, source_idx):
        '''Remove all information regarding labels for source_id'''
        # Compute idx to go back to
        
        # Don't touch history_pairs or first_scores or has_results
    
        if source_idx in self.first_is_match:
            del self.first_is_match[source_idx]
            del self.any_is_match[source_idx]
        else:
            assert source_idx not in self.any_is_match

        # Check whether or not to re-compute scores 


    def compute_metrics(self, t_p=0.95, t_r=0.3):
        ''' 
        Compute the optimal threshold and the associated metrics 
        
        INPUT:
            - t_p: target_precision
            - t_r: target_recall
            
        OUTPUT:
            - thresh: This refers to the threshold on the Elasticsearch score
                    for this query. If the ES score of the best hit is above
                    this threshold, it can be considered to be a match; if not, 
                    it can be assumed that it the pair is not a match (or at 
                    least very uncertain).
                    
                    The threshold is computed so as to maximize the estimated 
                    query score (based on precision and recall).
            - precision: The precision corresponding to that computed using the 
                    optimal threshold.
            - recall: The recall corresponding to that computed using the 
                    optimal threshold
            - score: The optimal score (that yields the optimal threshold)
        '''
        summaries = [{'score': self.first_scores[source_idx], 
                      'has_results': self.has_results[source_idx],
                      'first_is_match': self.first_is_match[source_idx], 
                      'any_is_match': self.any_is_match[source_idx]} \
                      for source_idx in self.any_is_match.keys()]
        
        # Filter out relevent summaries only (line not forgotten)
        summaries = [summary for summary in summaries if summary['first_is_match'] is not None]
        
        #
        num_summaries = len(summaries)
    
        # Sort summaries by score / NB: 0 score deals with empty hits
        summaries = sorted(summaries, key=lambda x: x['score'], reverse=True)
        
        # score_vect = np.array([x['_score_first'] for x in sorted_summaries])
        is_first_vect = np.array([bool(x['first_is_match']) for x in summaries])
        
        # If no matches are present
        # TODO: eliminates things from start ?
        if sum(is_first_vect) == 0:
            self.thresh = 10**3
            self.precision = 0
            self.recall = 0
            self.score = 0 
            return self.thresh, self.precision, self.recall, self.score
     
        
        rolling_precision = is_first_vect.cumsum() / (np.arange(num_summaries) + 1)
        rolling_recall = is_first_vect.cumsum() / num_summaries
        
        # TODO: WHY is recall same as precision ?
        
        # Compute score
        _f_precision = lambda x: max(x - t_p, 0) + min(t_p*(x/t_p)**3, t_p)
        _f_recall = lambda x: max(x - t_r, 0) + min(t_p*(x/t_r)**3, t_r)
        a = np.fromiter((_f_precision(xi) for xi in rolling_precision), rolling_precision.dtype)
        b = np.fromiter((_f_recall(xi) for xi in rolling_recall), rolling_recall.dtype)
        rolling_score = a*b
    
        # Find best index for threshold
        MIN_OBSERVATIONS = 4 # minimal number of values on which to comupute threshold etc...
        idx = max(num_summaries - rolling_score[::-1].argmax() - 1, min(MIN_OBSERVATIONS, num_summaries-1))
        
        # If the best score is obtained using all results (no ES score 
        # thresholding), then set threshold to ~0 
        if idx == len(summaries) - 1:
            thresh = 0.0001
        else:
            thresh = summaries[idx]['score']
        
        precision = rolling_precision[idx]
        recall = rolling_recall[idx]
        score = rolling_score[idx]
        
        # Compute 
        inclusion_ratio = sum(x['any_is_match'] for x in summaries) / len(summaries)
        
        self.thresh = thresh
        self.precision = precision
        self.recall = recall
        self.score = score
        
        return self.thresh, self.precision, self.recall, self.score


    def multiply_by_boost(self, boost_multiplier=2):
        '''Takes a single compound query and returns a list of multiple queries'''
        new_query_templates = []
        
        new_query = copy.deepcopy(self) 
        new_query_templates.append(new_query) # Keep original query
        
        og_boost_total = sum(x.boost for x in self.musts + self.shoulds)
        for level in ['shoulds', 'musts']:
            for i in range(len(self.__dict__[level])):
                new_query = copy.deepcopy(self)
                new_boost_total = og_boost_total + new_query.__dict__[level][i].boost
                new_query.__dict__[level][i].boost *= 2
                
                # Normalize total boost back to original level
                for level_2 in ['shoulds', 'musts']:
                    for j in range(len(self.__dict__[level_2])):
                        new_query.__dict__[level_2][j].boost *= og_boost_total/new_boost_total
                new_query_templates.append(new_query)
                
        return new_query_templates
    
    def multiply_by_core(self, core_queries, bool_levels):
        '''Takes a single compound query and returns a list of multiple queries'''
        
        new_query_templates = []
        
        new_query = copy.deepcopy(self)
        new_query_templates.append(new_query) # Keep original query
        
        # Multiply by core queries and bool levels
        for core_query in core_queries:
            for level in bool_levels:
                if level == 'must':
                    if all(core_query.core != q.core for q in self.musts):
                        new_query = copy.deepcopy(self)
                        new_query.add_must(core_query)
                        new_query_templates.append(new_query)
                elif level == 'should':
                    if all(core_query.core != q.core for q in self.musts):
                        new_query = copy.deepcopy(self)
                        new_query.add_should(core_query)
                        new_query_templates.append(new_query)   
                else:
                    raise ValueError('Invalid level: {0}; should be must or should')
        return new_query_templates


class CoreScorerQueryTemplate(SingleQueryTemplate):
    '''Scoer'''
    
    def __init__(self, *argv, **kwargs):
        super().__init__(*argv, **kwargs)    
        
        self.source_lens = []
        self.ref_lens = []
        self.intersect_lens = []
    
    def _sanity_check(self):
        assert len(self.source_lens) \
                == len(self.ref_lens) \
                == len(self.intersect_lens)
    
    @staticmethod
    def _analyze(es, index_name, analyzer, text):
        '''Use the current analyzer in the given index_name to analyze text'''
        ic = client.IndicesClient(es)
        if analyzer:
            return ic.analyze(index_name, body={'text': text, 'analyzer': analyzer})
        else:
            return ic.analyze(index_name, body={'text': text})
        
    def analyze_pair_items(self, es, index_name, source_item, ref_item):
        '''Count tokens for source, ref and the intersection fof both'''  
        analyzer = self.analyzer_suffix.strip('.')
        
        text_source = ' '.join(source_item[col] for col in self.source_col if source_item[col] is not None)
        tokens_source = {x['token'] for x in self._analyze(es, index_name, analyzer, text_source)['tokens']}
        
        text_ref = ' '.join(ref_item[col] for col in self.ref_col  if ref_item[col] is not None)
        tokens_ref = {x['token'] for x in self._analyze(es, index_name, analyzer, text_ref)['tokens']}

        return len(tokens_source), len(tokens_ref), len(tokens_source & tokens_ref)
    
    def add_labelled_pair_items(self, es, index_name, source_item, ref_item):
        ''''''
        l_s, l_r, l_i = self.analyze_pair_items(es, index_name, source_item, ref_item)
        
        self.source_lens.append(l_s)
        self.ref_lens.append(l_r)
        self.intersect_lens.append(l_i)
        
        self.score = sum(x > 0 for x in self.intersect_lens) / len(self.intersect_lens)
        
        
        

def _gen_suffix(columns_to_index, s_q_t_2):
    '''
    Yields suffixes to add to field_names for the given analyzers
    
    INPUT:
        - columns_to_index:
        - s_q_t_2: column or columns from ref
        
    NB: s_q_t_2: element two of the single_query_template
    
    '''
    
    if isinstance(s_q_t_2, str):
        analyzers = columns_to_index[s_q_t_2]
    elif isinstance(s_q_t_2, tuple):
        analyzers = set.union(*[columns_to_index[col] for col in s_q_t_2])
    else:
        raise ValueError('s_q_t_2 should be str or tuple (not list)')
    yield '' # No suffix for standard analyzer
    for analyzer in analyzers:
        yield '.' + analyzer


def DETUPLIFY_TODO_DELETE(arg):
    if isinstance(arg, tuple) and (len(arg) == 1):
        return arg[0]
    return arg


def _gen_all_single_query_template_tuples(match_cols, columns_to_index, bool_levels,
                                          boost_levels):
    '''Generate a list of single query tuples'''
    return list(((bool_lvl, DETUPLIFY_TODO_DELETE(x['source']), DETUPLIFY_TODO_DELETE(x['ref']), suffix, boost) \
                                       for x in match_cols \
                                       for suffix in _gen_suffix(columns_to_index, x['ref']) \
                                       for bool_lvl in bool_levels.get(suffix, ['must']) \
                                       for boost in boost_levels))

def _gen_all_query_template_tuples(match_cols, columns_to_index, bool_levels, 
                            boost_levels, max_num_levels):
    ''' 
    Generate query templates (as tuples). This will generate most combinations
    of queries that are possible using the different input.
    
    INPUT:
        match_cols: # TODO: see elsewhere
        columns_to_index: # TODO: see elsewhere
        bool_levels: list of the bool levels to generate on: ['must'], or ['must', 'should']
        boost_levels: all original boosts (see ES boost) to try. It is strongly
                    recommended to initiate with [1] (all boosts at 1) to avoid 
                    generating too many query templates
        max_num_levels: The maximal number of single queries in the compound
                    queries that are generated
                    
    OUTPUT:
        all_query_templates: A list of tuples representing compound queries
    '''
    single_queries = _gen_all_single_query_template_tuples(match_cols, 
                                    columns_to_index, bool_levels, boost_levels)
    all_query_templates = list(itertools.chain(*[list(itertools.combinations(single_queries, x)) \
                                        for x in range(2, max_num_levels+1)][::-1]))
    # Queries must contain at least two distinct columns (unless one match is mentionned)
    if len(match_cols) > 1:
        all_query_templates = list(filter(lambda query: len(set((x[1], x[2]) for x in query)) >= 2, \
                                    all_query_templates))

    all_query_templates = list(filter(lambda query: 'must' in [s_q_t[0] for s_q_t in query], \
                                all_query_templates))
    return all_query_templates


"""

"""


"""
Initialise core queries (list)
On match:
    - Analyze source and ref Using the query analyzer
    - Add number Of common tokens 
"""





class Labeller():
    '''
    Labeller object that stores previous labels and generates new matches
    for the user to label
    
    0.1) Initiate queries/ metrics/ history
    0.2) Read first row
    
    1) Perform queries / update history of hits
    2) Generate pairs to propose (based on sorted queries)
    3...) Until row is over: User inputs label
    4) Update metrics and history 
    5) Sort queries
    6) Gen new row and back to 1)
        
        
    HOW TO USE:
    
    >>> labeller = Labeller()
    
    >>> for x in range(100):
            to_display = labeller.to_emit()
            labeller.update('y')
    
    >>> labeller.export_params()    
        
    '''
    NUM_RESULTS = 3
    
    MAX_NUM_SAMPLES = 100
    
    VALID_ANSWERS = {'yes': 'y',
                 'y': 'y',
                 '1': 'y',
                 
                 'no': 'n',
                 'n': 'n',
                 '0': 'n',
                 
                    #                 'uncertain': 'u',
                    #                 'u': 'u',
                 
                 'forget_row': 'f',
                 'f': 'f',
                 
                 'previous': 'p',
                 'p': 'p',

                 'quit': 'q'             
                 }

    #min_precision_tab = [(20, 0.7), (10, 0.5), (5, 0.3)]
    #max_num_keys_tab = [(20, 10), (10, 50), (7, 200), (5, 500), (0, 1000)]
    #num_results_labelling = 3
    
    MAX_NUM_LEVELS = 3 # Number of match clauses
    BOOL_LEVELS = {'.integers': ['must', 'should'], 
                   '.city': ['must', 'should']}
    BOOST_LEVELS = [1]
    
    # Target precision and target recall
    t_p = 0.95
    t_r = 0.3    
    
    
    def __init__(self, source, ref_index_name, 
                 match_cols, columns_to_index, 
                 certain_column_matches=None, 
                 must={}, must_not={}):
        '''        
        INPUT: 
          - source: pandas DataFrame containing input data to match
          - ref_index_name: name of the Elasticsearch index to search in
          - match_cols: Indication of columns to try to match:
              EX: [{'source': ['city'], 'ref': ['CITY_ADD']},
                   {'source': ['first_name', 'last_name'], 'ref': ['FULL_NAME']}]
          - columns_to_index: What columns were indexed and how:
              EX:  {
                    'CITY_ADD': {},
                    'FULL_NAME': {'french', 'integers', 'n_grams', 'city'}
                    }
          - must_filters
        '''
        
        self.source = source
        self.ref_index_name = ref_index_name
        self.match_cols = match_cols      
        self.columns_to_index = columns_to_index
        
        self._init_es()
        
        self.labelled_pairs = [] # Flat list of labelled pairs
        self.labels = {} # Flat list of labels
        self.num_rows_labelled = [] # Flat list: at given label, how many were labelled. NB: starts at 0, becomes 1 at first yes/forgotten, or when next_row
        self.num_positive_rows_labelled = [] # Flat list: at given label, how_many were matches
        self.labelled_pairs_match = [] # For each row, the resulting match: (A, B) / no-match: (A, None) or forgotten: None
         
        
        self.must_filters = must
        self.must_not_filters = must_not
        
           
        self._init_queries(match_cols, columns_to_index) # creates self.current_queries
        self._init_core_queries(match_cols, columns_to_index) # creates self.single_core_queries
        #self._init_history() # creates self.history
        <commit_msg>distinction of forget and all no<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Oct  4 18:43:40 2017

@author: m75380

Directions:
    - Parents do not include filters

Filter by core when should is almost never used

TODO:

- ref_gen
- buffer_for_future_matches in same row when previous is called:
    overwrite ref_gen
- query_filter   


for query in labeller.current_queries:
    print('\n*******')
    print(query._as_tuple())
    print('precision:', query.precision)
    print('recall', query.recall)
    print('score', query.score)
    print('thresh', query.thresh)


TODO: check why precision is 1 even after no/yes

TODO: try multiple training orders on data

"""
import copy
import itertools
import pickle
import random

from elasticsearch import Elasticsearch, client
import numpy as np

from es_helpers import _gen_body, _bulk_search
from my_json_encoder import MyEncoder

# =============================================================================
# Dev decorators
# =============================================================================

INDENT_LVL = 0

def print_name(function):
    '''Print the function being run with indentation for inclusion'''
    def wrapper(*args, **kwargs):
        global INDENT_LVL
        my_id = random.randint(0, 100000)
        ind = 2 * INDENT_LVL * ' '
        INDENT_LVL += 1
        print('{0}<< Starting: {1} ({2})'.format(ind, function.__name__, my_id))
        res = function(*args, **kwargs)
        INDENT_LVL -= 1
        print('{0}>> Ending: {1} ({2})'.format(ind, function.__name__, my_id))
        return res
    return wrapper

# =============================================================================
# 
# =============================================================================

class SingleQueryTemplate():
    '''Class representing an Elasticsearch query template'''
    def __init__(self, bool_lvl, source_col, ref_col, analyzer_suffix, boost):
        self.bool_lvl = bool_lvl
        
        if isinstance(source_col, str):
            source_col = [source_col]
            
        if isinstance(ref_col, str):
            ref_col = [ref_col]
            
        self.source_col = tuple(sorted(source_col))
        self.ref_col = tuple(sorted(ref_col))
        self.analyzer_suffix = analyzer_suffix
        self.boost = boost
        
        self.core = self._core()
        

    def to_dict(self):
        '''Returns a dict representation of the class'''
        return {'bool_lvl': self.bool_lvl,
                'source_col': self.source_col,
                'ref_col': self.ref_col,
                'analyzer_suffix': self.analyzer_suffix,
                'boost': self.boost}
        
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        query_template = cls(dict_['bool_lvl'], 
                                   dict_['source_col'], 
                                   dict_['ref_col'], 
                                   dict_['analyzer_suffix'], 
                                   dict_['boost'])
        query_template.core = query_template._core()
        return query_template

    def _core(self):
        '''
        Return the minimal compononent that guarantees equivalence of the claim:
        "this query has results"    
        '''
        return (self.source_col, self.ref_col, self.analyzer_suffix)
        
    def __hash__(self):
        try:
            return self.__hash
        except:
            self.__hash = hash((self.bool_lvl, tuple(sorted(self.source_col)), tuple(sorted(self.source_col)), 
                     self.analyzer_suffix, self.boost))
            return self.__hash
    
    def __eq__(self, other):
        return self._as_tuple() == other._as_tuple()
    

    def __gt__(self, other):
        return self._as_tuple() > other._as_tuple()

#    def __str__(self):
#        return self._as_tuple().__str__()
    
    def _as_tuple(self):
        return (self.bool_lvl, self.source_col, self.ref_col, self.analyzer_suffix, self.boost)


class CompoundQueryTemplate():
    '''Information regarding a query to be used in the labeller'''
    def __init__(self, single_query_templates_tuple):        
        self.musts = tuple(sorted([SingleQueryTemplate(*x) for x in single_query_templates_tuple if x[0] == 'must']))
        self.shoulds = tuple(sorted([SingleQueryTemplate(*x) for x in single_query_templates_tuple if x[0] == 'should']))
        
        assert self.musts
        
        self.core = self._core()
        self.parent_cores = self._parent_cores()

    def add_must(self, query):
        self.musts = tuple(sorted(list(self.musts) + [query]))
        self.core = self._core()
        
    def add_should(self, query):
        self.shoulds = tuple(sorted(list(self.shoulds) + [query]))

    def to_dict(self):
        '''Returns a dict representation of the instance'''
        return [x.to_dict() for x in self.musts + self.shoulds]
        
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        single_query_templates_tuple = [SingleQueryTemplate.from_dict(d)._as_tuple() \
                                        for d in dict_]
        return cls(single_query_templates_tuple)

    def _core(self):
        '''Same as in SingleQueryTemplate'''
        cores = []
        for must in self.musts:
            cores.append(must._core())
        return tuple(sorted(cores))

    def __hash__(self):
        try:
            return self.__hash
        except:
            self.__hash = hash((q.__hash__() for q in self.musts + self.shoulds))    
            return self.__hash

    def __eq__(self, other):        
        bool_1 = self.musts == other.musts
        bool_2 = self.shoulds == other.shoulds
        
        return bool_1 and bool_2


    def __len__(self):
        return len(self.musts) + len(self.shoulds)

#    def __str__(self):
#        return [x.__str__() for x in self.musts + self.shoulds].__str__()

    def _parent_cores(self):
        '''Returns the core of all possible parents'''
        if len(self.core) == 1:
            return []
        else:
            to_return = []
            for i in range(1, len(self.core)-1):
                to_return.extend(list(itertools.combinations(self.core, i)))
            return to_return        
    
    
    def _gen_body(self, row, must_filter={}, must_not_filter={}, num_results=3):
        '''
        Generate the string to pass to Elastic search for it to execute query
        
        INPUT:
            - row: pandas.Series from the source object
            - must: terms to filter by field (AND: will include ONLY IF ALL are in text)
            - must_not: terms to exclude by field from search (OR: will exclude if ANY is found)
            - num_results: Max number of results for the query
        
        OUTPUT:
            - body: the query as string
        '''
        
        # TODO: do we use this ?
        body = _gen_body(self.musts + self.shoulds, row, must_filter, must_not_filter, num_results)
        return body
        
    def _as_tuple(self):
        return tuple(x._as_tuple() for x in self.musts + self.shoulds)
     
        
class LabellerQueryTemplate(CompoundQueryTemplate):
    '''
    Extension of CompoundQueryTemplate which can also store data relative
    to the performance of the query in the context of labelling.
    '''
    
    def __init__(self, *argv, **kwargs):
        super().__init__(*argv, **kwargs)
        
        self.reset()
        
    def _sanity_check(self):
        assert len(self.history_pairs) \
                == len(self.first_scores) \
                == len(self.has_results)
                
        assert len(self.first_is_match) == len(self.any_is_match)
        

    def reset(self):
        self.history_pairs = {} # {source_id': [(), (), ]}
        #self.steps = [] # Num rows labelled at which this pair was added (in case of later creation)
        self.first_scores = {} # ES scores of first result of queies
        self.has_results = {} # Whether the query returned any result
        self.first_is_match = {} # First hit is the actual match
        self.any_is_match = {} # Any of the hits is the actual match
        
        self.thresh = None
        self.precision = None
        self.recall = None
        self.inclusion_ratio = None
        self.score = 0.1 # General score of the template based on precision and recall        

    def to_dict(self):
        '''Returns a dict representation of the instance'''
        dict_ = super().to_dict()
        dict_['history_pairs'] = self.history_pairs
        dict_['first_scores'] = self.first_scores
        dict_['has_results'] = self.has_results
        dict_['first_is_match'] = self.first_is_match
        dict_['any_is_match'] = self.any_is_match
        
        dict_['thresh'] = self.thresh
        dict_['precision'] = self.precision
        dict_['recall'] = self.recall
        dict_['inclusion_ratio'] = self.inclusion_ratio
        dict_['score'] = self.score
        return dict_
    
    @classmethod
    def from_dict(cls, dict_):
        '''
        Returns an instance of the class using a representation generated by to_dict
        '''
        labeller = super(LabellerQueryTemplate, cls).from_dict(cls, dict_)
        labeller['history_pairs'] = dict_['history_pairs']
        labeller['first_scores'] = dict_['first_scores']
        labeller['has_results'] = dict_['has_results']
        labeller['first_scores'] = dict_['first_scores']
        labeller['any_is_match'] = dict_['any_is_match']
        
        labeller['thresh'] = dict_['thresh']
        labeller['precision'] = dict_['precision']
        labeller['recall'] = dict_['recall']
        labeller['inclusion_ratio'] = dict_['inclusion_ratio']
        labeller['score'] = dict_['score']
        return labeller
        
    def add_pairs(self, source_idx, pairs, first_score):
        '''Add pairs to history_pairs (pairs that were generated by this query)'''
        self.history_pairs[source_idx] = pairs
        self.has_results[source_idx] = bool(pairs)
        self.first_scores[source_idx] = first_score
    
    def add_labelled_pair(self, source_idx, labelled_pair):
        ''' Update first_is_match, any_is_match with a labelled_pair (commes after add_pairs)'''
        
        assert len(self.first_scores) == len(self.history_pairs)
        
        # If row is forgotten or no matches were found in all queries, mark as unaplicable
        if labelled_pair[1] == '__FORGET': 
            self.first_is_match[source_idx] = None # None if does not apply
            self.any_is_match[source_idx] = None # None if does not apply
        elif labelled_pair[1] == '__NO_RESULT': # Consider that no results <-> False
            self.first_is_match[source_idx] = False 
            self.any_is_match[source_idx] = False
        else:
            if self.history_pairs[source_idx]:        
                self.first_is_match[source_idx] = labelled_pair==self.history_pairs[source_idx][0]
                self.any_is_match[source_idx] = labelled_pair in self.history_pairs[source_idx]
            else:
                self.first_is_match[source_idx] = False
                self.any_is_match[source_idx] = False
            
    def unlabel(self, source_idx):
        '''Remove all information regarding labels for source_id'''
        # Compute idx to go back to
        
        # Don't touch history_pairs or first_scores or has_results
    
        if source_idx in self.first_is_match:
            del self.first_is_match[source_idx]
            del self.any_is_match[source_idx]
        else:
            assert source_idx not in self.any_is_match

        # Check whether or not to re-compute scores 


    def compute_metrics(self, t_p=0.95, t_r=0.3):
        ''' 
        Compute the optimal threshold and the associated metrics 
        
        INPUT:
            - t_p: target_precision
            - t_r: target_recall
            
        OUTPUT:
            - thresh: This refers to the threshold on the Elasticsearch score
                    for this query. If the ES score of the best hit is above
                    this threshold, it can be considered to be a match; if not, 
                    it can be assumed that it the pair is not a match (or at 
                    least very uncertain).
                    
                    The threshold is computed so as to maximize the estimated 
                    query score (based on precision and recall).
            - precision: The precision corresponding to that computed using the 
                    optimal threshold.
            - recall: The recall corresponding to that computed using the 
                    optimal threshold
            - score: The optimal score (that yields the optimal threshold)
        '''
        summaries = [{'score': self.first_scores[source_idx], 
                      'has_results': self.has_results[source_idx],
                      'first_is_match': self.first_is_match[source_idx], 
                      'any_is_match': self.any_is_match[source_idx]} \
                      for source_idx in self.any_is_match.keys()]
        
        # Filter out relevent summaries only (line not forgotten)
        summaries = [summary for summary in summaries if summary['first_is_match'] is not None]
        
        #
        num_summaries = len(summaries)
    
        # Sort summaries by score / NB: 0 score deals with empty hits
        summaries = sorted(summaries, key=lambda x: x['score'], reverse=True)
        
        # score_vect = np.array([x['_score_first'] for x in sorted_summaries])
        is_first_vect = np.array([bool(x['first_is_match']) for x in summaries])
        
        # If no matches are present
        # TODO: eliminates things from start ?
        if sum(is_first_vect) == 0:
            self.thresh = 10**3
            self.precision = 0
            self.recall = 0
            self.score = 0 
            return self.thresh, self.precision, self.recall, self.score
     
        
        rolling_precision = is_first_vect.cumsum() / (np.arange(num_summaries) + 1)
        rolling_recall = is_first_vect.cumsum() / num_summaries
        
        # TODO: WHY is recall same as precision ?
        
        # Compute score
        _f_precision = lambda x: max(x - t_p, 0) + min(t_p*(x/t_p)**3, t_p)
        _f_recall = lambda x: max(x - t_r, 0) + min(t_p*(x/t_r)**3, t_r)
        a = np.fromiter((_f_precision(xi) for xi in rolling_precision), rolling_precision.dtype)
        b = np.fromiter((_f_recall(xi) for xi in rolling_recall), rolling_recall.dtype)
        rolling_score = a*b
    
        # Find best index for threshold
        MIN_OBSERVATIONS = 4 # minimal number of values on which to comupute threshold etc...
        idx = max(num_summaries - rolling_score[::-1].argmax() - 1, min(MIN_OBSERVATIONS, num_summaries-1))
        
        # If the best score is obtained using all results (no ES score 
        # thresholding), then set threshold to ~0 
        if idx == len(summaries) - 1:
            thresh = 0.0001
        else:
            thresh = summaries[idx]['score']
        
        precision = rolling_precision[idx]
        recall = rolling_recall[idx]
        
        # Initial version
        score = rolling_score[idx]
        
        # Version taking into account the number of matches...
        
        
        # Compute 
        inclusion_ratio = sum(x['any_is_match'] for x in summaries) / len(summaries)
        
        self.thresh = thresh
        self.precision = precision
        self.recall = recall
        self.score = score
        
        return self.thresh, self.precision, self.recall, self.score


    def multiply_by_boost(self, boost_multiplier=2):
        '''Takes a single compound query and returns a list of multiple queries'''
        new_query_templates = []
        
        new_query = copy.deepcopy(self) 
        new_query_templates.append(new_query) # Keep original query
        
        og_boost_total = sum(x.boost for x in self.musts + self.shoulds)
        for level in ['shoulds', 'musts']:
            for i in range(len(self.__dict__[level])):
                new_query = copy.deepcopy(self)
                new_boost_total = og_boost_total + new_query.__dict__[level][i].boost
                new_query.__dict__[level][i].boost *= 2
                
                # Normalize total boost back to original level
                for level_2 in ['shoulds', 'musts']:
                    for j in range(len(self.__dict__[level_2])):
                        new_query.__dict__[level_2][j].boost *= og_boost_total/new_boost_total
                new_query_templates.append(new_query)
                
        return new_query_templates
    
    def multiply_by_core(self, core_queries, bool_levels):
        '''Takes a single compound query and returns a list of multiple queries'''
        
        new_query_templates = []
        
        new_query = copy.deepcopy(self)
        new_query_templates.append(new_query) # Keep original query
        
        # Multiply by core queries and bool levels
        for core_query in core_queries:
            for level in bool_levels:
                if level == 'must':
                    if all(core_query.core != q.core for q in self.musts):
                        new_query = copy.deepcopy(self)
                        new_query.add_must(core_query)
                        new_query_templates.append(new_query)
                elif level == 'should':
                    if all(core_query.core != q.core for q in self.musts):
                        new_query = copy.deepcopy(self)
                        new_query.add_should(core_query)
                        new_query_templates.append(new_query)   
                else:
                    raise ValueError('Invalid level: {0}; should be must or should')
        return new_query_templates


class CoreScorerQueryTemplate(SingleQueryTemplate):
    '''Scoer'''
    
    def __init__(self, *argv, **kwargs):
        super().__init__(*argv, **kwargs)    
        
        self.source_lens = []
        self.ref_lens = []
        self.intersect_lens = []
    
    def _sanity_check(self):
        assert len(self.source_lens) \
                == len(self.ref_lens) \
                == len(self.intersect_lens)
    
    @staticmethod
    def _analyze(es, index_name, analyzer, text):
        '''Use the current analyzer in the given index_name to analyze text'''
        ic = client.IndicesClient(es)
        if analyzer:
            return ic.analyze(index_name, body={'text': text, 'analyzer': analyzer})
        else:
            return ic.analyze(index_name, body={'text': text})
        
    def analyze_pair_items(self, es, index_name, source_item, ref_item):
        '''Count tokens for source, ref and the intersection fof both'''  
        analyzer = self.analyzer_suffix.strip('.')
        
        text_source = ' '.join(source_item[col] for col in self.source_col if source_item[col] is not None)
        tokens_source = {x['token'] for x in self._analyze(es, index_name, analyzer, text_source)['tokens']}
        
        text_ref = ' '.join(ref_item[col] for col in self.ref_col  if ref_item[col] is not None)
        tokens_ref = {x['token'] for x in self._analyze(es, index_name, analyzer, text_ref)['tokens']}

        return len(tokens_source), len(tokens_ref), len(tokens_source & tokens_ref)
    
    def add_labelled_pair_items(self, es, index_name, source_item, ref_item):
        ''''''
        l_s, l_r, l_i = self.analyze_pair_items(es, index_name, source_item, ref_item)
        
        self.source_lens.append(l_s)
        self.ref_lens.append(l_r)
        self.intersect_lens.append(l_i)
        
        self.score = sum(x > 0 for x in self.intersect_lens) / len(self.intersect_lens)
        
        
        

def _gen_suffix(columns_to_index, s_q_t_2):
    '''
    Yields suffixes to add to field_names for the given analyzers
    
    INPUT:
        - columns_to_index:
        - s_q_t_2: column or columns from ref
        
    NB: s_q_t_2: element two of the single_query_template
    
    '''
    
    if isinstance(s_q_t_2, str):
        analyzers = columns_to_index[s_q_t_2]
    elif isinstance(s_q_t_2, tuple):
        analyzers = set.union(*[columns_to_index[col] for col in s_q_t_2])
    else:
        raise ValueError('s_q_t_2 should be str or tuple (not list)')
    yield '' # No suffix for standard analyzer
    for analyzer in analyzers:
        yield '.' + analyzer


def DETUPLIFY_TODO_DELETE(arg):
    if isinstance(arg, tuple) and (len(arg) == 1):
        return arg[0]
    return arg


def _gen_all_single_query_template_tuples(match_cols, columns_to_index, bool_levels,
                                          boost_levels):
    '''Generate a list of single query tuples'''
    return list(((bool_lvl, DETUPLIFY_TODO_DELETE(x['source']), DETUPLIFY_TODO_DELETE(x['ref']), suffix, boost) \
                                       for x in match_cols \
                                       for suffix in _gen_suffix(columns_to_index, x['ref']) \
                                       for bool_lvl in bool_levels.get(suffix, ['must']) \
                                       for boost in boost_levels))

def _gen_all_query_template_tuples(match_cols, columns_to_index, bool_levels, 
                            boost_levels, max_num_levels):
    ''' 
    Generate query templates (as tuples). This will generate most combinations
    of queries that are possible using the different input.
    
    INPUT:
        match_cols: # TODO: see elsewhere
        columns_to_index: # TODO: see elsewhere
        bool_levels: list of the bool levels to generate on: ['must'], or ['must', 'should']
        boost_levels: all original boosts (see ES boost) to try. It is strongly
                    recommended to initiate with [1] (all boosts at 1) to avoid 
                    generating too many query templates
        max_num_levels: The maximal number of single queries in the compound
                    queries that are generated
                    
    OUTPUT:
        all_query_templates: A list of tuples representing compound queries
    '''
    single_queries = _gen_all_single_query_template_tuples(match_cols, 
                                    columns_to_index, bool_levels, boost_levels)
    all_query_templates = list(itertools.chain(*[list(itertools.combinations(single_queries, x)) \
                                        for x in range(2, max_num_levels+1)][::-1]))
    # Queries must contain at least two distinct columns (unless one match is mentionned)
    if len(match_cols) > 1:
        all_query_templates = list(filter(lambda query: len(set((x[1], x[2]) for x in query)) >= 2, \
                                    all_query_templates))

    all_query_templates = list(filter(lambda query: 'must' in [s_q_t[0] for s_q_t in query], \
                                all_query_templates))
    return all_query_templates


"""

"""


"""
Initialise core queries (list)
On match:
    - Analyze source and ref Using the query analyzer
    - Add number Of common tokens 
"""





class Labeller():
    '''
    Labeller object that stores previous labels and generates new matches
    for the user to label
    
    0.1) Initiate queries/ metrics/ history
    0.2) Read first row
    
    1) Perform queries / update history of hits
    2) Generate pairs to propose (based on sorted queries)
    3...) Until row is over: User inputs label
    4) Update metrics and history 
    5) Sort queries
    6) Gen new row and back to 1)
        
        
    HOW TO USE:
    
    >>> labeller = Labeller()
    
    >>> for x in range(100):
            to_display = labeller.to_emit()
            labeller.update('y')
    
    >>> labeller.export_params()    
        
    '''
    NUM_RESULTS = 3
    
    MAX_NUM_SAMPLES = 100
    
    VALID_ANSWERS = {'yes': 'y',
                 'y': 'y',
                 '1': 'y',
                 
                 'no': 'n',
                 'n': 'n',
                 '0': 'n',
                 
                    #                 'uncertain': 'u',
                    #                 'u': 'u',
                 
                 'forget_row': 'f',
                 'f': 'f',
                 
                 'previous': 'p',
                 'p': 'p',

                 'quit': 'q'             
                 }

    #min_precision_tab = [(20, 0.7), (10, 0.5), (5, 0.3)]
    #max_num_keys_tab = [(20, 10), (10, 50), (7, 200), (5, 500), (0, 1000)]
    #num_results_labelling = 3
    
    MAX_NUM_LEVELS = 3 # Number of match clauses
    BOOL_LEVELS = {'.integers': ['must', 'should'], 
                   '.city': ['must', 'should']}
    BOOST_LEVELS = [1]
    
    # Target precision and target recall
    t_p = 0.95
    t_r = 0.3    
    
    @staticmethod
    def _dedupe_source(source, match_cols):
        '''Dedupe source on columns used for matching to avoid multiple search'''
         # Dedupe source on matching columns
        source_cols_for_match = set()
        for match in match_cols:
            if isinstance(match['source'], str):
                source_cols_for_match.add(match['source'])
            else:
                source_cols_for_match.update(match['source'])
        source_cols_for_match = list(source_cols_for_match)       
        
        smaller_source = source.drop_duplicates(subset=source_cols_for_match)
        
        return smaller_source
    
    def __init__(self, source, ref_index_name, 
                 match_cols, columns_to_index, 
                 certain_column_matches=None, 
                 must={}, must_not={}):
        '''        
        INPUT: 
          - source: pandas DataFrame containing input data to match
          - ref_index_name: name of the Elasticsearch index to search in
          - match_cols: Indication of columns to try to match:
              EX: [{'source': ['city'], 'ref': ['CITY_ADD']},
                   {'source': ['first_name', 'last_name'], 'ref': ['FULL_NAME']}]
          - columns_to_index: What columns were indexed and how:
              EX:  {
                    'CITY_ADD': {},
                    'FULL_NAME': {'french', 'integers', 'n_grams', 'city'}
                    }
          - must_filters
        '''
            
        self.source = self._dedupe_source(source, match_cols)
        
        self.ref_index_name = ref_index_name
        self.match_cols = match_cols      
        self.columns_to_index = columns_to_index
        
        self._init_es()
        
        self.labelled_pairs = [] # Flat list of labelled pairs
        self.labels = {} # Flat list of labels
        self.num_rows_labelled = [] # Flat list: at given label, how many were labelled. NB: starts at 0, becomes 1 at first yes/forgotten, or when next_row
        self.num_positive_rows_labelled = [] # Flat list: at given label, how_many were matches
        self.labelled_pairs_match = [] # For each row, the resulting match: (A, B) / no-match: (A, None) or forgotten: None
         
        
        self.must_filters = must
        self.must_not_filters = must_not
        
           
        self._init_queries(match_cols, columns_to_index) # creates self.current_queries
        self._init_core_queries(match_cols, columns_to_index) # creates self.single_core_queries
        #self._init_history() # creates self.history
        
{'fraction_non_alphanumeric': 0.10544712007211317, 'fraction_numerical': 0.006199525377582369, 'mean_word_length': 2.0173179396092364, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 46, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6316990', 'n_tokens_mistral': 16141, 'n_tokens_neox': 15389, 'n_words': 4176}