from AlgorithmImports import *

class IntradayArbitrage(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)

        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.iwm = self.AddEquity("IWM", Resolution.Minute).Symbol

        self.pairRatio = None
        self.deviationThreshold = 0.002  # Example threshold

        self.SetWarmup(10, Resolution.Daily)

    def OnData(self, data):
        if not self.Securities[self.spy].Price or not self.Securities[self.iwm].Price:
            return

        current_ratio = self.Securities[self.spy].Price / self.Securities[self.iwm].Price

        if self.pairRatio is None:
            self.pairRatio = current_ratio
            return

        deviation = (current_ratio - self.pairRatio) / self.pairRatio

        if abs(deviation) > self.deviationThreshold:
            if deviation > 0:
                self.SetHoldings(self.spy, 0.5)
                self.SetHoldings(self.iwm, -0.5)
            else:
                self.SetHoldings(self.spy, -0.5)
                self.SetHoldings(self.iwm, 0.5)

        self.pairRatio = current_ratio

    def OnEndOfDay(self):
        self.Liquidate()

from AlgorithmImports import *
import numpy as np 

class IntradayArbitrage(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)

        self.symbols = [self.AddEquity(ticker, Resolution.Second).Symbol for ticker in ['IVV', 'SPY']]
        self.spread_adjusters = [0, 0]
        self.entry_timer = [0, 0]
        self.exit_timer = [0, 0]
        self.long_side = -1
        self.history = {symbol: {'bids': np.zeros(400), 'asks': np.zeros(400)} for symbol in self.symbols}
        self.order_delay = 3
        self.pct_threshold = 0.02 / 100
        self.window_size = 400

        self.SetWarmup(self.window_size, Resolution.Daily)

        for symbol in self.symbols:
            consolidator = QuoteBarConsolidator(1)
            consolidator.DataConsolidated += self.CustomDailyHandler
            self.SubscriptionManager.AddConsolidator(symbol, consolidator)
    
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            if security.Symbol not in self.symbols:
                continue
            self.history[security.Symbol] = {'bids': np.zeros(self.window_size), 'asks': np.zeros(self.window_size)}

    def CustomDailyHandler(self, sender, consolidated):
        self.history[consolidated.Symbol]['bids'] = np.append(self.history[consolidated.Symbol]['bids'][1:], consolidated.Bid.Close)
        self.history[consolidated.Symbol]['asks'] = np.append(self.history[consolidated.Symbol]['asks'][1:], consolidated.Ask.Close)
        self.update_spread_adjusters()

    def update_spread_adjusters(self):
        for i in range(2):
            numerator_history = self.history[self.symbols[i]]['bids']
            denominator_history = self.history[self.symbols[abs(i-1)]]['asks']
            self.spread_adjusters[i] = (numerator_history / denominator_history).mean()

    def OnData(self, data):
        if self.IsWarmingUp: return
        quotebars = {symbol: data[symbol] for symbol in self.symbols}

        # Search for entries
        for i in range(2):
            if quotebars[self.symbols[abs(i-1)]].Bid.Close / quotebars[self.symbols[i]].Ask.Close - self.spread_adjusters[abs(i-1)] >= self.pct_threshold:
                self.entry_timer[i] += 1
                if self.entry_timer[i] == self.order_delay:
                    self.exit_timer = [0, 0]
                    if self.long_side == i:
                        return
                    self.long_side = i
                    self.SetHoldings(self.symbols[i], 0.5)
                    self.SetHoldings(self.symbols[abs(i-1)], -0.5)
                else:
                    return
            self.entry_timer[i] = 0

        # Search for an exit
        if self.long_side >= 0: # In a position
            if quotebars[self.long_side].Bid.Close / quotebars[abs(self.long_side-1)].Ask.Close - self.spread_adjusters[self.long_side] >= 0:
                self.exit_timer[self.long_side] += 1
                if self.exit_timer[self.long_side] == self.order_delay:
                    self.exit_timer[self.long_side] = 0
                    self.Liquidate()
                    self.long_side = -1
                else:
                    return

