#region imports
from AlgorithmImports import *
from utils import *
#endregion

class BreakoutForecastAlphaModel(AlphaModel):

    def __init__(self, algorithm, symbols, periods, multiplier=1):
        self.algorithm = algorithm
        self.symbols = symbols
        self.periods = periods
        self.multiplier = multiplier
        self.day = -1

    def update(self, algorithm, data):
        if self.day == algorithm.time.day:
            return []
        self.day = algorithm.time.day

        signals = {period: {} for period in self.periods}
        period_weight = 1 / len(self.periods)

        # Get signals for each Future
        for period in self.periods:
            for symbol in self.symbols:
                if data.bars.contains_key(symbol):
                    if self.get_breakout_signal(algorithm, symbol, period):
                        signals[period][symbol] = 1
                    else:
                        signals[period][symbol] = -1

        return get_insights(algorithm, signals, period_weight, self.multiplier)

    def get_breakout_signal(self, algorithm, symbol, period):
        security = algorithm.securities[symbol.canonical]
        return security[f"EMA{period}"].current.value > 0
        
    def on_securities_changed(self, algorithm, changes):
        for removed in changes.removed_securities:
            if removed.type != SecurityType.FUTURE or not removed.symbol.canonical in self.symbols: 
                continue
            
            symbol = removed.symbol.canonical
            security = algorithm.securities[symbol]
            algorithm.subscription_manager.remove_consolidator(symbol, security["consolidator"])
            for period in self.periods:
                security[f"Maximum{period}"].reset()
                security[f"Minimum{period}"].reset()
                security[f"EMA{period}"].reset()

        for added in changes.added_securities:
            if added.type != SecurityType.FUTURE or not added.symbol.canonical in self.symbols: 
                continue
            
            security = algorithm.securities[added.symbol.canonical]
            security["consolidator"] = TradeBarConsolidator(timedelta(1))
            for period in self.periods:
                security[f"Maximum{period}"] = Maximum(period)
                security[f"Minimum{period}"] = Minimum(period)
                security[f"EMA{period}"] = ExponentialMovingAverage(period // 4)
            security["consolidator"] = reset_and_warm_up_breakout(algorithm, security, self.periods)
            security["consolidator"].data_consolidated += self.on_consolidated
            algorithm.subscription_manager.add_consolidator(added.symbol.canonical, security["consolidator"])

    def on_consolidated(self, _, bar):
        security = self.algorithm.securities[bar.symbol.canonical]
        for period in self.periods:
            max_ = security[f'Maximum{period}'].current.value
            min_ = security[f'Minimum{period}'].current.value
            if not max_ or not min_ or max_ == min_:
                continue
            mean = (max_ + min_) / 2
            security[f'EMA{period}'].update(bar.end_time, 40 * (bar.close - mean) / (max_ - min_))

# region imports
from AlgorithmImports import *
from universe import FrontMonthFutureUniverseSelectionModel
from breakout_alpha import BreakoutForecastAlphaModel
from trend_following_alpha import TrendFollowingAlphaModel
# endregion

class BreakoutFuturesAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2019, 1, 1)
        self.set_end_date(2023, 6, 30)
        self.SetCash(10000000000)

        breakout_weight = self.get_parameter("breakout_weight", 0.5)
        trend_following_weight = self.get_parameter("trend_following_weight", 0.5)

        self.symbols = [
            Symbol.create(Futures.Grains.WHEAT, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Grains.SOYBEANS, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Meats.FEEDER_CATTLE, SecurityType.FUTURE, Market.CME),
            Symbol.create(Futures.Softs.COTTON_2, SecurityType.FUTURE, Market.ICE),
            Symbol.create(Futures.Financials.Y_30_TREASURY_BOND, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Financials.Y_10_TREASURY_NOTE, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Financials.Y_2_TREASURY_NOTE, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Financials.FIVE_YEAR_USDMAC_SWAP, SecurityType.FUTURE, Market.CBOT),
            Symbol.create(Futures.Indices.SP_500_E_MINI, SecurityType.FUTURE, Market.CME),
            Symbol.create(Futures.Indices.MSCI_EUROPE_NTR, SecurityType.FUTURE, Market.NYSELIFFE),
            Symbol.create(Futures.Indices.NIKKEI_225_DOLLAR, SecurityType.FUTURE, Market.CME),
            Symbol.create(Futures.Currencies.BTC, SecurityType.FUTURE, Market.CME),
            Symbol.create(Futures.Currencies.ETH, SecurityType.FUTURE, Market.CME),
            Symbol.create(Futures.Metals.GOLD, SecurityType.FUTURE, Market.COMEX),
            Symbol.create(Futures.Metals.PLATINUM, SecurityType.FUTURE, Market.NYMEX),
            Symbol.create(Futures.Metals.PALLADIUM, SecurityType.FUTURE, Market.NYMEX),
            Symbol.create(Futures.Energies.NATURAL_GAS, SecurityType.FUTURE, Market.NYMEX),
            Symbol.create(Futures.Energies.CRUDE_OIL_WTI, SecurityType.FUTURE, Market.NYMEX),
            Symbol.create(Futures.Energies.BRENT_CRUDE, SecurityType.FUTURE, Market.ICE),
            Symbol.create(Futures.Indices.VIX, SecurityType.FUTURE, Market.CFE)
        ]

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        
        self.universe_settings.extended_market_hours = True
        self.universe_settings.data_normalization_mode = DataNormalizationMode.BACKWARDS_RATIO
        self.universe_settings.data_mapping_mode = DataMappingMode.OPEN_INTEREST
        self.universe_settings.contract_depth_offset = 0
        self.universe_settings.resolution = Resolution.MINUTE
        
        # Seed initial price data
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        # We only want front month contract
        self.add_universe_selection(FrontMonthFutureUniverseSelectionModel(self.select_future_chain_symbols))

        # breakout +/- trend following strategies
        breakout_periods = [5, 10, 20, 40, 80, 160, 320]
        self.add_alpha(BreakoutForecastAlphaModel(self, self.symbols, breakout_periods, breakout_weight))
        trend_periods = [2, 4, 8, 16, 32, 64, 128]
        self.add_alpha(TrendFollowingAlphaModel(self.symbols, trend_periods, trend_following_weight))

        # Position sizing is controlled by alpha model
        self.set_portfolio_construction(InsightWeightingPortfolioConstructionModel())

        # Avoid massive bid-ask spread friction
        self.set_execution(SpreadExecutionModel(0.03))

        # Set benchmark: since it is a buy and hold portfolio, using SPY will be appropriate
        self.add_equity("SPY")
        self.set_benchmark("SPY")

    def select_future_chain_symbols(self, utcTime):
        return self.symbols

