import numpy as np  
import pandas as pd
from itertools import groupby
from math import ceil

class CalmAsparagusAnt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 3, 11)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        ief = self.AddEquity("IEF", Resolution.Daily).Symbol
        tlt = self.AddEquity("TLT", Resolution.Daily).Symbol
        self.UniverseSettings.Resolution = Resolution.Daily
        self.BONDS = [ief, tlt]
        self.TARGET_SECURITIES = 25  
        self.TOP_ROE_QTY = 100 #First sort by ROE
        self.numberOfSymbolsCoarse = 3000
        self.numberOfSymbolsFine = 1500
        self.dollarVolumeBySymbol = {}
        self.activeUniverse = []
        self.lastMonth = -1
        self.trend_up = False
        
        # This is for the trend following filter  
        self.SPY = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        self.TF_LOOKBACK = 126  
        self.TF_CURRENT_LOOKBACK = 20
        
        history = self.History(self.SPY, self.TF_LOOKBACK, Resolution.Daily)
        self.spy_ma50_slice = self.SMA(self.SPY, self.TF_CURRENT_LOOKBACK, Resolution.Daily)
        self.spy_ma200_slice = self.SMA(self.SPY, self.TF_LOOKBACK, Resolution.Daily)
        for tuple in history.loc[self.SPY].itertuples():
            self.spy_ma50_slice.Update(tuple.Index, tuple.close)
            self.spy_ma200_slice.Update(tuple.Index, tuple.close)    
        
        self.trend_up = self.spy_ma50_slice.Current.Value > self.spy_ma200_slice.Current.Value
    
        # This is for the determining momentum  
        self.MOMENTUM_LOOKBACK_DAYS = 126 #Momentum lookback  
        self.MOMENTUM_SKIP_DAYS = 2  
    
        # Initialize any other variables before being used  
        self.stock_weights = pd.Series()  
        self.bond_weights = pd.Series()
    
        # Should probably comment out the slippage and using the default  
        # set_slippage(slippage.FixedSlippage(spread = 0.0))  
        # Create and attach pipeline for fetching all data  
        # Schedule functions  
        # Separate the stock selection from the execution for flexibility  
        
        self.Schedule.On(self.DateRules.MonthEnd("SPY", 0),
                 self.TimeRules.BeforeMarketClose(self.SPY, 0),       
                 self.SelectStocksAndSetWeights)
 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose(self.SPY, 0),       
                 self.RecordVars)
                 
                 
        self.lastMonth = -1
        self.AddUniverse(self.Coarse, self.Fine)

    def OnData(self, data):
        pass
    
    def SelectStocksAndSetWeights(self):
        # Get pipeline output and select stocks  
        #df = algo.pipeline_output('pipeline')  
        ''' Pick a new Universe ?????? every week '''
        
        current_holdings = [x.Key for x in self.Portfolio if x.Value.Invested]
        # Define our rule to open/hold positions  
        # top momentum and don't open in a downturn but, if held, then keep  
        rule = 'top_quality_momentum & (trend_up or (not trend_up & index in @current_holdings))'  
        stocks_to_hold = self.activeUniverse
        self.trend_up = self.spy_ma50_slice.Current.Value > self.spy_ma200_slice.Current.Value
        if self.trend_up == False:
            return
        # Set desired stock weights  
        # Equally weight  
        stock_weight = 1.0 / (self.TARGET_SECURITIES)  
        self.weights = {}
        for x in stocks_to_hold:
            self.weights[x] = stock_weight
        # Set desired bond weight  
        # Open bond position to fill unused portfolio balance  
        # But always have at least 1 'share' of bonds  
        ### bond_weight = max(1.0 - context.stock_weights.sum(), stock_weight) / len(context.BONDS)  
        bond_weight = (1.0 - stock_weight) / len(self.BONDS)  
        for x in self.BONDS:
            self.weights[x] = bond_weight
        self.Trade()
    
    def Trade(self):
        makeInvestments = sorted([x for x in self.weights.keys()], key = lambda x: self.weights[x], reverse = False)
        total = 0
        for stock in makeInvestments:
            weight = self.weights[stock]
            if weight == 0:
                self.Liquidate(stock)
            else:
                self.SetHoldings(stock, weight)
            total += weight
        self.Debug(total)
    
    def RecordVars(self):
        pass
    
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol == self.SPY or symbol in self.BONDS:
                continue
            if symbol not in self.activeUniverse:
                self.activeUniverse.append(symbol)
            
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            self.Liquidate(symbol)
            if symbol in self.activeUniverse:
                self.activeUniverse.remove(symbol)

    def Coarse(self, coarse):
        if self.lastMonth == self.Time.month:
            return Universe.Unchanged
        
        sortedByDollarVolume = sorted([x for x in coarse 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 Fine(self, fine):
        sortedBySector = sorted([x for x in fine if x.CompanyReference.CountryId == "USA"
                                        and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                        and (self.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 = self.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)
        Q1500US = [x for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]
        
        df = {}
        stocks = []
        df["Stocks"] = []
        df["Cash Return"] = []
        df["FcfYield"] = []
        df["ROIC"] = []
        df["LtdToEq"] = []
        for x in Q1500US:
            symbol = x.Symbol
            stocks.append(symbol)
            df["Stocks"].append(symbol)
            df["Cash Return"].append(x.ValuationRatios.CashReturn)
            df["FcfYield"].append(x.ValuationRatios.FCFYield)
            df["ROIC"].append(x.OperationRatios.ROIC.ThreeMonths)
            df["LtdToEq"].append(x.OperationRatios.LongTermDebtEquityRatio.ThreeMonths)
        df = pd.DataFrame.from_dict(df)
        df["Cash Return"] = df["Cash Return"].rank()
        df["FcfYield"] = df["FcfYield"].rank()
        df["ROIC"] = df["ROIC"].rank()
        df["LtdToEq"] = df["LtdToEq"].rank()
        df["Value"] = (df["Cash Return"] + df["FcfYield"]).rank()
        df["Quality"] = df["ROIC"] + df["LtdToEq"] + df["Value"]
        
        top_quality = df.sort_values(by=["Quality"])
        top_quality = top_quality["Stocks"][:self.TOP_ROE_QTY]
        
        topStocks = [x for x in top_quality]
        history = self.History(topStocks, self.MOMENTUM_LOOKBACK_DAYS+self.MOMENTUM_SKIP_DAYS, Resolution.Daily)
        returns_overall = {}
        returns_recent = {}
        for symbol in topStocks:
            mompLong = MomentumPercent(self.MOMENTUM_LOOKBACK_DAYS+self.MOMENTUM_SKIP_DAYS)
            mompShort = MomentumPercent(self.MOMENTUM_SKIP_DAYS)
            for tuple in history.loc[symbol].itertuples():
                mompLong.Update(tuple.Index, tuple.close)
                mompShort.Update(tuple.Index, tuple.close)
            returns_overall[symbol] = mompLong.Current.Value
            returns_recent[symbol] = mompShort.Current.Value
        
        final = sorted(topStocks, key = lambda x: returns_overall[x] - returns_overall[x], reverse = True)
        final = final[:self.TARGET_SECURITIES]
        return final