import random
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.formula.api as sm

class OvernightMomo(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2018, 1, 30)
        
        # cash and buffer
        self.Settings.FreePortfolioValuePercentage = 0.1
        self.SetCash(20000)

        # Reality modelling
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
        
        # Universe settings
        self.Reso = Resolution.Daily
        self.UniverseSettings.Resolution = self.Reso
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))     
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.TotalReturn))
        
        # check if its a new month (don't need to select by cap more than once a monty)
        self.lastMonth = -1
        
        # number of symbols in Fine universe
        self.count = 10
        
        # minimum market cap of companies
        self.MinCap = 3e9
        
        # array for symbols to trade
        self.selected = []
        
        # scheduling for trades
        self.AddEquity("SPY", self.Reso)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose("SPY", 20), self.HistLookup) 
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", -20), self.ClosePositions)
    
    
    def OnData(self, data):
        pass
    
    def OnSecuritiesChanged(self, changes):
        # self.Debug("ADDED" + str([security.Symbol.Value for security in changes.AddedSecurities]))
        # self.Debug("REMOVED" + str([security.Symbol.Value for security in changes.RemovedSecurities]))
        pass

    def CoarseSelectionFunction(self, coarse):
        if self.Time.month == self.lastMonth: 
            return Universe.Unchanged
        
        # sort descending by daily dollar volume
        self.Log('Refreshing COARSE Universe >> ' + str(self.Time))
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 5], 
            key=lambda x: x.DollarVolume, reverse=True)[:1000]
        
        # return the symbol objects of the top entries from our sorted collection
        return [ x.Symbol for x in sortedByDollarVolume]
        
    # sort the data by MarketCap and take the bottom 'NumberOfSymbolsFine'
    def FineSelectionFunction(self, fine):
        """
        sortedByMktCap = sorted([x for x in fine if x.MarketCap > self.MinCap], 
            key=lambda x: x.MarketCap, reverse=False)[:self.count]
        
        count = len(sortedByMktCap)
    
        if count == 0:
            return Universe.Unchanged
        
        # update if all checks passed
        self.lastMonth = self.Time.month
        
        # this is a list of symbol objects
        selected_symbols = [ x.Symbol for x in sortedByMktCap ]
        
        # convert universe to DF (https://www.quantconnect.com/forum/discussion/2644/how-do-i-get-a-percentile-of-dollarvolume-universe/p1)
        intercepts = []

        for symbol in selected_symbols:
            # make history call for past 12 months of daily prices
            hist = self.History(symbol, 253, Resolution.Daily)
            
            # need thorough history check is performing regression correctly
            # Warning: when performing history requests, the start date will be adjusted if it is before the first known date for the symbol.
            # think something to do with this
            
            # Check history returns correct number of datapoints, then do regression
            if len(hist) == 253:
                self.Log(hist)
                
                # Define overnight returns: r_o =  ln(open_t / close_t-1)
                hist['r_o'] = pd.Series(np.log(hist['open']/hist['close'].shift(1)))
            
                # Define total returns: r = ln(close / close_t-1)
                hist['r_t'] = pd.Series(np.log(hist['close']/hist['open']))
            
                # regression of overnight on total
                model = sm.ols(formula = 'r_o~r_t',data = hist).fit()
                
                # append intercept param to intercepts list
                intercepts.append(model.params[0])
            
        # Make dataframe
        data_df = pd.DataFrame(list(zip(selected_symbols,intercepts)), columns = ['symbol','intercepts'])
        self.Log(data_df)
        
        # Get quantiles
        lower_percent = data_df.intercepts.quantile(.9)
        upper_percent = data_df.intercepts.quantile(1.0)

        self.selected = (data_df.
            query('(intercepts >= @lower_percent) & (intercepts <= @upper_percent)'))
        """
        return self.selected
        
    def HistLookup(self):
        for symbol in self.selected:
            quantity = self.CalculateOrderQuantity(symbol, 1/len(self.selected))
            marketCloseOrderTicket = self.MarketOnCloseOrder(symbol,quantity)
            
        self.Debug(f"{self.Time} {str([symbol.Value for symbol in self.selected])} will be bot and held overnight")
        
    def ClosePositions(self):
        self.Liquidate()