import numpy as np
from scipy import stats
import pandas as pd

from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Risk.NullRiskManagementModel import NullRiskManagementModel
from QuantConnect.Indicators import *

class A0001(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 20)  # Set Start Date
        # self.SetEndDate(2017, 1, 10)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Benchmark
        self.benchmark = Symbol.Create('SPY', SecurityType.Equity, Market.USA)
        self.AddEquity('SPY', Resolution.Daily)
        
        # Data resolution
        self.UniverseSettings.Resolution = Resolution.Daily
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False

        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.AddAlpha(ClenowMomentumAlphaModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(self.DateRules.Every(DayOfWeek.Wednesday)))
        self.SetRiskManagement(NullRiskManagementModel())
        
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil

class QC500UniverseSelectionModel(FundamentalUniverseSelectionModel):
    '''Defines the QC500 universe as a universe selection model for framework algorithm
    For details: https://github.com/QuantConnect/Lean/pull/1663'''

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the QC500UniverseSelectionModel'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = 500
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1
        self.referenceSMAperiod = 200
        self.referenceTicker = "SPY"

    def SelectCoarse(self, algorithm, coarse):
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged

        # Do not invest in stocks generating errors messages in logs
        filteredErrors = [x for x in coarse if x.Symbol.Value != "VIAC" and x.Symbol.Value != "BEEM" and x.Symbol.Value != "LSI" and x.Symbol.Value != "IQV" and x.Symbol.Value != "GBT" and x.Symbol.Value != "VTRS" and x.Symbol.Value != "FUBO" and x.Symbol.Value != "SPCE" and x.Symbol.Value != "TFC" and x.Symbol.Value != "PEAK"]

        sortedByDollarVolume = sorted([x for x in filteredErrors if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
                                     key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        # return the symbol objects our sorted collection
        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''Performs fine selection for the QC500 constituents
        The company's headquarter must in the U.S.
        The stock must be traded on either the NYSE or NASDAQ
        At least half a year since its initial public offering
        The stock's market cap must be greater than 500 million'''

        sortedBySector = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        and x.MarketCap > 5e8],
                               key = lambda x: x.CompanyReference.IndustryTemplateCode)

        count = len(sortedBySector)

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if count == 0:
            return Universe.Unchanged

        # Update self.lastMonth after all QC500 criteria checks passed
        self.lastMonth = algorithm.Time.month

        percent = self.numberOfSymbolsFine / count
        sortedByDollarVolume = []

        # select stocks with top dollar volume in every single sector
        for code, g in groupby(sortedBySector, lambda x: x.CompanyReference.IndustryTemplateCode):
            y = sorted(g, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
            c = ceil(len(y) * percent)
            sortedByDollarVolume.extend(y[:c])

        sortedByDollarVolume = sorted(sortedByDollarVolume, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)
        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]

from QuantConnect.Algorithm.Framework.Alphas import *

