#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import logging

_logger = logging.getLogger(__name__)


class PerformanceManager(object):
    """
    https://www.quantopian.com/docs/api-reference/pyfolio-api-reference
    Record equity, positions, and trades in accordance to pyfolio format
    First date will be the first data start date
    """

    def __init__(self, instrument_meta):
        self._symbols = []
        self.instrument_meta = instrument_meta  # sym ==> meta

        self._equity = None
        self._df_positions = None
        self._df_trades = None

    def add_watch(self, data_key, data):
        if "Close" in data.columns:  # OHLCV
            self._symbols.append(data_key)
        else:  # CLZ20, CLZ21
            self._symbols.extend(data.columns)

    #  or each sid
    def reset(self):
        self._realized_pnl = 0.0
        self._unrealized_pnl = 0.0

        self._equity = pd.Series(dtype=np.float64)  # equity line
        self._equity.name = "total"

        self._df_positions = pd.DataFrame(
            columns=self._symbols + ["cash"], dtype=np.float64
        )
        self._df_trades = pd.DataFrame(
            np.empty(
                0,
                dtype=np.dtype(
                    [("amount", np.int64), ("price", np.float64), ("symbol", np.str)]
                ),
            )
        )
        # self._df_trades.amount = self._df_trades.amount.astype(int)     # pyfolio transactions

    def on_fill(self, fill_event):
        # self._df_trades.loc[fill_event.timestamp] = [fill_event.fill_size, fill_event.fill_price, fill_event.full_symbol]
        self._df_trades = self._df_trades.append(
            pd.DataFrame(
                {
                    "amount": [int(fill_event.fill_size)],
                    "price": [fill_event.fill_price],
                    "symbol": [fill_event.full_symbol],
                },
                index=[fill_event.fill_time],
            )
        )

    def update_performance(self, current_time, position_manager, data_board):
        """
        update previous time/date
        """
        if self._equity.empty:  # no previous day
            self._equity[current_time] = 0.0

        # on a new time/date, calculate the performances for previous time/date
        if current_time != self._equity.index[-1]:
            performance_time = self._equity.index[-1]
        else:
            # When a new data date comes in, it calcuates performances for the previous day
            # This leaves the last date not updated.
            # So we call the update explicitly
            performance_time = current_time

        equity = 0.0
        self._df_positions.loc[performance_time] = [0] * len(self._df_positions.columns)
        for sym, pos in position_manager.positions.items():
            if sym in self.instrument_meta.keys():
                multiplier = self.instrument_meta[sym]["Multiplier"]
            else:
                multiplier = 1

            # data_board (timestamp) hasn't been updated yet
            if pos.size == 0:  # skip empty size
                continue
            equity += pos.size * data_board.get_last_price(sym) * multiplier
            self._df_positions.loc[performance_time, sym] = (
                pos.size * data_board.get_last_price(sym) * multiplier
            )

        self._df_positions.loc[performance_time, "cash"] = position_manager.cash
        self._equity[performance_time] = equity + position_manager.cash
        self._df_positions.loc[performance_time, "total"] = self._equity[
            performance_time
        ]

        if performance_time != current_time:  # not final day
            self._equity[current_time] = 0.0  # add new date
        else:  # final day, re-arrange column order
            self._df_positions = self._df_positions[self._symbols + ["cash"]]
