#region imports
from AlgorithmImports import *
#endregion


class IntradayMomentumAlphaModel(AlphaModel):
    """
    This class emits insights to take positions for the last `return_bar_count` minutes 
    of the day in the direction of the return for the first `return_bar_count` minutes of the day.
    """
    _intraday_momentum_by_symbol = {}
    sign = lambda _, x: int(x and (1, -1)[x < 0])
    
    def __init__(self, algorithm, return_bar_count = 30):
        """
        Input:
         - return_bar_count
            Number of minutes to calculate the morning return over and the number of minutes
            to hold before the close (0 < return_bar_count < 195)
        """
        if return_bar_count <= 0 or return_bar_count >= 195:
            algorithm.quit(f"Requirement violated:  0 < return_bar_count < 195")
        self._return_bar_count = return_bar_count
    
    def update(self, algorithm, slice):
        """
        Called each time our alpha model receives a new data slice.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - data
            A data structure for all of an algorithm's data at a single time step
        
        Returns a list of Insights to the portfolio construction model
        """
        insights = []
        
        for symbol, intraday_momentum in self._intraday_momentum_by_symbol.items():
            if slice.contains_key(symbol) and slice[symbol] is not None:
                intraday_momentum.bars_seen_today += 1
        
                # End of the morning return
                if intraday_momentum.bars_seen_today == self._return_bar_count:
                    intraday_momentum.morning_return = (slice[symbol].close - intraday_momentum.yesterdays_close) / intraday_momentum.yesterdays_close
                
                ## Beginning of the close
                next_close_time = intraday_momentum.exchange.hours.get_next_market_close(slice.time, False)
                mins_to_close = int((next_close_time - slice.time).total_seconds() / 60)
                
                if mins_to_close == self._return_bar_count + 1:
                    insight = Insight.price(intraday_momentum.symbol, 
                                            next_close_time, 
                                            self.sign(intraday_momentum.morning_return))
                    insights.append(insight)
                    continue
                
                # End of the day
                if not intraday_momentum.exchange.date_time_is_open(slice.time):
                    intraday_momentum.yesterdays_close = slice[symbol].close
                    intraday_momentum.bars_seen_today = 0
        
        return insights
        
    def on_securities_changed(self, algorithm, changes):
        """
        Called each time our universe has changed.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and subtractions to the algorithm's security subscriptions
        """
        for security in changes.added_securities:
            self._intraday_momentum_by_symbol[security.symbol] = IntradayMomentum(security, algorithm)
            
        for security in changes.removed_securities:
            self._intraday_momentum_by_symbol.pop(security.symbol, None)
            

class IntradayMomentum:
    """
    This class manages the data used for calculating the morning return of a security.
    """
    def __init__(self, security, algorithm):
        """
        Input:
         - security
            The security to trade
         - algorithm
            Algorithm instance running the backtest
        """
        self.symbol = security.symbol
        self.exchange = security.exchange
        
        self.bars_seen_today = 0
        self.yesterdays_close = algorithm.history(self.symbol, 1, Resolution.DAILY).loc[self.symbol].close[0]
        self.morning_return = 0
        

#region imports
from AlgorithmImports import *
#endregion


class CloseOnCloseExecutionModel(ExecutionModel):
    """
    Provides an implementation of IExecutionModel that immediately submits a market order to achieve 
    the desired portfolio targets and an associated market on close order.
    """

    def __init__(self):
        self._targets_collection = PortfolioTargetCollection()
        self._invested_symbols = []

    def execute(self, algorithm, targets):
        """
        Immediately submits orders for the specified portfolio targets.
        Input:
         - algorithm
            Algorithm instance running the backtest
         - targets
            The portfolio targets to be ordered
        """
        # for performance we check count value, OrderByMarginImpact and ClearFulfilled are expensive to call
        self._targets_collection.add_range(targets)
        if self._targets_collection.count > 0:
            for target in self._targets_collection.order_by_margin_impact(algorithm):
                # calculate remaining quantity to be ordered
                quantity = OrderSizing.get_unordered_quantity(algorithm, target)
                if quantity == 0:
                    continue
                
                algorithm.market_order(target.symbol, quantity)
                algorithm.market_on_close_order(target.symbol, -quantity)
                
            self._targets_collection.clear_fulfilled(algorithm)
        

#region imports
from AlgorithmImports import *

from alpha import IntradayMomentumAlphaModel
from execution import CloseOnCloseExecutionModel
#endregion


class IntradayETFMomentumAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        
        tickers = [
            'SPY',  # S&P 500
            'IWM',  # Russell 2000
            'IYR'   # Real Estate ETF
        ]
        symbols = [ Symbol.create(ticker, SecurityType.EQUITY, Market.USA) for ticker in tickers ]
        self.set_universe_selection(ManualUniverseSelectionModel(symbols))
        
        self.set_alpha(IntradayMomentumAlphaModel(self))
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(lambda _: None))
        self.set_execution(CloseOnCloseExecutionModel())