class ClenowMomentumAlphaModel(AlphaModel):
    def __init__(self, resolution = Resolution.Daily):
        self.resolution = resolution
        self.symbolDataBySymbol = {}
        
        self.linRegPeriod = 90
        self.SMAperiod = 100
        self.ATRperiod = 20
        
        self.referenceSMAperiod = 200
        self.referenceTicker = "SPY"
        
        self.riskFactor = 0.1/100

        self.isPositionsRebalancingWeek = True
        self.previousStocksBought = []
        
    def Update(self, algorithm, data):
        insights = []
        topSelection = []
        stocksToBuy = []
        
        
        # Weekly update of insights every Wednesday + rebalancing every 2 weeks
        if algorithm.Time.isoweekday() != 3:
            return insights
        else:
            algorithm.Log("Weekly Wednesday investing start")
            for symbol, symbolData in self.symbolDataBySymbol.items():
                security_data = algorithm.History([symbol], self.linRegPeriod, Resolution.Daily)
                
                # check if candle has close component(assume others?)
                if 'close' not in security_data.columns:
                    algorithm.Log("History had no Close for %s"%symbol)
                    security_data = None
                      
                if security_data is not None:
                    # we need enough for the np.diff which removes 1 from length
                    if len(security_data.close.index) < self.linRegPeriod: 
                        algorithm.Log("Close test had too few or no values for %s with period %d"%(symbol, self.linRegPeriod))
                        security_data = None

                if security_data is None:
                    continue
                else:
                    y = np.log(security_data.close.values)
                    x = np.arange(len(y))
                    slope, intercept, r_value, p_value, std_err = stats.linregress(x,y)
        
                    momentum = ((1+slope)**252)*(r_value**2)

                    if data.ContainsKey(symbolData.Symbol):
                        if data[symbolData.Symbol] is None:
                            continue
                        else:
                            weight = self.riskFactor*data[symbolData.Symbol].Close/symbolData.ATR.Current.Value
             
                            selectedStock = []
                            selectedStock.append(symbolData)
                            selectedStock.append(momentum)
                            selectedStock.append(weight)

                            topSelection.append(selectedStock)
                    
            #Rank and keep 20% of stock based on momentum indicator
            topSelection = sorted(topSelection, key=lambda x: x[1], reverse=True)            
            topSelection = topSelection[:int(20*len(topSelection)/100)]

            AvailablePorfolioPercent = 1
            
            # buy the previous stock first or sell according to conditions
            # algorithm.Log("buying previous starting with % "+str(AvailablePorfolioPercent) )
            for previousStockBought in self.previousStocksBought:
                for selectedStock in topSelection:
                    if previousStockBought[0] == selectedStock[0] and data[selectedStock[0].Symbol].Close > selectedStock[0].SMA.Current.Value:
                        if self.isPositionsRebalancingWeek == True:
                            if selectedStock[2] < AvailablePorfolioPercent:
                                # algorithm.Log("Rebalancing with fresh stock")
                                
                                AvailablePorfolioPercent = AvailablePorfolioPercent - selectedStock[2]
                                stocksToBuy.append(selectedStock)
                                
                                # algorithm.Log(str(selectedStock[0].Symbol) + " / " + str(selectedStock[1]) + " / " + str(selectedStock[2]) )
                                # algorithm.Log("previous stock / available % "+str(AvailablePorfolioPercent))
                        else:
                            if previousStockBought[2] < AvailablePorfolioPercent:
                                # algorithm.Log("Rebalancing with previous stock")
                                
                                AvailablePorfolioPercent = AvailablePorfolioPercent - previousStockBought[2]
                                stocksToBuy.append(previousStockBought)
                                
                                # algorithm.Log(str(previousStockBought[0].Symbol) + " / " + str(previousStockBought[1]) + " / " + str(previousStockBought[2]) )
                                # algorithm.Log("previous stock / available % "+str(AvailablePorfolioPercent))
                        
                        topSelection.remove(selectedStock)
                    elif data[selectedStock[0].Symbol].Close < selectedStock[0].SMA.Current.Value:
                        topSelection.remove(selectedStock)
            
            # buy the rest with money left if bull market condition is ok
            SPYSymbol = algorithm.AddEquity(self.referenceTicker, Resolution.Daily)

            referenceHistory = algorithm.History([self.referenceTicker], self.referenceSMAperiod, Resolution.Daily)
            referenceSMA = algorithm.SMA(self.referenceTicker, self.referenceSMAperiod, Resolution.Daily)
            referencePrice = None
            
            if not referenceHistory.empty:
                    for tuple in referenceHistory.loc[self.referenceTicker].itertuples():
                        referenceSMA.Update(tuple.Index, tuple.close)
                        referencePrice = tuple.close
    
            if referencePrice < referenceSMA.Current.Value:
                algorithm.Log(str(referencePrice)+ " / "+ str(referenceSMA.Current.Value) + " / bear market detected")
            else:
                # algorithm.Log(str(referencePrice)+ " / "+ str(referenceSMA.Current.Value) + " / bull market validation")
                # algorithm.Log("buying extra starting with % "+str(AvailablePorfolioPercent))
                
                for selectedStock in topSelection:
                    if selectedStock[2] < AvailablePorfolioPercent:
                        stocksToBuy.append(selectedStock)
                        AvailablePorfolioPercent = AvailablePorfolioPercent - selectedStock[2]
                        
                        # algorithm.Log(str(selectedStock[0].Symbol) + " / " + str(selectedStock[1]) + " / " + str(selectedStock[2]) )
                        # algorithm.Log("fresh stock / available % "+str(AvailablePorfolioPercent))
                    else:
                        break

            for stockToBuy in stocksToBuy:
                insights.append(Insight.Price(stockToBuy[0].Symbol, timedelta(days=1), InsightDirection.Up, None, 1.00, None, stockToBuy[2]))
                # magnitude, confidence, weight
            
            self.previousStocksBought = stocksToBuy
            
            if self.isPositionsRebalancingWeek == True:
                self.isPositionsRebalancingWeek = False
            else:
                self.isPositionsRebalancingWeek = True

            return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        # clean up data for removed securities
        symbols = [x.Symbol for x in changes.RemovedSecurities]
        if len(symbols) > 0:
            
            for subscription in algorithm.SubscriptionManager.Subscriptions:
                if subscription.Symbol in symbols:
                    self.symbolDataBySymbol.pop(subscription.Symbol, None)
                    subscription.Consolidators.Clear()

        # initialize data for added securities
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        if len(addedSymbols) == 0: return

        SMAhistory = algorithm.History(addedSymbols, self.SMAperiod, self.resolution)
        ATRhistory = algorithm.History(addedSymbols, self.ATRperiod, self.resolution)

        for symbol in addedSymbols:
            sma = algorithm.SMA(symbol, self.SMAperiod, self.resolution)
            
            if not SMAhistory.empty:
                ticker = SymbolCache.GetTicker(symbol)

                if ticker not in SMAhistory.index.levels[0]:
                    Log.Trace(f'SMA added on securities changed: {ticker} not found in history data frame.')
                    continue

                for tuple in SMAhistory.loc[ticker].itertuples():
                    sma.Update(tuple.Index, tuple.close)
                    
            atr = algorithm.ATR(symbol, self.ATRperiod, MovingAverageType.Simple, self.resolution)
            
            if not ATRhistory.empty:
                ticker = SymbolCache.GetTicker(symbol)

                if ticker not in ATRhistory.index.levels[0]:
                    Log.Trace(f'ATR added on securities changed: {ticker} not found in history data frame.')
                    continue
                
                for tuple in ATRhistory.loc[ticker].itertuples():
                    bar = TradeBar(tuple.Index, symbol, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)     
                    atr.Update(bar)

            self.symbolDataBySymbol[symbol] = SymbolData(symbol, sma, atr)

class SymbolData:
    def __init__(self, symbol, sma, atr):
        self.Symbol = symbol
        self.SMA = sma
        self.ATR = atr