import os
import math
import json
import random
import datetime

from collections import defaultdict
from collections import OrderedDict

import QuantConnect
from QuantConnect import Time
from QuantConnect import Orders
from QuantConnect.Data.UniverseSelection import *


kepler_url   = 'http://api.analyticsventures.com/kepler.json'
earnings_url = 'http://api.analyticsventures.com/earnings.json'


class Kepler:

    def __init__(self, json_data):
        kepler_items = sorted(json_data.items())
        self.lookup = OrderedDict(kepler_items)


    def getFraction(self, ticker, direction, generation = 1):

        """ The kepler model file from Yury maps tickers to a list
            of twolists: one list is used if the price went up,
            while the other is used if the price went down.
            In either case, the 0th and 1st element of each sublist
            contain the fraction that the stock's price must
            increase (or decrease) in order for the 1st and
            2nd generation limit orders to be triggered,
            respectively. To illustrate, when passed the key
            'AAPL', the kepler model file returns the following
            list of two lists:

            [[0.075, 0.024, 1, -0.03661209418309422, -0.009153023545773555],
             [0.02, 0.032, 3, 0.025277211129414216, 0.006319302782353554]]

            These numbers can be interpreted as follows.

            Suppose 'AAPL' reports earnings overnight. The next morning,
            we compare the previous day's close price to today's open price.
            Consider the following two scenarios:

            (1) If the price went up, we'll be using the list:
            [0.075, 0.024, 1, -0.03661209418309422, -0.009153023545773555]

            (2) If the price went down, we'll be using the list:
            [0.02, 0.032, 3, 0.025277211129414216, 0.006319302782353554]

            In the first case, the price went up, and we predict it will
            continue to go up, so we place a limit order to *sell* some
            number of AAPL shares if the price reaches close*(1 + 0.075),
            where 'close' is the closing price of AAPL on the previous
            day (i.e., the day before earnings were reported).
            If that order successfully executes, we'll then place a limit
            order to *buy* the same number of shares of AAPL if the price
            reaches fill*(1 - 0.024), where 'fill' is the price at which
            the first order filled.

            In the second case, the price went down, and we predict it will
            continue to go down, so we place a limit order to *buy* some
            number of AAPL shares if the price reaches close*(1 - 0.02),
            where 'close' is the closing price of AAPL on the previous
            day (i.e., the day before earnings were reported).
            If that order successfully executes, we'll then place a limit
            order to *sell* the same number of shares of AAPL if the price
            reaches fill*(1 + 0.032), where 'fill' is the price at which
            the first order filled.

            In all cases, the fraction in question is applied to the
            current (or very recent) price of the stock, and thus always
            represents a relative change in the price of that stock,
            regardless of whether we're buying or selling, or whether
            the earnings report was followed by an increase or decrease
            in the price of the stock.
        """
        kdata = self._getlist(ticker, direction)

        if generation == 1:
            fraction = kdata[0]
        elif generation == 2:
            fraction = kdata[1]
        else:
            raise ValueError(f"Unknown generation: '{generation}'")
        return fraction


    def getLifetime(self, ticker, direction, generation = 1):

        """ The kepler model file maps tickers to a list of
            two lists, one to be used if the price went up,
            the other to be used if the price went down. In
            either case, the 2nd element of each sublist
            contains the lifetime in days that the 2nd limit
            order should be allowed to live before we expire it.
            The lifetime of the first generation limit order in
            the kepler model is always set to a single day.
        """
        kdata = self._getlist(ticker, direction)

        if generation == 1:
            lifetime = 1
        elif generation == 2:
            lifetime = kdata[2] # index 2 in model data is lifetime in days
        else:
            raise ValueError(f"Unknown generation: '{generation}'")
        return lifetime


    def _getlist(self, ticker, direction):

        """ The kepler model file from provides a dictionary
            mapping each ticker to a list of two lists, each sublist
            containing 5 numbers. The first of the two sublists contains
            parameters to use if the overnight reporting of earnings
            was followed by a stock's price increasing, and the second
            list contains the parameters to use if earnings led to the
            price decreasing. This method simply returns the relevant
            sublist when passed (a) a ticker, and (b) a string specifying
            whether the price went up or down. 
        """
        if direction == 'up':
            kdata = self.lookup[ticker][0]
        elif direction == 'down':
            kdata = self.lookup[ticker][1]
        else:
            raise ValueError(f"Unrecognized direction: {direction}")
        return kdata


