import asyncio
import time as tm
from datetime import datetime
from os.path import join
from typing import List

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from appdirs import user_cache_dir
from matplotlib.ticker import ScalarFormatter

from backtesterrb30.libs.communication_broker.broker_base import BrokerBase
from backtesterrb30.libs.interfaces.python_backtester.data_finish import (
    CustomChart,
    DataFinish,
)
from backtesterrb30.libs.interfaces.python_backtester.data_start import DataStart
from backtesterrb30.libs.interfaces.python_backtester.debug_breakpoint import (
    DebugBreakpoint,
)
from backtesterrb30.libs.interfaces.python_backtester.last_feed import LastFeed
from backtesterrb30.libs.interfaces.python_backtester.positions import Position
from backtesterrb30.libs.interfaces.python_backtester.trade import Trade
from backtesterrb30.libs.interfaces.utils.config import Config
from backtesterrb30.libs.interfaces.utils.data_schema import DataSchema
from backtesterrb30.libs.interfaces.utils.data_symbol import DataSymbol
from backtesterrb30.libs.utils.list_of_services import SERVICES, SERVICES_ARRAY
from backtesterrb30.libs.utils.service import Service


class Backtester(Service):
    downloaded_data_path = user_cache_dir("rb30_cache")
    _broker: BrokerBase

    def __init__(
        self, config: Config, data_schema: DataSchema, loop=None, logger=print
    ):
        super().__init__(config, logger)
        self.config: Config = config
        self.__loop = loop
        self.__custom_event_loop = False
        if self.__loop is None:
            self.__loop = asyncio.get_event_loop()
            self.__custom_event_loop = True
        self.data_schema: DataSchema = data_schema
        for dat in self.data_schema.data:
            dat.additional_properties["position"] = None
            dat.additional_properties["chart_data_frame"] = None

        # self.trading_instruments_charts = []
        self.cumulated_money_chart = []
        # self.positions: List[Position] = []
        self.__last_feed: LastFeed = {}
        self.__last_timestamp: int = 0

        self.__backtest_start_time = 0
        self.__fig = None
        self.__chart_displayed = False

    # # override
    # def _asyncio_loop(self, loop: asyncio.AbstractEventLoop):
    #     self._create_listeners(loop)
    #     loop.create_task(self.__update_chart())

    # # override
    # def _handle_zmq_message(self, message):
    #     pass

    def _loop(self):
        # self._broker.run()
        self._broker.create_listeners(self.__loop)
        self.__loop.create_task(self.__update_chart())
        if self.__custom_event_loop:
            self.__loop.run_forever()
            self.__loop.close()

    # def _send(self, service: SERVICES, msg: str, *args):
    #     self._broker.send(service, msg, *args)

    def _configure(self):
        super()._configure()
        self._broker.register("trade", self.__trade_event)
        self._broker.register("data_finish", self.__data_finish_event)
        self._broker.register("close_all_trades", self.__close_all_trades_event)
        self._broker.register("data_start", self.__data_start_event)
        self._broker.register("debug_breakpoint", self.__debug_breakpoint_event)
        self._broker.register("last_feed", self.__last_feed_event)

    # def _trigger_event(self, event):
    #     pass

    async def __stop_all_services(self):
        for service in SERVICES_ARRAY:
            if service != self.name:
                await self._broker.send(getattr(SERVICES, service), "stop")
        self._broker.stop()

    def __axis_format(self, ax, title):
        ax.set_title(title)
        ax.set_axisbelow(True)
        ax.yaxis.grid(color="gray", linestyle="dashed")

    async def __print_charts(self, custom_charts: List[CustomChart]):
        plt.close()

        # plot trades chart
        # position: Position = self.positions[0]
        number_of_custom_charts = 0
        if custom_charts is not None:
            number_of_custom_charts = len(
                [ch for ch in custom_charts if not ch.display_on_price_chart]
            )

        chartable_symbols = [
            sym for sym in self.data_schema.data if sym.display_chart_in_summary
        ]
        axs_number_used = 0

        # prepare axes
        self.__fig, axs = plt.subplots(
            nrows=1 + len(chartable_symbols) + number_of_custom_charts,
            ncols=1,
            sharex=True,
            figsize=(13, 13),
        )
        if not isinstance(axs, np.ndarray):
            # if only 1 axis its returned not in array
            axs = [axs]
        ax = None
        # plot instrment charts
        for sym in chartable_symbols:
            position: Position = sym.additional_properties["position"]
            # if position:
            if self.__last_timestamp is not None:
                main_chart = sym.additional_properties["chart_data_frame"].loc[
                    sym.additional_properties["chart_data_frame"]["timestamp"]
                    <= self.__last_timestamp
                ]
            ax = main_chart.plot(
                x="timestamp", y="price", kind="line", ax=axs[axs_number_used]
            )
            axs_number_used += 1
            self.__axis_format(ax, sym.symbol + " - " + sym.historical_data_source)
            if self.data_schema.log_scale_valuation_chart:
                ax.yaxis.set_major_formatter(ScalarFormatter())
            if position:
                normalized_quants = self.__normalize(
                    [abs(trade[2]) for trade in position.trades], (5, 15)
                )
                for trade, quant in zip(position.trades, normalized_quants):
                    ax.plot(
                        trade[0], trade[1], ".g" if trade[2] > 0 else ".r", ms=quant
                    )

        if len(self.cumulated_money_chart) > 0:
            # plot money chart
            money_df = pd.DataFrame(
                self.cumulated_money_chart, columns=["timestamp", "income"]
            )
            ax = money_df.plot(
                x="timestamp",
                y="income",
                kind="line",
                ax=axs[axs_number_used],
                sharex=ax,
            )
            axs_number_used += 1

        # plot custom charts
        if custom_charts is not None:
            for i, ch in enumerate(custom_charts):
                chart = [[c.timestamp, c.value] for c in ch.chart]
                custom_df = pd.DataFrame(chart, columns=["timestamp", ch.name])
                if ch.display_on_price_chart:
                    custom_df.plot(
                        x="timestamp",
                        y=ch.name,
                        kind="line",
                        ax=axs[0],
                        sharex=ax,
                        color=ch.color,
                    )
                else:
                    ax = custom_df.plot(
                        x="timestamp",
                        y=ch.name,
                        kind="line",
                        ax=axs[axs_number_used],
                        sharex=ax,
                        color=ch.color,
                    )
                    axs_number_used += 1
                    self.__axis_format(ax, ch.name)
                    if ch.log_scale:
                        ax.yaxis.set_major_formatter(ScalarFormatter())
        if not self.__chart_displayed:
            plt.ion()
            plt.show(block=False)
            self.__chart_displayed = True

    def __add_position(self, data_symbol: DataSymbol) -> Position:
        pos = Position()
        data_symbol.additional_properties["position"] = pos
        self.__update_last_instrument_prices()
        return pos

    async def __update_chart(self):
        while True:
            if self.__fig is not None:
                try:
                    self.__fig.canvas.draw()
                    self.__fig.canvas.flush_events()
                except Exception:
                    pass
            await asyncio.sleep(0.1)

    def __normalize(
        self, x, newRange=(0, 1)
    ):  # x is an array. Default range is between zero and one
        if len(x) == 0:
            return []
        xmin, xmax = np.min(x), np.max(x)  # get max and min from input array
        if xmin != xmax:
            norm = (x - xmin) / (xmax - xmin)  # scale between zero and one

            if newRange == (0, 1):
                return norm  # wanted range is the same as norm
            elif newRange != (0, 1):
                return (
                    norm * (newRange[1] - newRange[0]) + newRange[0]
                )  # scale to a different range.
            # add other conditions here. For example, an error messag
        return np.ones_like(x) * 15

    async def __recalculate_positions(self):
        positions: List[Position] = [
            elem.additional_properties["position"] for elem in self.data_schema.data
        ]
        positions = [pos for pos in positions if pos is not None]
        for pos in positions:
            pos.position_outcome = (
                -pos.buy_summary_cost
                - pos.sell_summary_cost
                + pos.number_of_actions * pos.last_instrument_price
            )
        current_capital = sum([pos.position_outcome for pos in positions])
        current_invested = sum(
            [
                abs(pos.number_of_actions) * pos.last_instrument_price
                for pos in positions
            ]
        )
        self.cumulated_money_chart.append([self.__last_timestamp, current_capital])
        await self._broker.send(
            SERVICES.python_executor, "set_current_capital_event", current_capital
        )
        await self._broker.send(
            SERVICES.python_executor, "set_current_invested_event", current_invested
        )

    def __get_data_symbol(self, symbol: str, source: str):
        arr = [
            sym
            for sym in self.data_schema.data
            if sym.symbol == symbol and sym.historical_data_source == source
        ]
        if len(arr) == 0:
            raise Exception("No data symbol found")
        if len(arr) > 1:
            raise Exception("More than one symbol found")
        return arr[0]

    async def __trade(self, trade: Trade):
        symbol: DataSymbol = self.__get_data_symbol(trade.symbol, trade.source)
        position = symbol.additional_properties["position"]
        if not position:
            position = self.__add_position(symbol)
        if trade.price == 0 and trade.timestamp == 0:
            trade.price = position.last_instrument_price
            trade.timestamp = self.__last_timestamp
        if trade.price <= 0 or trade.timestamp <= 0:
            self._log(
                "Cannot make trade with this price and timestamp:",
                trade.price,
                ";",
                trade.timestamp,
            )
            return
        self._log(
            f"Trade for [{symbol.symbol}]: time={datetime.utcfromtimestamp(trade.timestamp/1000)}, | value={trade.value}, price={trade.price}"
        )
        position.trades.append([trade.timestamp, trade.price, trade.value])
        # print('trade price', trade.price)
        position.number_of_actions += trade.value / trade.price
        if trade.value > 0:
            position.buy_summary_cost += trade.value
        else:
            position.sell_summary_cost += trade.value
        await self.__recalculate_positions()

    async def __print_summary(
        self,
        custom_charts: List[CustomChart],
        display_charts=True,
        # file_names: List[str],
        # main_instrument_price: float,
        # last_timestamp: Union[int, None] = None
    ):
        self._log("==========================")
        self._log(
            ("BREAKPOINT " if self.__last_timestamp is not None else "") + "SUMMARY"
        )
        finish_time = tm.time()
        time_of_backtest = finish_time - self.__backtest_start_time
        self._log("time of backtest:", round(time_of_backtest, 2), "[s]")
        # self._log('number of trades:', len(self.trades))
        # self._log('buy_summary_cost:', self.buy_summary_cost)
        # self._log('sell_summary_cost:', self.sell_summary_cost)
        # self._log('number of unrealized actions:', self.number_of_actions)
        # self._log('biggest investment: ', self.biggest_investment)
        # self._log('actual price:', main_instrument_price)
        # income = - self.buy_summary_cost - self.sell_summary_cost + self.number_of_actions * main_instrument_price
        if len(self.cumulated_money_chart) > 0:
            self._log("income:", self.cumulated_money_chart[-1][1])
        self._log("==========================")
        self._log("")
        if display_charts:
            await self.__print_charts(custom_charts)

    def __load_instrument_chart(instrument):
        loaded_df = None

        return loaded_df

    # COMMANDS

    async def __data_finish_event(self, finish_params):
        finish_params = DataFinish(**finish_params)
        self._log("")
        self._log("====================================================")
        self._log("BACKTEST FINISHED")
        await self.__print_summary(finish_params.custom_charts)
        self._log("====================================================")
        plt.show(block=True)
        await self.__stop_all_services()

    async def __trade_event(self, msg):
        trade: Trade = Trade(**msg)
        await self.__trade(trade)

    async def __last_feed_event(self, msg):
        # print('last feed event')
        self.__last_feed = LastFeed(**msg)
        self.__last_timestamp = self.__last_feed.last_feed[0]
        self.__update_last_instrument_prices()

    def __update_last_instrument_prices(self):
        for data_symbol, last_feed in zip(
            self.data_schema.data, self.__last_feed.last_feed[1:]
        ):
            position: Position = data_symbol.additional_properties["position"]
            if position:
                if isinstance(last_feed, (float, int)):
                    position.last_instrument_price = last_feed
                    continue
                if isinstance(last_feed, list):
                    print("feed is 3d")
                    if isinstance(last_feed[0], (float, int)):
                        position.last_instrument_price = last_feed[0]
                    continue

                raise Exception("Error during setting last instrument price")

    async def __close_all_trades_event(self):
        self._log("Received close all trades command NOT IMPLEMENTED")

        # TODO

    async def __data_start_event(self, start_params):
        start_params = DataStart(**start_params)
        self.__backtest_start_time = start_params.start_time

        for element in start_params.file_names:
            dfs = []
            for file in element.files:
                df = pd.read_csv(
                    join(self.downloaded_data_path, file.to_filename()),
                    index_col=None,
                    header=None,
                    names=["timestamp", "price"],
                )
                dfs.append(df)
            symbol = self.__get_data_symbol(element.symbol, element.source)
            symbol.additional_properties["chart_data_frame"] = pd.concat(dfs)

    async def __debug_breakpoint_event(self, breakpoint_params):
        breakpoint_params = DebugBreakpoint(**breakpoint_params)
        await self.__print_summary(
            breakpoint_params.custom_charts, breakpoint_params.display_charts
        )
