from itertools import cycle
from typing import Dict, Union

from matplotlib import pyplot as plt

from graphs import TestStore
from structure import StandardTest


def plot_iperf_results(
        store: TestStore,
        series: Dict[str, StandardTest],
        fast_tests: bool,
        title: str = None,
        direction='both',
        error_bars=True,
        filename=None,
):
    if direction == 'both':
        plot_iperf_results(store, series, fast_tests, title, 'inbound', error_bars, filename)
        plot_iperf_results(store, series, fast_tests, title, 'outbound', error_bars, filename)
        return

    if filename in ['png', 'eps']:
        filename = 'output/{}{}{}.{}'.format(
            'I' if direction == 'inbound' else 'O',
            'E' if error_bars else '',
            ''.join(['S{}-{}'.format(i, x.name()) for (i, x) in enumerate(series.values())]),
            filename,
        )
        print(filename)

    series = {
        k: (store.get_inbound(v) if direction == 'inbound' else store.get_outbound(v))
        for (k, v) in series.items()
    }

    fig = plt.figure()
    axes = fig.add_axes([0, 0, 1, 1])

    if title is not None:
        axes.set_title(title, pad=20.0 if True in [len(x.test.events) > 0 for x in series.values()] else None)

    axes.set_ylabel('Throughput (Mbps)')

    for k, v in series.items():
        axes.bar(
            k,
            v.bandwidth_mean() / 1e6,
            yerr=1.5 * v.bandwidth_standard_deviation() / 1e6,
            width=0.4,
            capsize=15,
        )

    if fast_tests:
        fig.text(0.95, 0.05, 'Draft', fontsize=50, color='gray', ha='right', va='bottom', alpha=0.5)

    if filename is not None:
        fig.savefig(filename, bbox_inches='tight', pad_inches=0.3)


def plot_iperf_results_time(
        store: TestStore,
        series: Dict[str, StandardTest],
        fast_tests: bool,
        title: str = None,
        direction='both',
        error_bars_x=False,
        error_bars_y=True,
        filename=None,
        start_at_zero=True,
):
    if direction == 'both':
        plot_iperf_results_time(store, series, fast_tests, title, 'outbound', error_bars_x, error_bars_y, filename,
                                start_at_zero)
        plot_iperf_results_time(store, series, fast_tests, title, 'inbound', error_bars_x, error_bars_y, filename,
                                start_at_zero)
        return

    if filename in ['png', 'eps']:
        filename = 'output/T{}{}{}{}.{}'.format(
            'I' if direction == 'inbound' else 'O',
            'Ex' if error_bars_x else '',
            'Ey' if error_bars_y else '',
            ''.join(['S{}-{}'.format(i, x.name()) for (i, x) in enumerate(series.values())]),
            filename,
        )
        print(filename)

    series = {
        k: (store.get_inbound(v) if direction == 'inbound' else store.get_outbound(v))
        for (k, v) in series.items()
    }

    cycol = cycle('brgy')

    fig = plt.figure()
    axes = fig.add_axes([0, 0, 1, 1])

    if title is not None:
        axes.set_title(title, pad=20.0 if True in [len(x.test.events) > 0 for x in series.values()] else None)

    axes.set_xlabel('Time (s)')
    axes.set_ylabel('Throughput (Mbps)')

    for k, v in series.items():
        data = v.interval_means()

        axes.errorbar(
            data.keys(),
            [x / 1e6 for x in data.values()],
            xerr=(
                [x[0] for x in v.interval_time_ranges().values()],
                [x[1] for x in v.interval_time_ranges().values()]) if error_bars_x else None,
            yerr=[x * 1.5 / 1e6 for x in v.interval_standard_deviations().values()] if error_bars_y else None,
            capsize=3,
            ecolor='grey',
            color=next(cycol),
            label=k,
        )

    legend = axes.legend()

    if start_at_zero:
        axes.set_ylim(bottom=0)
        axes.set_xlim(left=0)

    if fast_tests:
        fig.text(0.95, 0.05, 'Draft', fontsize=50, color='gray', ha='right', va='bottom', alpha=0.5)

    if filename is not None:
        fig.savefig(filename, bbox_extra_artists=(legend,), bbox_inches='tight', pad_inches=0.3)