class Earnings:


    def __init__(self, json_data):

        """ This class contains utility functions for accessing
            simple information about historical earnings dates.
        """
        earnings_items = sorted(json_data.items())
        self.datemap = OrderedDict(earnings_items)

        all_tickers = []
        for date, ticker_list in self.datemap.items():
            all_tickers += ticker_list
        self.all_tickers = sorted(set(all_tickers))


    def _check_date(self, iso_date):

        """ Since the historical earnings data is stored in a dictionary
            whose keys are date strings, raise an exception to alert
            ourselves if for some reason we ever request a date that isn't
            formatted in the same manner as the dict keys, rather than simply
            handing back an empty list and allowing the error to pass silently.
        """
        year, month, day = iso_date.split('-')
        if len(year) != 4:
            raise ValueError(f"Improperly formatted year: {year}")
        if len(month) != 2:
            raise ValueError(f"Improperly formatted month: {year}")
        if len(day) != 2:
            raise ValueError(f"Improperly formatted day: {year}")


    def next_date_after(self, iso_date):

        """ Find the first date in our historical earnings dates
            after the passed date, on which some company reported.
        """
        self._check_date(iso_date)

        for date_key in self.datemap.keys():
            if date_key > iso_date: break
        else:
            raise KeyError(f"No known earnings date after {iso_date}")
        return date_key


    def next_date_before(self, iso_date):

        """ Find the first date in our historical earnings dates
            before the passed date, on which some company reported.
        """
        self._check_date(iso_date)

        for date_key in reversed(self.datemap.keys()):
            if date_key < iso_date: break
        else:
            raise KeyError(f"No known earnings date before {iso_date}")
        return date_key


    def next_companies_after(self, iso_date):

        """ Find the set of companies that will be reporting on the
            first date in our catalogue of historical earnings dates
            after the passed date.
        """
        next_date = self.next_date_after(iso_date)
        return self.datemap[next_date]


    def next_companies_before(self, iso_date):

        """ Find the set of companies that will be reporting on the
            first date in our catalogue of historical earnings dates
            before the passed date.
        """
        prev_date = self.next_date_before(iso_date)
        return self.datemap[prev_date]


    def on(self, iso_date):

        """ Returns a list of companies reporting on a given date. """

        self._check_date(iso_date)
        try:
            return self.datemap[iso_date]
        except KeyError:
            return []