from AlgorithmImports import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
#This is a Template of dynamic stock selection.
#You can try your own fundamental factor and ranking method by editing the CoarseSelectionFunction and FineSelectionFunction
from QuantConnect.Data.UniverseSelection import *
class BasicTemplateAlgorithm(QCAlgorithm):
    
    def __init__(self):
    # set the flag for rebalance
        self.reb = 1
    # Number of stocks to pass CoarseSelection process
        self.num_coarse = 250
    # Number of stocks to long/short
        self.num_fine = 10
        self.symbols = None
        self.spy_momentum = None
        self.spy_performance = []
        self.init_cash = 100000
        
    def Initialize(self):
        self.SetCash(self.init_cash)
        self.SetStartDate(2017,1,1)        
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        
    # Schedule the rebalance function to execute at the begining of each month
        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
                         self.TimeRules.AfterMarketOpen(self.spy,5), 
                         Action(self.rebalance))
        
        self.spy_momentum = self.MOM(self.spy, 14, Resolution.Daily)
        
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.BeforeMarketClose(self.spy, 0), 
                         Action(self.record_vars))
        
    def CoarseSelectionFunction(self, coarse):
    # if the rebalance flag is not 1, return null list to save time.
        if self.reb != 1:
            return self.long + self.short
            
    # make universe selection once a month
    # drop stocks which have no fundamental data or have too low prices
        selected = [x for x in coarse if (x.HasFundamentalData) 
                    and (float(x.Price) > 5)]
        
        sortedByDollarVolume = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
        top = sortedByDollarVolume[:self.num_coarse]
        return [i.Symbol for i in top]
    def FineSelectionFunction(self, fine):
    # return null list if it's not time to rebalance
        if self.reb != 1:
            return self.long + self.short
            
        self.reb = 0
            
    # drop stocks which don't have the information we need.
    # you can try replacing those factor with your own factors here
    
        filtered_fine = [x for x in fine if x.OperationRatios.FCFGrowth.Value
                                        and x.ValuationRatios.PriceChange1M 
                                        and x.ValuationRatios.PERatio]
                                        
        self.Log('remained to select %d'%(len(filtered_fine)))
        
        # rank stocks by three factor.
        sortedByfactor1 = sorted(filtered_fine, key=lambda x: x.OperationRatios.FCFGrowth.Value, reverse=True)
        sortedByfactor2 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PriceChange1M, reverse=True)
        sortedByfactor3 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        
        stock_dict = {}
        
        # assign a score to each stock, you can also change the rule of scoring here.
        for i,ele in enumerate(sortedByfactor1):
            rank1 = i
            rank2 = sortedByfactor2.index(ele)
            rank3 = sortedByfactor3.index(ele)
            score = sum([rank1*0.2,rank2*0.4,rank3*0.4])
            stock_dict[ele] = score
        
        # sort the stocks by their scores
        self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
        sorted_symbol = [x[0] for x in self.sorted_stock]
        # sotre the top stocks into the long_list and the bottom ones into the short_list
        self.long = [x.Symbol for x in sorted_symbol[:self.num_fine]]
        self.short = [x.Symbol for x in sorted_symbol[-self.num_fine:]]
    
        return self.long + self.short
    def OnData(self, data):
        pass
    
    def rebalance(self):
        if self.spy_momentum.Current.Value <= 0:
            self.Log("SPY momentum is not positive. Skipping rebalancing.")
            return
        
    # if this month the stock are not going to be long/short, liquidate it.
        long_short_list = self.long + self.short
        for i in self.Portfolio.Values:
            if (i.Invested) and (i.Symbol not in long_short_list):
                self.Liquidate(i.Symbol)
                
    # Alternatively, can liquidate all the stocks at the end of each month.
    # Which method to choose depends on investment philosiphy 
    # if prefer to realized the gain/loss each month can choose this method.
    
        #self.Liquidate()
        
    # Assign each stock equally.
        for i in self.long:
            self.SetHoldings(i, 0.9/self.num_fine)
        
        for i in self.short:
            self.SetHoldings(i, -0.9/self.num_fine)
        self.reb = 1
            
    def record_vars(self):
        # Record the performance of SPY
        hist = self.History(self.spy, 2, Resolution.Daily)['close'].unstack(level=0).dropna()
        self.spy_performance.append(hist[self.spy].iloc[-1])
        
        if len(self.spy_performance) > 1:
            spy_perf = self.spy_performance[-1] / self.spy_performance[0] * self.init_cash
            self.Plot('Strategy Equity', 'SPY', spy_perf)
            
        # Record the performance of the portfolio
        portfolio_value = self.Portfolio.TotalPortfolioValue
        self.Plot('Strategy Equity', 'Portfolio', portfolio_value)