%load_ext autoreload
%autoreload 2

import importlib
import os
import sys

from elasticsearch import Elasticsearch
from skopt.plots import plot_objective

# project library
sys.path.insert(0, os.path.abspath('..'))

import qopt
importlib.reload(qopt)

from qopt.notebooks import evaluate_mrr100_dev, optimize_query_mrr100
from qopt.optimize import Config

# use a local Elasticsearch or Cloud instance (https://cloud.elastic.co/)
# es = Elasticsearch('http://localhost:9200')
es = Elasticsearch('http://35.246.195.44:9200')

# set the parallelization parameter `max_concurrent_searches` for the Rank Evaluation API calls
# max_concurrent_searches = 10
max_concurrent_searches = 30

index = 'msmarco-document'
template_id = 'combined'

"""
Explanation: Appendix B: Combining queries
In this example, we'll walk through the process of building a more complex query for the MSMARCO Document dataset. This assumes that you are familiar with the other query tuning notebooks. We'll be using the cross_fields and best_fields queries and the optimal parameters found as the foundation on which we build a more complex query.
As with the previous notebook and in accordance with the MSMARCO Document ranking task, we'll continue to use MRR@100 on the dev dataset for evaluation and comparison with other approaches.
End of explanation
"""


cross_fields_params = {
    'operator': 'OR',
    'minimum_should_match': 50,
    'tie_breaker': 0.25,
    'url|boost': 1.0129720302556104,
    'title|boost': 5.818478716515356,
    'body|boost': 3.736613263685484,
}

best_fields_params = {
    'tie_breaker': 0.3936135232328522,
    'url|boost': 0.0,
    'title|boost': 8.63280262513067,
    'body|boost': 10.0,
}

"""
Explanation: Combining cross_fields and best_fields
Based on previous tuning, we have the following optimal parameters for each multi_match query type.
End of explanation
"""


def prefix_keys(d, prefix):
    return {f'{prefix}{k}': v for k, v in d.items()}

# prefix key of each sub-query
# add default boosts
all_params = {
    **prefix_keys(cross_fields_params, 'cross_fields|'),
    'cross_fields|boost': 1.0,
    **prefix_keys(best_fields_params, 'best_fields|'),
    'best_fields|boost': 1.0,
}
all_params

"""
Explanation: We've seen the process to optimize field boosts on two different multi_match queries but it would be interesting to see if combining them in some way might actually result in even better MRR@100. Let's give it a shot and find out.
Side note: Combining queries where each sub-query is always executed may improve relevance but it will hurt performance and the query times will be quite a lot higher than with a single, simpler query. Keep this in mind when building complex queries for production!
End of explanation
"""


%%time

_ = evaluate_mrr100_dev(es, max_concurrent_searches, index, template_id, params=all_params)

"""
Explanation: Baseline evaluation
End of explanation
"""


%%time

_, _, final_params_boosts, metadata_boosts = optimize_query_mrr100(es, max_concurrent_searches, index, template_id,
    config_space=Config.parse({
        'num_iterations': 30,
        'num_initial_points': 15,
        'space': {
            'cross_fields|boost': { 'low': 0.0, 'high': 5.0 },
            'best_fields|boost': { 'low': 0.0, 'high': 5.0 },
        },
        'default': all_params,
    }))

_ = plot_objective(metadata_boosts, sample_source='result')

"""
Explanation: Query tuning
Here we'll just tune the boosts for each sub-query. Note that this takes twice as long as tuning individual queries because we have two queries combined.
End of explanation
"""


%%time

_ = evaluate_mrr100_dev(es, max_concurrent_searches, index, template_id, params=final_params_boosts)

"""
Explanation: Seems that there's not much to tune here, but let's keep going.
End of explanation
"""


import matplotlib.pyplot as plt
from itertools import chain

from qopt.notebooks import ROOT_DIR
from qopt.search import temporary_search_template, search_template
from qopt.trec import load_queries_as_tuple_list, load_qrels

def collect_scores():
    
    def _search(template_id, query_string, params, doc_id):
        res = search_template(es, index, template_id, query={
            'id': 0,
            'params': {
                'query_string': query_string,
                **params,
            },
        })
        return [hit['score'] for hit in res['hits'] if hit['id'] == doc_id]

    queries = load_queries_as_tuple_list(os.path.join(ROOT_DIR, 'data', 'msmarco-document-sampled-queries.1000.tsv'))
    qrels = load_qrels(os.path.join(ROOT_DIR, 'data', 'msmarco', 'document', 'msmarco-doctrain-qrels.tsv'))
    template_file = os.path.join(ROOT_DIR, 'config', 'msmarco-document-templates.json')
    size = 100
    
    cross_field_scores = []
    best_field_scores = []
    
    with temporary_search_template(es, template_file, 'cross_fields', size) as cross_fields_template_id:
        with temporary_search_template(es, template_file, 'best_fields', size) as best_fields_template_id:
            for query in queries:
                doc_id = list(qrels[query[0]].keys())[0]
                
                cfs = _search(cross_fields_template_id, query[1], cross_fields_params, doc_id)
                bfs = _search(best_fields_template_id, query[1], best_fields_params, doc_id)

                # keep just n scores to make sure the lists are the same length
                length = min(len(cfs), len(bfs))
                cross_field_scores.append(cfs[:length])
                best_field_scores.append(bfs[:length])

    return cross_field_scores, best_field_scores

cfs, bfs = collect_scores()

# plot scores
cfs_flat = list(chain(*cfs))
bfs_flat = list(chain(*bfs))

plt.scatter(cfs_flat, bfs_flat)
plt.show()

"""
Explanation: So that's the same as without tuning. What's going on?
Debugging
Plot scores from each sub-query to determine why we don't really see an improvement over individual queries.
End of explanation
"""


