#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/strategies/the-52-week-high-and-short-term-reversal-in-stock-returns/
#
# The investment universe consists of all common stocks (share codes 10 or 11) trading at the NYSE, AMEX and NASDAQ. Exclude stocks with prices
# that are less than $5 at the end of month t. Divide stocks into three groups: the micro-cap, small-cap, and large-cap firms, using the NYSE 20th
# and 50th percentile as breakpoints. From now on focus purely on the large-cap firms. Firstly, compute PTH, the ratio of the current price of a
# stock to its highest price within the past 52 weeks to measure nearness to the 52-week high. At the end of month t, stocks are sorted into 
# quintile portfolios based on their past 1-month returns. Secondly, stocks are also sorted into quintile portfolios based on their PTH ranking
# based on the price information up to month t−1. The intersection of these reversal and PTH portfolios produces 25 portfolios. Long low PTH past
# losers portfolio and short low PTH past winners portfolio. Portfolio is equally weighted and rebalanced on a monthly basis.
#
# QC implementation changes:
#   - Instead of all listed stock, we select 500 most liquid stocks traded on NYSE, AMEX, or NASDAQ.
#   - Instead of historical stock highs, stock closes are used as a proxy.

class ReversalCombinedwithVolatility(QCAlgorithm):

    def Initialize(self):
        # self.SetStartDate(2000, 1, 1)
        self.SetStartDate(2017, 12, 1)
        # self.SetEndDate(2017, 1, 31)
        self.SetEndDate(2024, 9, 30)
        self.SetCash(100000)

        self.saltare_allocation = 0.25
        self.max_short_size = 0.05

        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

        self.coarse_count = 500
        
        self.data = {}
        self.period = 52 * 5
        
        self.symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        
        self.long = []
        self.short = []

        # self.meme_months = [6, 7]
        self.meme_months = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
        # self.restricted_stocks = ['CVNA', 'MARA', 'PLUG', 'AI', 'NIO', 'RIVN', 'DOCU']
        # self.restricted_stocks = ['NKLA', 'CVNA', 'AMC', 'BYND', 'UPST', 'MARA', 'RIOT', 
        #                         'COIN', 'FUBO', 'HUT', 'SNAP', 'FSR', 'SPCE', 'TLRY', 'PACW', 
        #                         'SOFI', 'ENVX', 'AFRM', 'GME', 'PLUG', 'ASTS', 'MULN', 'SIRI', 
        #                         'LCID', 'CCL', 'LAZR', 'CLOV', 'MSTR', 'RKLB', 'XXII', 'VRM', 
        #                         'RGTI', 'EOSE', 'MVST', 'RBT', 'EDTX', 'MVIS', 'IONQ', 'NVAX', 
        #                         'BRCC', 'TMC', 'QBTS', 'CHPT', 'PGY', 'DWAC', 'NVTA', 'SQL', 
        #                         'DNA', 'WISH', 'CPA']

        # This list is produced using both a Quiver Quant and a Roundhill MEME ETF Holdings Scrape
        # 'https://www.quiverquant.com/scores/memestocks'
        # 'https://www.roundhillinvestments.com/etf/meme/'
        self.restricted_stocks = ['HOOD', 'DNA', 'SOFI', 'TLRY', 'VRM', 'SPCE', 'EDTX', 'ENVX', 
                                'MULN', 'AMC', 'PLTR', 'CHPT', 'MSTR', 'PGY', 'UPST', 'EOSE', 
                                'CVNA', 'IONQ', 'XXII', 'SIRI', 'MVST', 'CLOV', 'BRCC', 'LCID', 
                                'HUT', 'ASTS', 'SQL', 'ENPH', 'TSLA', 'DWAC', 'DAL', 'BYND', 'CPA', 
                                'QBTS', 'NKLA', 'IBM', 'RGTI', 'RIVN', 'NVTA', 'TMC', 'GME', 'LAZR', 
                                'MARA', 'MVIS', 'AAL', 'PLUG', 'FUBO', 'COIN', 'RBT', 'RIOT', 'RKLB', 
                                'AFRM', 'FSR', 'SNAP', 'CCL', 'NVAX', 'NIO', 'MAT', 'PACW', 'WISH', 
                                'DISH']

        # self.restricted_stocks = ['HOOD', 'DNA', 'SOFI', 'TLRY', 'VRM', 'SPCE', 'EDTX', 'ENVX', 
        #                         'MULN', 'AMC', 'PLTR', 'CHPT', 'MSTR', 'PGY', 'UPST', 'EOSE', 
        #                         'CVNA', 'IONQ', 'XXII', 'SIRI', 'MVST', 'CLOV', 'BRCC', 'LCID', 
        #                         'HUT', 'ASTS', 'SQL', 'ENPH', 'TSLA', 'DWAC', 'DAL', 'BYND', 'CPA', 
        #                         'QBTS', 'NKLA', 'IBM', 'RGTI', 'RIVN', 'NVTA', 'TMC', 'GME', 'LAZR', 
        #                         'MARA', 'MVIS', 'AAL', 'PLUG', 'FUBO', 'COIN', 'RBT', 'RIOT', 'RKLB', 
        #                         'AFRM', 'FSR', 'SNAP', 'CCL', 'NVAX', 'NIO', 'MAT', 'PACW', 'WISH', 
        #                         'DISH', 'MVST', 'NVAX', 'QBTS', 'MARA', 'TSLA', 'CAVA', 'BYND', 'FSR', 
        #                         'SNAP', 'SMCI', 'RBLX', 'LAZR', 'BRCC', 'CPA', 'RIVN', 'FUBO', 'ASTS', 
        #                         'ENVX', 'MSTR', 'SNOW', 'ZM', 'DWAC', 'NKLA', 'PLTR', 'COIN', 'MRNA', 
        #                         'MVIS', 'AMC', 'RIOT', 'CLOV', 'RGTI', 'MULN', 'GME', 'SOFI', 'TMC', 
        #                         'SIRI', 'EDTX', 'HUT', 'JNJ', 'IONQ', 'RBT', 'CVNA', 'NVTA', 'TLRY', 
        #                         'CCL', 'CLF', 'PLUG', 'PACW', 'AFRM', 'RKLB', 'SE', 'EOSE', 'LCID', 
        #                         'PGY', 'ULTA', 'SQL', 'DNA', 'PTON', 'WISH', 'UPST', 'PANW', 'HOOD', 
        #                         'DKS', 'SPCE', 'VRM', 'CHPT', 'XXII']

        self.meme_etf = self.AddEquity('MEME', Resolution.Daily).Symbol
        
        self.selection_flag = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.Schedule.On(self.DateRules.MonthEnd(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol), self.Selection)
        
        self.settings.daily_precise_end_time = False

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            # security.SetFeeModel(CustomFeeModel(self))
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(10)

    def CoarseSelectionFunction(self, coarse):
        # Update the rolling window every day.
        for stock in coarse:
            symbol = stock.Symbol

            # Store monthly price.
            if symbol in self.data:
                self.data[symbol].update(stock.AdjustedPrice)

        if not self.selection_flag:
            return Universe.Unchanged

        # selected = [x.Symbol for x in coarse if x.HasFundamentalData and x.Market == 'usa']
        selected = [x.Symbol
            for x in sorted([x for x in coarse if x.HasFundamentalData and x.Market == 'usa'],
                key = lambda x: x.DollarVolume, reverse = True)[:self.coarse_count]]
        
        # Warmup price rolling windows.
        for symbol in selected:
            if symbol in self.data:
                continue
            
            self.data[symbol] = SymbolData(symbol, self.period)
            history = self.History(symbol, self.period, Resolution.Daily)
            if history.empty:
                self.Log(f"Not enough data for {symbol} yet.")
                continue
            closes = history.loc[symbol].close
            # for time, close in closes.iteritems():
            for time, close in closes.items():
                self.data[symbol].update(close)
            
        return [x for x in selected if self.data[x].is_ready()]

    def FineSelectionFunction(self, fine):
        fine = [x for x in fine if ((x.SecurityReference.ExchangeId == "NYS") or (x.SecurityReference.ExchangeId == "NAS") or (x.SecurityReference.ExchangeId == "ASE"))]

        # if len(fine) > self.coarse_count:
        #     sorted_by_market_cap = sorted(fine, key = lambda x: x.MarketCap, reverse=True)
        #     top_by_market_cap = sorted_by_market_cap[:self.coarse_count]
        # else:
        #     top_by_market_cap = fine
        top_by_market_cap = fine
        
        pth_performance = {x.Symbol : (self.data[x.Symbol].pth(), self.data[x.Symbol].performance()) for x in top_by_market_cap}

        sorted_by_pth = sorted(pth_performance.items(), key = lambda x: x[1][0], reverse = True)
        sorted_by_pth = [x[0] for x in sorted_by_pth]
        
        sorted_by_ret = sorted(pth_performance.items(), key = lambda x: x[1][1], reverse = True)
        sorted_by_ret = [x[0] for x in sorted_by_ret]

        quintile  = int(len(sorted_by_ret) / 5)
        
        low_pth = sorted_by_pth[-quintile:]
        top_ret = sorted_by_ret[:quintile]
        low_ret = sorted_by_ret[-quintile:]
        
        self.long = [x for x in low_pth if x in low_ret]
        self.short = [x for x in low_pth if x in top_ret]
        
        return self.long + self.short

    def OnData(self, data):
        if not self.selection_flag:
            return
        self.selection_flag = False
        
        # Trade execution.
        long_count = len(self.long)
        short_count = len(self.short)

        stocks_invested = [x.Key for x in self.Portfolio if x.Value.Invested]
        for symbol in stocks_invested:
            if symbol not in self.long + self.short:
                self.Liquidate(symbol)

        for symbol in self.long:
            self.SetHoldings(symbol, 1 / long_count)
            # self.Log(str(symbol.Value))
        for symbol in self.short:
            if self.Time.year == 2023:
                if self.Time.month in self.meme_months:
                    if symbol.Value not in self.restricted_stocks:
                        # self.SetHoldings(symbol, -1 / short_count)
                        # self.SetHoldings(symbol, max(-0.20, -1 / short_count))
                        self.SetHoldings(symbol, max(-(self.max_short_size/self.saltare_allocation), -1 / short_count))
                        # self.SetHoldings(symbol, max(-0.12, -1 / short_count))
            else:
                # self.SetHoldings(symbol, -1 / short_count)
                # self.SetHoldings(symbol, max(-0.20, -1 / short_count))
                self.SetHoldings(symbol, max(-(self.max_short_size/self.saltare_allocation), -1 / short_count))
                # self.SetHoldings(symbol, max(-0.12, -1 / short_count))
            # self.SetHoldings(symbol, -1 / short_count)
            
        # if self.Time.year == 2022 or self.Time.year == 2023:
        # if self.Time.year == 2023:
            # for symbol in self.short:
                # self.SetHoldings(symbol, -1 / short_count)
                # self.SetHoldings(symbol, -0.5 / short_count)
            # self.SetHoldings(self.meme_etf, 0.1)
        # else:
            # for symbol in self.short:
                # self.SetHoldings(symbol, -1 / short_count)

        self.long.clear()
        self.short.clear()

    def Selection(self):
        self.selection_flag = True

class SymbolData():
    def __init__(self, symbol, period):
        self.Symbol = symbol
        self.Price = RollingWindow[float](period)
    
    def update(self, value):
        self.Price.Add(value)
    
    def is_ready(self) -> bool:
        return self.Price.IsReady
    
    def pth(self):
        high_proxy = [x for x in self.Price]
        symbol_price = high_proxy[0]
        return  symbol_price / max(high_proxy[21:])
    
    def performance(self) -> float:
        closes = [x for x in self.Price][:21]
        return (closes[0] / closes[-1] - 1)

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