#region imports
from AlgorithmImports import *
from utils import *
#endregion

class TrendFollowingAlphaModel(AlphaModel):

    def __init__(self, symbols, periods, multiplier=1):
        self.symbols = symbols
        self.periods = periods
        self.multiplier = multiplier
        self.month = -1

    def update(self, algorithm, data):
        if self.month == algorithm.time.month:
            return []
        self.month = algorithm.time.month

        signals = {period: {} for period in self.periods}
        period_weight = 1 / len(self.periods)

        # Get signals for each Future
        for period in self.periods:
            for symbol in self.symbols:
                if data.bars.contains_key(symbol):
                    if self.get_trend_signal(algorithm, symbol, period):
                        signals[period][symbol] = 1
                    else:
                        signals[period][symbol] = -1

        return get_insights(algorithm, signals, period_weight, self.multiplier)

    def get_trend_signal(self, algorithm, symbol, period):
        security = algorithm.securities[symbol.canonical]
        return security[f"EMA{period}"].current.value < security.price
    
    def on_securities_changed(self, algorithm, changes):
        for removed in changes.removed_securities:
            if removed.type != SecurityType.FUTURE or not removed.symbol.canonical in self.symbols: 
                continue
            
            symbol = removed.symbol.canonical
            security = algorithm.securities[symbol]
            for period in self.periods:
                security[f"EMA{period}"].reset()

        for added in changes.added_securities:
            if added.type != SecurityType.FUTURE or not added.symbol.canonical in self.symbols: 
                continue
            
            security = algorithm.securities[added.symbol.canonical]
            security["consolidator"] = TradeBarConsolidator(timedelta(1))
            for period in self.periods:
                security[f"EMA{period}"] = ExponentialMovingAverage(period)
            security["consolidator"] = reset_and_warm_up_trend_following(algorithm, security, self.periods)

