from elasticsearch_client import elasticsearch_client
from snowflake_client import snowflake_client
from utils import time_function, time_n_calls_v2
import example_data
import elastic_poc
import snowflake_poc
import matplotlib.pyplot as plt
import numpy as np

es_client = elasticsearch_client()
sf_client = snowflake_client()


def compare(operation, args):
    print('comparing', operation)
    snowflake_times = []
    elastic_times = []

    for arg in args:
        print('\t', arg)

        snowflake_time = time_function(
            'snowflake', operation, arg, sf_client)
        elastic_time = time_function('elastic', operation, arg, es_client)

        snowflake_times.append(snowflake_time)
        elastic_times.append(elastic_time)

    return snowflake_times, elastic_times


def compare_v2(title, functions, test_cases, num_iterations=5):
    print('-' * 50)
    print(title)
    print('-' * 50)
    results = {}
    clients = {'sf_client': sf_client, 'es_client': es_client}
    print('\n' * 2)

    for test in test_cases:
        print('Testing: {}'.format(test))

        for function in functions:
            print('Executing:{}.{}'.format(
                function.__module__, function.__name__))
            time_req = time_n_calls_v2(
                function, {**test, **clients}, n=num_iterations)
            results.setdefault('{}.{}'.format(
                function.__module__, function.__name__), []).append(time_req)

        print('\n' * 2)

    return results


def graphit(results, labels, xlabel='', filename='', title=''):
    width = 0.2
    fig, ax = plt.subplots()
    x = np.arange(len(labels))
    plots = []
    items = list(results.items())
    total_items = len(items)

    if total_items % 2 == 0:
        mid = total_items // 2
        for key, val in items[:mid]:
            plots.append(ax.bar(x - width/total_items, val, width, label=key))

        for key, val in items[mid:]:
            plots.append(ax.bar(x + width/total_items, val, width, label=key))
    else:
        mid = total_items // 2

        for key, val in items[:mid]:
            plots.append(ax.bar(x - width, val, width, label=key))

        plots.append(ax.bar(x, items[mid][1], width, label=items[mid][0]))

        for key, val in items[mid + 1:]:
            plots.append(ax.bar(x + width, val, width, label=key))

    ax.set_ylabel('avg query time(n iterations) in seconds')
    ax.set_xlabel(xlabel)

    ax.set_title(title)
    ax.set_xticks(x)
    ax.legend()

    fig.tight_layout()
    plt.savefig(filename)


def main():
    playlist_isrc_pairs = [
        {
            'isrc': isrc,
            'playlist_id': example_data.todays_top_hits_playlist_id
        } for isrc in list(example_data.isrcs.values())
    ]

    """
    res = compare_v2('Hybrid vs Pure ES vs Pure SF: Placements + Streams',
                     [
                         snowflake_poc.playlist_positions_over_time_with_streams,
                         elastic_poc.placements_and_streams_overtime_hybrid,
                         elastic_poc.placements_and_streams_overtime_es
                     ],
                     playlist_isrc_pairs,
                     num_iterations=10)

    graphit(res, [name for name in list(example_data.isrcs.keys())],
            title='Pure SF vs Hybrid ES + SF vs Pure ES: Placements + Streams',
            xlabel='tracks in Todays top hits',
            filename='hybrid_vs_es_vs_sf_placements_overtime.pdf')

    res = compare_v2(
        'Pure SF vs Hybrid vs Pure ES: Current placements+streams in playlist',
        [
            snowflake_poc.current_placements_in_playlist_with_streams,
            elastic_poc.current_placements_in_playlist_with_streams_hybrid,
            elastic_poc.current_placements_in_playlist_with_streams
        ],
        [{'playlist_id': playlist}
            for playlist in example_data.top_playlist_ids.values()][:10],
        num_iterations=10
    )

    graphit(res,
            [name for name in list(example_data.top_playlist_ids.keys())][:10],
            title='Pure SF vs SF+ES vs Pure ES: Current placements',
            xlabel='Top playlists',
            filename='hybrid_vs_es_vs_sf_current_placements.pdf')
    """
    res = compare_v2(
        'SF vs ES: Current placements given isrc',
        [
            snowflake_poc.current_placements_given_isrc,
            elastic_poc.current_placements_given_isrc_hybrid,
            elastic_poc.current_placements_given_isrc
        ],
        [{'isrc': isrc} for isrc in example_data.isrcs.values()][:3],
        num_iterations=4
    )

    graphit(res, [name for name in list(example_data.isrcs.values())][:3],
            title='SF vs ES: current placements given isrc', xlabel='Tracks', filename='sf_vs_es_current_placements.pdf')


if __name__ == '__main__':
    main()