class Algorithm(QCAlgorithm):

    def Initialize(self):

        """ This class contains the main code for user story ch369.

            In the example that follows, we implement the basic logic
            required to run a simple backtest of the kepler model.
        """
        self.SetStartDate(2017, 9, 1)
        self.SetEndDate(2018, 1, 1)

        self.SetCash(1_000_000)

        self.openLimitOrders = defaultdict(list)

        self.currentUniverse = []

        # The following set contains the symbols that
        # QuantConnect complains about if we try to buy
        # them, with an error message simply saying that
        # they have been marked as "non-tradeable" without
        # providing any clear reason. We can just skip them.
        self.blacklistedTickers = {"BGC", "MON", "TWX", "STL"}

        earnings_data = self.Download(earnings_url)
        self.earnings = Earnings(json.loads(earnings_data))

        kepler_data = self.Download(kepler_url)
        self.kepler = Kepler(json.loads(kepler_data))

        # SPY is used as a reference symbol for scheduling events,
        # following the Quantconnect documentation and examples here:
        # https://www.quantconnect.com/docs/algorithm-reference/scheduled-events
        self.AddEquity("SPY", Resolution.Minute)

        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.AfterMarketOpen("SPY", 1),
            self.marketOpenCallback
        )

        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.BeforeMarketClose("SPY", 1),
            self.marketCloseCallback
        )

        self.UniverseSettings.Resolution = Resolution.Minute

        self.AddUniverse("earnings-universe", self.universeSelector)


    # def Debug(self, *args, **kwargs):
    #     pass


    def universeSelector(self, dt):

        today = dt.date().isoformat()

        earningsToday = self.earnings.on(today)

        currentUniverse = [t for t in earningsToday if t not in self.blacklistedTickers]

        self.currentUniverse = currentUniverse

        # We *absolutely* need this. Without manually re-adding each ticker to our
        # equities (even if it's already in our universe), everything will seem to
        # work *until* we add a ticker to our universe which was previously added
        # and then subsequently removed. This was the source of the spacebar ticker bug.
        for ticker in self.currentUniverse:
            self.AddEquity(ticker, Resolution.Minute)

        self.Debug(f"===========")
        self.Debug(f"Universe for {today} is {self.currentUniverse}")

        # Returning a non null value here is required by the framework,
        # even though we can access the same data most directly simply
        # by referring to the attribute we just set on self.
        return self.currentUniverse


    def isBusinessDay(self, day):
        tradingDay = self.TradingCalendar.GetTradingDay(day)
        isBusinessDay = tradingDay.BusinessDay
        return isBusinessDay


    def getPreviousTradingDay(self):
        day = self.Time
        while True:
            day -= datetime.timedelta(1)
            if self.isBusinessDay(day):
                return day


    def getPreviousDayClose(self, symbol):
        yest = self.getPreviousTradingDay().date()
        a = datetime.datetime(yest.year, yest.month, yest.day, 15, 59, 0)
        b = datetime.datetime(yest.year, yest.month, yest.day, 16, 0, 0)
        try:
            hist = self.History([symbol], a, b, Resolution.Minute)
            close = float(hist.loc[symbol, 'close'].iloc[-1])
        except KeyError:
            return None
        else:
            return close


    def getPrice(self, symbol):
        price = float(self.Securities[symbol].Price)
        return price


    def marketOpenCallback(self):
        """
            Scheduled to be called soon after market open.
        """

        if len(self.currentUniverse) == 0:
            return

        prices = {}
        closes = {}
        activeTickers = []

        for ticker in self.currentUniverse:

            try:
                price = float(self.Securities[ticker].Price)
            except Exception as e:
                self.Debug(f"[!] Getting price failed for ticker '{ticker}'.")
                continue

            try:
                close = self.getPreviousDayClose(ticker)
                if close is None: continue
            except Exception as e:
                self.Debug(f"[!] Getting previous day close failed for ticker '{ticker}: {e}'")
                continue

            if price < 0:
                self.Debug(f"[!] Price {price} for ticker {ticker} is negative.")
                continue

            if price == 0:
                # self.Debug(f"[!] Price {price} for ticker {ticker} is zero.")
                continue

            prices[ticker] = price
            closes[ticker] = close
            activeTickers.append(ticker)

        ### Can't do above and below in the same loop, since getting some of the prices
        ### begins randomly failing after any given backtest runs for more than 1 month,
        ### so we need to divide the cash evenly between the tickers we can actually
        ### *buy* without error. Hence, the division of this logic into two loops.

        ### Set daily invested fraction to a sort of "reverse ReLU"
        ### to avoid investing too much in a ticker just because it
        ### happens to be the only one reporting earnings on a given day.

        num_tickers = len(activeTickers)

        if num_tickers == 0:
            self.Debug(f"[!] Had {len(self.currentUniverse)} tickers for today but getting the price failed for all of them.")
            return


        # Intuitively, the amount of cash we want to divide up between all the
        # companies reporting earnings on a given day should be:
        # (a) Close to %100, so we're using all the money we have,
        # (b) Not *too* close to %100, so we have a safety net, and
        # (c) Smaller if only one (or a few) companies are reporting earnings
        #     on a given day, compared to days on which a larger number of companies
        #     is reporting, to avoid putting all our eggs in one basket on days
        #     when we have lots of money to spare but only (say) one company
        #     reports earnings.
        # 
        # The simplest tradeoff that achieves all of the above is to make the
        # fraction of our total remining money that we divide up on each day to
        # be a linearly increasing function of the number of companies reporting,
        # which then gets set to a constant value once it's close to 100% so we
        # don't try to allocate more money than we have if (say) 30 companies happen
        # to report earnings on a given day during the earnings season.
        #
        # The easiest way to do this is to just take the max of a line and a constant.
        # Let's do that:
        current_cash = float(self.Portfolio.Cash)

        max_frac_per_day = 0.90
        min_frac_per_day = 0.10
        
        frac_per_day = min(min_frac_per_day * num_tickers, max_frac_per_day)

        cash_per_ticker = (frac_per_day * current_cash) / num_tickers


        for ticker in activeTickers:

            price = prices[ticker]
            close = closes[ticker]

            direction = 'up' if (price >= close) else 'down'
            sign = +1 if direction == 'up' else -1

            # ^ sign being positive means the price went up, sign being negative means it went down

            # IMPORTANT NOTES:
            # 
            # 'shares' will be *positive* in generation 1 iff sign is *negative*,
            # since we're buying when the price goes down.
            #
            # 'fraction' will be *negative* in generation 1 iff sign is *negative*,
            # since we're always typing 'price * (1 + fraction)', and if sign is negative,
            # then the price went down, so we want to buy, so the limit price is *less* than the price.

            shares = (-sign) * math.floor(cash_per_ticker / price) # (-sign) is essential here!

            if shares == 0:
                continue

            fraction = (+sign) * self.kepler.getFraction(ticker, direction, generation = 1)
            lifetime = self.kepler.getLifetime(ticker, direction, generation = 1)

            limitPrice = price * (1 + fraction)

            try:
                order = self.LimitOrder(ticker, shares, limitPrice)
                order.shares = shares
                order.daysUntilExpiry = lifetime
                order.generation = 1
                order.previousDayClose = close
                order.direction = direction
                order.keplerType = 'initial'
                self.openLimitOrders[ticker].append(order)
                # self.Debug(f"marketOpenCallback: Placed order to buy {shares} shares of {ticker} at {limitPrice}")
            except Exception as e:
                self.Debug(f"marketOpenCallback: Placing limit order for {ticker} failed. Exception was: {e}")


    def marketCloseCallback(self):
        """
            Scheduled to be called soon before market close.

            Here we determine which of our orders have filled,
            and which of our orders have expired. If a given
            order has neither filled nor expired, we decrement
            its number of days until expiry by 1, and keep it
            around for the following trading day.
        """

        newOpenLimitOrders = defaultdict(list)

        # call list() here to *copy* the dict_view object on the off chance that
        # an order goes through while we're running the marketCloseCallback, which
        # (though rare), when it happens, causes an exception saying that the dict
        # self.openLimitOrders changed size while we were iterating over it, even
        # though we never change it ourselves in the loop below.
        # Tl;dr: Calling list() on something that's already an iterable looks dumb,
        # but it's actually required to cover an odd corner case here.
        for ticker, openOrders in list(self.openLimitOrders.items()):

            for order in openOrders:

                g = order.generation

                if order.Status == OrderStatus.Filled:

                    # self.Debug(f"marketCloseCallback: Generation {g} order for {ticker} was filled!")
                    pass

                elif order.Status == OrderStatus.Canceled:

                    # self.Debug(f"marketCloseCallback: Generation {g} order for {ticker} was canceled.")
                    pass

                elif order.daysUntilExpiry <= 0:

                    try:
                        order.Cancel(f"marketCloseCallback: Generation {g} order for {ticker} expired")
                    except:
                        orderType = self.decodeEventType(order)
                        self.Debug(f"[-] Failed to cancel generation {g} order of type {orderType}")

                    if g == 2:
                        try:
                            self.Liquidate(ticker)
                            # self.Debug(f"[+] Success liquidating on expiry for ticker {ticker}, generation {g}")
                        except:
                            self.Debug(f"[-] Failed to liquidate on expiry for ticker {ticker}, generation {g}")

                else:
                    order.daysUntilExpiry -= 1
                    newOpenLimitOrders[ticker].append(order)


        self.openLimitOrders = newOpenLimitOrders


    def decodeEventType(self, orderEvent):
        """
            Turns the numeric values we receive from
            orderEvent.Status into a human readable string so we
            can tell what kind of event is being passed when the
            OnOrderEvent method is automatically called on fills
        """

        statuses = ('New', 'Submitted', 'PartiallyFilled', 'Filled', 
                    'Canceled', 'None', 'Invalid', 'CancelPending')

        order_statuses = {
            getattr(Orders.OrderStatus, status): status
            for status in statuses
        }

        return order_statuses[orderEvent.Status]



    def tryFindOrder(self, ticker, orderEvent):

        thisOrderId = orderEvent.OrderId
        thisOrder = self.Transactions.GetOrderById(orderEvent.OrderId)

        ordersForTicker = self.openLimitOrders[ticker] # ordersForTicker is now a list of order tickets

        # self.Debug(f"Ticker is {ticker}. Number of open orders for {ticker} is {len(ordersForTicker)}")

        for thatOrder in ordersForTicker:

            thatOrderId = thatOrder.OrderId

            if thisOrderId == thatOrderId:
                # self.Debug(f"OnOrderEvent: Found the order we were looking for!")
                order = thatOrder
                break

        else:
            raise ValueError(f"[-] Couldn't find the order for {ticker}! Not sure why...")

        # Why do we need to do all this?
        #
        # Well, if you happen to have QuantConnect open at the moment, and you
        # click on the "API" button over on the left and then scroll down to
        # "OnOrderEvent", you'll see that the documentation begins with the
        # delightful little snippet:
        # 
        # "EXPERTS ONLY :: [-!-Async Code-!-]"
        # 
        # Now for Kepler, we need to be able to place new Limit Orders when
        # other orders go through, so we've got to use that function with
        # the big warning in all caps.
        #
        # So here's the deal: When we wake up and find ourselves in OnOrderEvent,
        # we need to figure out which order we're being called *on*. That is,
        # which order just generated this event that we now need to handle.
        # 
        # To do that, all we have access to is this "orderEvent" object that
        # gets passed to the OnOrderEvent function. Unfortunately, calling
        # 
        # self.Transactions.GetOrderById(orderEvent.OrderId)
        # 
        # as suggested by the QuantConnect documentation does *not* return the
        # original order ticket object that was returned when we originally
        # placed the LimitOrder that just filled. (That's right.) It returns a
        # different object, for no obvious reason. That means the object we're
        # given is missing all the essential metadata we've been keeping track of by
        # attaching it to the original orderTicket object.
        #
        # Despite that inconvenience, we need some way to track arbitrary
        # metadata about these orders. The simplest solution appears to be
        # just keeping track of a list of limit orders we've placed, and then
        # checking to see which of the objects we've been keeping track of has the
        # same order id as the orderEvent we get passed in the event handler.
        # The need for the code in this function is illustrated by the following
        # three assertions.
        assert hasattr(thisOrder, 'generation') is False
        assert hasattr(thatOrder, 'generation') is True
        assert thisOrderId == thatOrderId

        return order


    def OnOrderEvent(self, orderEvent):
        """
            This method is called by the QuantConnect framework whenever
            an "order event" occurs. Common events occur at the moment
            when an order is submitted, fully filled, partially filled,
            cancelled, or at the moment cancellation is requested, but
            before the cancellation occurs. For our purposes, the most
            relevant event types are the ones that are emitted when a
            limit order either (a) fills or (b) is canceled.
        """

        ticker = orderEvent.Symbol.Value # Need the ".Value" to actually make this a real str object
        filled = orderEvent.FillQuantity
        eventType = self.decodeEventType(orderEvent)

        thisOrder = self.Transactions.GetOrderById(orderEvent.OrderId)
        thisOrderId = orderEvent.OrderId
        createdTime = thisOrder.CreatedTime


        if eventType == 'Submitted':
            # These "submitted" events mostly just end up adding noise to the logfile.
            # We know when we submit a LimitOrder, and these events occur as
            # deterministically as we could hope for after each such order that
            # we place, so any logging we need to do concerning when an order is
            # submitted can just be done at the moment we submit it.
            return


        elif eventType == 'Canceled':
            return


        elif eventType == 'Filled':

            try:
                order = self.tryFindOrder(ticker, orderEvent)
                g = order.generation
            except:
                # self.Debug(f"[*] Got a postmortem 'filled' event for '{ticker}', which apparently means we just liquidated it.")
                return

            # We now have the order, complete with its modified metadata.
            # Time to place the generation 2 orders (if we just filled generation 1)
            # or else die happy (if we just filled generation 2)

            # self.Debug(f"OnOrderEvent: Filled {ticker}'s generation {g} order!")

            if g == 2:
                assert hasattr(order, 'keplerType')
                # Now attempt to cancel the order's partner
                partner = order.partner
                CANCEL_PARTNER_MESSAGE  = f"[+] Generation {g} order of type '{order.keplerType}' filled for {ticker}. "
                CANCEL_PARTNER_MESSAGE += f"Cancelling partner order of type '{partner.keplerType}'"
                # self.Debug(CANCEL_PARTNER_MESSAGE)
                partner.Cancel(CANCEL_PARTNER_MESSAGE)
                return

            assert order.generation == 1, f"Got bad generation value: {order.generation}"

            try:
                price = float(order.AverageFillPrice) # Yury said to do this
            except:
                self.Debug(f"[-] Getting price failed for ticker '{ticker}'. Could not place generation 2 limit orders.")
                return

            # This is the *original* direction that the price moved after earnings.
            # Note: The decisions we make here are flipped relative to the ones we made originally
            direction = order.direction

            sign = +1 if direction == 'up' else -1
            # ^ sign being positive means the price went up, sign being negative means it went down

            # IMPORTANT NOTES, GENERATION 2:
            # 
            # * shares will be *negative* in generation 2 iff sign is *negative*, since we bought in generation 1 if the
            #   price went down, which means we're going to be trying to *sell* now.
            # 
            # * fraction will be *positive* in generation 2 iff sign is *negative*, since we're always typing 'price * (1 + fraction)',
            #   and if sign is negative, then the price went down originally, so we bought, but now we're trying to sell,
            #   so the limit price is *more than* the price.

            shares = -1 * order.shares # at generation 2, we place the inverse order of what we placed originally
            fraction = (-sign) * self.kepler.getFraction(ticker, direction, generation = 2)
            lifetime = self.kepler.getLifetime(ticker, direction, generation = 2)

            limitPriceHappy = price * (1 + fraction)
            limitPriceStop  = price * (1 - fraction)

            try:
                newOrderHappy = self.LimitOrder(ticker, shares, limitPriceHappy)
                newOrderHappy.daysUntilExpiry = lifetime
                newOrderHappy.generation = 2
                # self.Debug(f"OnOrderEvent: Placed generation 2 happy order for {shares} shares of {ticker} at {limitPriceHappy}, relative to price {price}")

                newOrderStop = self.LimitOrder(ticker, -shares, limitPriceStop)
                newOrderStop.daysUntilExpiry = lifetime
                newOrderStop.generation = 2
                # self.Debug(f"OnOrderEvent: Placed generation 2 stop order for {-shares} shares of {ticker} at {limitPriceStop}, relative to price {price}")

                # Now set partnership between the stop-loss order and the 2nd generation limit order.
                newOrderHappy.partner = newOrderStop
                newOrderStop.partner = newOrderHappy

                newOrderHappy.keplerType = 'happy'
                newOrderStop.keplerType  = 'stop'

                self.openLimitOrders[ticker].append(newOrderHappy)
                self.openLimitOrders[ticker].append(newOrderStop)

            except:
                self.Debug(f"OnOrderEvent: Placing generation 2 order for {ticker} failed.")

        elif eventType == 'CancelPending':
            pass

        else:
            # self.Debug(f"OnOrderEvent: Got '{eventType}' event for {ticker}")
            pass