# https://quantpedia.com/strategies/overnight-stock-trading/
#
# The investment univese consists of stocks from NYSE/AMEX/Nasdaq. Each year investor picks stocks which are the most susceptible to have 
# elevated returns during overnight session. These stocks could be picked by simple approach (investor sorts all stocks by their nightly
# performance during previous year and picks the best performing decile) or by using regression model stated on page 19 (equation number 
# 8, linear regression model in which daily log overnight returns are regressed on daily log total returns, investor sorts stocks into 4 
# groups based on their susceptibility to have elevated/depressed overnight returns). Each day at close, investor buys stocks which are 
# in a group of significant night performers and short stocks from group of significant night laggards. Long short portfolio is held
# overnight and is liquidated at open. Stocks are weighted equally.
# 
# QC implementation:
#   - The investment univese consists of 100 most liquid US stocks.
#   - Each year, stocks are picked by their overnight performance during previous year.

import numpy as np
from AlgorithmImports import *

class OvernightStockTrading(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 1, 1)
        self.SetCash(10000)

        self.period = 12 * 21
        
        self.symbol = self.AddEquity('SPY', Resolution.Minute).Symbol
        
        self.course_count = 100
        
        self.selected = [] # symbols of selected stocks from CoarseSelectionFunction
        
        self.long = []
        self.short = []
        self.data = {}
        
        self.months_counter = 1
        self.selection_flag = True
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction)
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), self.TimeRules.BeforeMarketClose(self.symbol, 20), self.MarketClose)
        self.Schedule.On(self.DateRules.MonthStart(self.symbol), self.TimeRules.BeforeMarketClose(self.symbol), self.Selection)

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(10)
                
    def CoarseSelectionFunction(self, coarse):
        # updating overnight data and stock price every day
        for stock in coarse:
            symbol = stock.Symbol
            
            if symbol in self.data:
                # update stock price
                self.data[symbol].price = stock.AdjustedPrice
                
                # get history data
                history = self.History(symbol, 1, Resolution.Daily)
                # update overnight return and change prev_close_price
                self.UpdateOvernightReturns(history, symbol)
                
        # one year rebalance
        if not self.selection_flag:
            return Universe.Unchanged
        self.selection_flag = False
        
        # sort stocks by dollar volume
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5 and x.Market == 'usa'],
            key=lambda x: x.DollarVolume, reverse=True)
        
        # get symbols of top n stocks by dollar volume
        selected = [x.Symbol for x in selected[:self.course_count]]
        
        for symbol in selected:
            if symbol in self.data:
                continue
            
            # create object of SymbolData class for current stock
            self.data[symbol] = SymbolData(self.period)
            
            # get history data
            history = self.History(symbol, self.period + 1, Resolution.Daily)
            # update overnight return and change prev_close_price
            self.UpdateOvernightReturns(history, symbol)
            
        # change self.selected list on rebalance
        self.selected = selected
            
        return selected
            
    def MarketClose(self):
        total_performance = {} # storing total overnight returns performance for self.period overnight returns
        
        # calculate total overnight performances
        for symbol in self.selected:
            if not self.data[symbol].is_overnight_returns_ready():
                continue
            
            # calculate and store total overnight performance
            total_performance[symbol] = self.data[symbol].total_overnight_performance()
            
        if len(total_performance) >= 10:
            # decile selection
            decile = int(len(total_performance) / 10)    
            sorted_by_total_perf = [x[0] for x in sorted(total_performance.items(), key=lambda item: item[1])]
            
            # long top decile stocks and short bottom decile stocks
            self.long = sorted_by_total_perf[-decile:]
            self.short = sorted_by_total_perf[:decile]
            
            long_length = len(self.long)
            short_length = len(self.short)
            
            # trade execution
            for symbol in self.long:
                current_price = self.data[symbol].price
                if current_price != 0 and self.Securities[symbol].Price != 0 and self.Securities[symbol].IsTradable: 
                    quantity = np.floor((self.Portfolio.TotalPortfolioValue / long_length) / current_price)
                    self.MarketOnCloseOrder(symbol, quantity)
                    self.MarketOnOpenOrder(symbol, -quantity)
                
            for symbol in self.short:
                current_price = self.data[symbol].price
                if current_price != 0 and self.Securities[symbol].Price != 0 and self.Securities[symbol].IsTradable: 
                    quantity = np.floor((self.Portfolio.TotalPortfolioValue / short_length) / current_price)
                    self.MarketOnCloseOrder(symbol, -quantity)
                    self.MarketOnOpenOrder(symbol, quantity)
    
    def UpdateOvernightReturns(self, history, symbol):
        ''' update overnight returns for specific stock according to history data '''
        
        # check if history isn't empty and history dataframe has required attributes
        if not history.empty and hasattr(history, 'close') and hasattr(history, 'open'):
            # get open and close prices from dataframe
            opens = history['open']
            closes = history['close']
            
            # update overnight return 
            for (_, open_price), (_, close_price) in zip(opens.iteritems(), closes.iteritems()):
                # update overnight return and change prev_close_price
                self.data[symbol].update(open_price, close_price)
        
    def Selection(self):
        if self.months_counter % 3 == 0:
            self.selection_flag = True
        self.months_counter += 1
    
class SymbolData():
    def __init__(self, period):
        self.overnight_returns = RollingWindow[float](period)
        self.prev_close_price = None
        self.price = 0
        
    def update(self, open_price, close_price):
        # update overnight returns only if prev_close_price isn't None
        if self.prev_close_price:
            overnight_return = open_price / self.prev_close_price - 1
            self.overnight_returns.Add(overnight_return)
        # change previous close price to current close price
        self.prev_close_price = close_price
        
    def total_overnight_performance(self):
        overnight_returns = [x for x in self.overnight_returns]
        return sum(overnight_returns)
        
    def is_overnight_returns_ready(self):
        return self.overnight_returns.IsReady

# Custom fee model
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters):
        fee = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))