#region imports
from AlgorithmImports import *
from Selection.FutureUniverseSelectionModel import FutureUniverseSelectionModel
#endregion

class FrontMonthFutureUniverseSelectionModel(FutureUniverseSelectionModel):
    '''Creates futures chain universes that select the front month contract and runs a user
    defined futureChainSymbolSelector every day to enable choosing different futures chains'''
    def __init__(self, select_future_chain_symbols, rebalance_period = 7):
        super().__init__(timedelta(rebalance_period), select_future_chain_symbols)

    def filter(self, filter):
        '''Defines the futures chain universe filter'''
        return (filter.front_month().only_apply_filter_at_market_open())

#region imports
from AlgorithmImports import *
#endregion

def reset_and_warm_up_breakout(algorithm, security, periods):
    consolidator = security['consolidator']
    lookback = max(periods)

    # historical request to update the consolidator that will warm up the indicator
    history = algorithm.history[consolidator.input_type](security.symbol, lookback, Resolution.DAILY,
        data_normalization_mode = DataNormalizationMode.SCALED_RAW)
    
    # Replace the consolidator, since we cannot reset it
    # Not ideal since we don't the consolidator type and period
    algorithm.subscription_manager.remove_consolidator(security.symbol, consolidator)
    consolidator = TradeBarConsolidator(timedelta(1))
    for period in periods:
        for indicator in [security[f'Maximum{period}'], security[f'Minimum{period}']]:
            indicator.reset()
            algorithm.register_indicator(security.symbol, indicator, consolidator)
        security[f'EMA{period}'].reset()
    
    for bar in list(history)[:-1]:
        consolidator.update(bar)
        for period in periods:
            max_ = security[f'Maximum{period}'].current.value
            min_ = security[f'Minimum{period}'].current.value
            mean = (max_ + min_) / 2
            security[f'EMA{period}'].update(bar.end_time, 40 * (bar.close - mean) / (max_ - min_))

    return consolidator

def reset_and_warm_up_trend_following(algorithm, security, periods):
    indicators = [security[f'EMA{period}'] for period in periods]
    consolidator = security['consolidator']
    lookback = max(periods)

    # historical request to update the consolidator that will warm up the indicator
    history = algorithm.history[consolidator.input_type](security.symbol, lookback, Resolution.DAILY,
        data_normalization_mode = DataNormalizationMode.SCALED_RAW)
    
    # Replace the consolidator, since we cannot reset it
    # Not ideal since we don't the consolidator type and period
    algorithm.subscription_manager.remove_consolidator(security.symbol, consolidator)
    consolidator = TradeBarConsolidator(timedelta(1))
    for indicator in indicators:
        indicator.reset()
        algorithm.register_indicator(security.symbol, indicator, consolidator)
    
    for bar in list(history)[:-1]:
        consolidator.update(bar)

    return consolidator

def get_insights(algorithm, signals, period_weight, alpha_weight=1):
    # Equal risk capital per signal group
    symbol_signals = {}
    for period, selected in signals.items():
        if not selected:
            continue
        symbol_weight = 1 / len(selected)
        
        for symbol, sign in selected.items():
            mapped = algorithm.securities[symbol].mapped
            multiplier = algorithm.securities[symbol].symbol_properties.contract_multiplier
            weight = sign * period_weight * symbol_weight * alpha_weight / multiplier
            if mapped not in symbol_signals:
                symbol_signals[mapped] = 0
            symbol_signals[mapped] += weight

    return [
        Insight.price(mapped, Expiry.END_OF_WEEK, InsightDirection.UP, weight = weight)
        for mapped, weight in symbol_signals.items() if abs(weight) >= 0.005
    ]