import asyncio
import math
import random
import os
import time

import matplotlib.pyplot as plt

from traderion2.bot import TraderionBot, BID, ASK


# Use a very large LOOP_SLEEP so the main loop gets called only once at the beginning.
LOOP_SLEEP = 2 * 60 * 60  # 2 hours

PLOTS_DIR = 'plots/'

LATENCY_NUM_RUNS = 500

THROUGHPUT_SECONDS = 60
THROUGHPUT_MAX_PENDING_QUERIES = 200

SERVER_COOLDOWN_SECONDS = 60


class StressBot(TraderionBot):
    """
    A bot that stress tests the simulator, measuring latency and throughput.

    It prints statistics to stdout and exports plots in PLOTS_DIR directory.
    """

    def __init__(self, username, password, room_id):
        super().__init__(username, password, room_id, LOOP_SLEEP)

        self.swift_id = self.api.get_swifts()[0]['swift_id']

    async def main_loop(self):
        await self.test_latency()
        await self.test_throughput()

        self.stop()

    async def test_latency(self):
        latency_tester = LatencyTester()

        await asyncio.sleep(SERVER_COOLDOWN_SECONDS)
        print('Measuring hit_price latency...')
        await latency_tester.test_async_func(self.hit_price_func, LATENCY_NUM_RUNS,
                                             plot_name='latencies_eb_hit_price')

        await asyncio.sleep(SERVER_COOLDOWN_SECONDS)
        print('Measuring eb_add_order latency...')
        await latency_tester.test_async_func(self.eb_add_order_func, LATENCY_NUM_RUNS,
                                             plot_name='latencies_eb_add_order')

        await self.api.cancel_all_orders(self.swift_id)

    async def test_throughput(self):
        throughput_tester = ThroughputTester()

        await asyncio.sleep(SERVER_COOLDOWN_SECONDS)
        print('Measuring hit_price throughput...')
        await throughput_tester.test_async_func(self.hit_price_func, THROUGHPUT_SECONDS,
                                                plot_name='throughput_eb_hit_price')

        await asyncio.sleep(SERVER_COOLDOWN_SECONDS)
        print('Measuring eb_add_order throughput...')
        await throughput_tester.test_async_func(self.eb_add_order_func, THROUGHPUT_SECONDS,
                                                plot_name='throughput_eb_add_order')

        await self.api.cancel_all_orders(self.swift_id)

    @property
    def hit_price_func(self):
        async def hit_price():
            await self.api.hit_price(self.swift_id, random.choice([BID, ASK]), 10)

        return eb_hit_price

    @property
    def eb_add_order_func(self):
        async def eb_add_order():
            direction = random.choice([BID, ASK])
            market_prices = self.api.get_market_prices(self.swift_id)
            price = {BID: market_prices['bid'], ASK: market_prices['ask']}[direction]
            await self.api.eb_add_order(self.swift_id, direction, 1, price)

        return eb_add_order


def plot_histogram(vals, xlabel, ylabel, plot_name):
    plt.hist(vals, bins='auto', color='#0504aa', alpha=0.7)
    plt.grid(axis='y', alpha=0.75)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(plot_name)

    if not os.path.exists(PLOTS_DIR):
        os.mkdir(PLOTS_DIR)

    plt.savefig(os.path.join(PLOTS_DIR, plot_name), dpi=600)
    plt.clf()


class LatencyTester(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.latencies = []
        self.min_lat = self.max_lat = self.avg_lat = None

    async def test_async_func(self, func, num_runs, plot_name=None):
        self.reset()

        for _ in range(num_runs):
            start_time = time.time()
            await func()
            end_time = time.time()

            latency = (end_time - start_time) * 1000.
            self.latencies.append(latency)

        self.compute_stats()
        self.print_stats()

        if plot_name is not None:
            self.plot_latency_histogram(plot_name)

    def compute_stats(self):
        self.avg_lat = sum(self.latencies) / len(self.latencies)
        self.min_lat = min(self.latencies)
        self.max_lat = max(self.latencies)

    def print_stats(self):
        print('Latency(ms): %.2f | %.2f | %.2f (min | avg | max)\n' % (self.min_lat, self.avg_lat, self.max_lat))

    def plot_latency_histogram(self, plot_name):
        plot_histogram(self.latencies, xlabel='Latency (ms)', ylabel='Count', plot_name=plot_name)


class ThroughputTester(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.throughputs = []
        self.pending_queries = 0
        self.answered_queries = 0

        self.should_finish = False
        self.finished_future = asyncio.Future()

    async def test_async_func(self, func, num_seconds, plot_name=None):
        self.reset()

        for _ in range(THROUGHPUT_MAX_PENDING_QUERIES):
            asyncio.ensure_future(self.send_query(func))

        await asyncio.sleep(1)  # Give some time to the server to get spinning

        last_answered_queries = self.answered_queries
        for _ in range(num_seconds):
            await asyncio.sleep(1)
            qps = self.answered_queries - last_answered_queries
            self.throughputs.append(qps)

            last_answered_queries = self.answered_queries

        self.should_finish = True

        self.compute_stats()
        self.print_stats()

        if plot_name is not None:
            self.plot_throughput_histogram(plot_name)

        await self.finished_future

    async def send_query(self, func):
        self.pending_queries += 1
        await func()
        self.pending_queries -= 1
        self.answered_queries += 1

        if not self.should_finish:
            asyncio.ensure_future(self.send_query(func))
        elif self.pending_queries == 0:
            self.finished_future.set_result(None)

    def compute_stats(self):
        self.avg_qps = sum(self.throughputs) / len(self.throughputs)
        self.min_qps = min(self.throughputs)
        self.max_qps = max(self.throughputs)

    def print_stats(self):
        print('Throughput(qps): %d | %d | %d (min | avg | max)\n' % (self.min_qps, self.avg_qps, self.max_qps))

    def plot_throughput_histogram(self, plot_name):
        plot_histogram(self.throughputs, xlabel='Throughput (qps)', ylabel='Count', plot_name=plot_name)


def main():
    username = 'YOUR_USERNAME'
    password = 'YOUR_PASSWORD'
    room_id = 'ROOM_ID'

    stress_bot = StressBot(username, password, room_id)
    loop = asyncio.get_event_loop()
    loop.run_until_complete(stress_bot.run())


if __name__ == '__main__':
    main()
