from __future__ import print_function
from abc import ABCMeta, abstractmethod
from datetime import datetime

import numpy as np
import os
import pandas as pd

from algo_trade.event import FillEvent, OrderEvent, MarketEvent
from tda.orders import equities


class Broker(metaclass=ABCMeta):
    '''
    Broker abstract class handles the interaction between a set of order objects generated by a Portfolio and the ultimate set of Fill objects that actually occur in the market.

    Handlers can be used to subclass simulated brokerages or live brokerages, with identical interfaces. Allows strategies to be backtested in a very similar manner to the live trading engine.
    '''
    @abstractmethod
    def execute_order(self, event):
        raise NotImplementedError('Should implement execute_order()')

class SimulatedBroker(Broker):
    '''
    Simulated execution handler converts all OrderEvent objects into their equivalent FillEvent objects without latency, slippage, or fill-ratio issues.

    Allows straightforward "first go" test of any strategy, before implementation with a more sophisticated execution handler.

    :param events: Event queue
    :param csv_dir: Absolute directory path to CSV files
    :param ticker_list: List of ticker strings
    :param test: Helper for unittest - blocks importing of CSV files
    '''
    def __init__(self, events, ticker_list, csv_dir=None):
        # Initialization states
        self.events = events
        self.csv_dir = csv_dir
        self.ticker_list = ticker_list

        # Data elements
        self.ticker_data = dict()
        self.latest_ticker_data = dict()
        self.ticker_generator = dict()
        
        self.continue_backtest = True

        if csv_dir:
            self._open_convert_csv_files()

    ########################################
    # Data Initialization methods
    
    def _open_convert_csv_files(self):
        '''
        Opens CSV files from the data directory, converting them in pandas DataFrames within a ticker dictionary. Current assumption is that the data is taken from YFinance.
        '''
        comb_index = None
        for t in self.ticker_list:
            self.ticker_data[t] = pd.io.parsers.read_csv(
                os.path.join(self.csv_dir, '%s.csv' % t),
                header=0,
                index_col=0,
                parse_dates=True,
                names=[
                    'datetime', 'open', 'high', 'low', 'close', 'volume', 'adj_close'
                ]
            ).sort_index()

            if comb_index is None:
                comb_index = self.ticker_data[t].index
            else:
                comb_index.union(self.ticker_data[t].index)

            self.latest_ticker_data[t] = []

        for t in self.ticker_list:
            self.ticker_data[t] = self.ticker_data[t].reindex(index=comb_index, method='pad')
            self.ticker_generator[t] = self.ticker_data[t].iterrows()

    def _get_new_bar(self, ticker):
        '''
        Returns latest bar from the data feed.
        '''
        return next(self.ticker_generator[ticker])

    ########################################
    # Data retrieval methods
    
    def get_latest_bar(self, ticker):
        '''
        Returns the last bar from the latest_ticker list.
        '''
        try:
            bars_list = self.latest_ticker_data[ticker]
        except KeyError:
            print('That ticker is not available in the historical dataset.')
            raise
        else:
            return bars_list[-1]

    def get_latest_bars(self, ticker, N=1):
        '''
        Returns the last N bars from the latest_ticker list, of N-k if less available.
        '''
        try:
            bars_list = self.latest_ticker_data[ticker]
        except KeyError:
            print('That ticker is not available in the historical dataset.')
            raise
        else:
            return bars_list[-N:]

    def get_latest_bar_datetime(self, ticker):
        '''
        Return Python datetime object for the last bar.
        '''
        try:
            bars_list = self.latest_ticker_data[ticker]
        except KeyError:
            print('That ticker is not available in the historical dataset.')
            raise
        else:
            return bars_list[-1][0]

    def get_latest_bar_value(self, ticker, val_type):
        '''
        Returns one of the OHLCVI values from the pandas Bar series object.
        '''
        try:
            bars_list = self.latest_ticker_data[ticker]
        except KeyError:
            print('That ticker is not available in the historical dataset.')
            raise
        else:
            return getattr(bars_list[-1][1], val_type)

    def get_latest_bars_values(self, ticker, val_type, N=1):
        '''
        Returns last N bar values from the latest_symbol list, or N-k if less available.
        '''
        try:
            bars_list = self.get_latest_bars(ticker, N)
        except KeyError:
            print('That ticker is not available in the historical dataset.')
            raise
        else:
            return np.array([getattr(b[1], val_type) for b in bars_list])

    def update_bars(self):
        '''
        Pushes latest bar to the latest_ticker_data structure for all tickers in the ticker list.
        '''
        for t in self.ticker_list:
            try:
                bar = self._get_new_bar(t)
            except StopIteration:
                self.continue_backtest = False
            else:
                if bar is not None:
                    self.latest_ticker_data[t].append((bar[0], bar[1]))

        self.events.put(MarketEvent())

    def execute_order(self, event):
        '''
        Converts OrderEvent into FillEvent objects naively
        '''
        if event.type == 'ORDER':
            fill_cost = self.get_latest_bar_value('SPY', 'adj_close') * event.quantity
            fill_event = FillEvent(
                datetime.utcnow(), event.ticker, 'ARCA', event.quantity, event.action, fill_cost
            )
            self.events.put(fill_event)

class TDABroker(Broker):
    '''
    TDAmeritrade execution handler, takes OrderEvent objects and submits trades via the TDA API. Current feature is HTTP trade submission only, does not return FillEvent information.

    Future - Research TDA API trade execution receipt

    :param client:  TDA API client that links handler to TDAmeritrade account
    :param acc_id:  TDA account number
    :param events:  Queue of Event objects
    '''
    def __init__(self, client, acc_id, events, ticker_list):
        self.client = client
        self.ACC_ID = acc_id
        self.events = events
        self.ticker_list = ticker_list

    ########################################
    # Data methods

    def get_data(self, period_type = 'year', period = 1, frequency_type = 'daily', frequency = 1):
        '''
        Returns dataframe dictionary of historical prices for a list of tickers
        '''
        comb_index = None
        for t in self.ticker_list:
            response = self.client.get_price_history(
                t,
                period_type=period_type,
                period=period,
                frequency_type=frequency_type,
                frequeny=frequency,
            ).json()

            price_history = pd.DataFrame(
                    data=response['candles'],
                    columns=['datetime', 'open', 'high', 'low', 'close'],
                )
            price_history.set_index('datetime', inplace=True)

            self.ticker_data[t] = price_history

            if comb_index is None:
                comb_index = self.ticker_data[t].index
            else:
                comb_index.union(self.ticker_data[t].index)

            self.latest_ticker_data[t] = list()

        for t in self.ticker_list:
            self.ticker_data[t] = self.ticker_data[t].reindex(index=comb_index, method='pad')
            self.ticker_data[t].sort_index(inplace=True)
            self.ticker_generator[t] = self.ticker_data[t].iterrows()

        return self.ticker_data

    ########################################
    # Rebalance Portfolio Methods

    def _create_order(self, ticker, quantity, action, trade_type, limit=None):
        '''
        Creates OrderEvent objects and adds them to the order_book set.

        :param ticker: Ticker symbol
        :param quantity: Number of shares to buy/sell
        :param action: BUY or SELL
        :param order_type: MARKET or LIMIT - currently only supports MARKET types
        '''
        return OrderEvent(
            ticker=ticker,
            quantity=quantity,
            action=action,
            trade_type=trade_type,
            limit=limit,
        )

    def _calculate_weight_change(self, target, current):
        '''
        Calculates the difference between tar_state and cur_state portfolios, returning how much each asset is under/over-weight

        :param tar_state: Pandas DataFrame of target asset weights in decimals
        :param cur_state: Pandas DataFrame of current asset weights in decimals
        '''
        diff_dict = dict()

        # Get difference from current holdings to target holdings
        for k, v in current.items():
            if k not in target:
                diff_dict[k] = -1 * v
            else:
                diff_dict[k] = target[k] - v

        # Check for new asset holdings in target
        for k, v in target.items():
            if k not in current:
                diff_dict[k] = v
        
        # Drop cash component
        diff_dict.pop('MMDA1', None)

        return diff_dict

    def rebalance(self, balance, price, target, current = None):
        '''
        Returns set of OrderEvents created from current TDA account balances and target asset weights. If cur_df provided, will generate BUY and SELL orders accordingly to rebalance portfolio.

        :param balance: Float value of current account balance
        :param price: Dictionary of relevant ticker prices - should include all unique tickers present in target and current
        :param tar_df: Dictionary of target asset weights in float
        :param cur_df: Dictionary of current asset weights in float
        '''

        if current is not None:
            target = self._calculate_weight_change(target, current)

        for k, v in target.items():
            quantity = int(v * balance / price[k])

            print(type(k), k)

            if k == 'MMDA1':
                continue

            if quantity > 0:
                self.events.put( self._create_order(k, quantity, 'BUY', 'MARKET') )
            elif quantity < 0:
                self.events.put( self._create_order(k, abs(quantity), 'SELL', 'MARKET') )

    ########################################
    # Execute order method

    def _submitBuy(self, event):
        '''
        Takes OrderEvent and submits a BUY order to the TDA account
        '''
        if event.trade_type == 'MARKET':
            self.client.place_order(
                self.ACC_ID,
                equities.equity_buy_market(
                    event.ticker,
                    event.quantity,
                ),
            )
        elif event.trade_type == 'LIMIT' and event.limit:
            self.client.place_order(
                self.ACC_ID,
                equities.equity_buy_limit(
                    event.ticker,
                    event.quantity,
                    event.limit
                ),
            )
        else:
            raise Exception('Invalid BUY order.')

    def _submitSell(self, event):
        '''
        Takes OrderEvent and submits a SELL order to the TDA account
        '''
        if event.trade_type == 'MARKET':
            self.client.place_order(
                self.ACC_ID,
                equities.equity_sell_market(
                    event.ticker,
                    event.quantity,
                ),
            )
        elif event.trade_type == 'LIMIT' and event.limit:
            self.client.place_order(
                self.ACC_ID,
                equities.equity_sell_limit(
                    event.ticker,
                    event.quantity, 
                    event.limit,
                ),
            )
        else:
            raise Exception('Invalid SELL order.')

    def execute_order(self, event):
        '''
        Virtual method - takes OrderEvent and calls self._submitBuy() and self._submitSell() in accordance to the order type
        '''
        if event.action == 'BUY':
            self._submitBuy(event)

        if event.action =='SELL':
            self._submitSell(event)