# region imports
from AlgorithmImports import *
class CustomImmediateExecutionModel(ImmediateExecutionModel):
    def __init__(self, leverage=2.0):
        self.leverage = leverage

    def Execute(self, algorithm, targets):
        leverage = self.leverage  # Assuming leverage of 2, adjust as needed

        for target in targets:
            if target.Symbol in algorithm.Portfolio:
                holding = algorithm.Portfolio[target.Symbol].Quantity
            else:
                holding = 0

            # Calculate the target quantity with leverage
            target_quantity = target.Quantity * leverage

            # Calculate the difference between target quantity and current holding
            adjustment = target_quantity - holding

            if target_quantity == 0:
                algorithm.Liquidate(target.Symbol, tag="Liquidate")
            elif adjustment > 0:
                if holding == 0:
                    tag = "new position"
                else:
                    tag = "Upsizing position"
                algorithm.MarketOrder(target.Symbol, adjustment, tag=tag)
            elif adjustment < 0:
                tag = "downsizing position"
                algorithm.MarketOrder(target.Symbol, adjustment, tag=tag)

#region imports
from AlgorithmImports import *
#endregion
class DualMomentumAlphaModel(AlphaModel):

    def __init__(self):
        self.sectors = {}
        self.securities_list = []
        self.day = -1

    def update(self, algorithm, data):

        insights = []

        for symbol in set(data.splits.keys() + data.dividends.keys()):
            security = algorithm.securities[symbol]
            if security in self.securities_list:
                security.indicator.reset()
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self._register_indicator(algorithm, security)

                history = algorithm.history[TradeBar](security.symbol, 7,
                                                      Resolution.DAILY,
                                                      data_normalization_mode=DataNormalizationMode.SCALED_RAW)
                for bar in history:
                    security.consolidator.update(bar)

        if data.quote_bars.count == 0:
            return []

        if self.day == algorithm.time.day:
            return []
        self.day = algorithm.time.day

        momentum_by_sector = {}
        security_momentum = {}

        for sector in self.sectors:
            securities = self.sectors[sector]
            security_momentum[sector] = {security: security.indicator.current.value
                              for security in securities if
                              security.symbol in data.quote_bars and security.indicator.is_ready}
            momentum_by_sector[sector] = sum(list(security_momentum[sector].values())) 

        target_sectors = [sector for sector in self.sectors if momentum_by_sector[sector] > 0]
        target_securities = []

        for sector in target_sectors:
            for security in security_momentum[sector]:
                if security_momentum[sector][security] > 0:
                    security.SetLeverage(2)
                    target_securities.append(security)

        target_securities = sorted(target_securities, key = lambda x: algorithm.securities[x.symbol].Fundamentals.MarketCap, reverse=True)[:10]

        for security in target_securities:
            insights.append(Insight.price(security.symbol, Expiry.END_OF_DAY, InsightDirection.UP))
        
        return insights

    def on_securities_changed(self, algorithm, changes):
        security_by_symbol = {}
        for security in changes.RemovedSecurities:
            if security in self.securities_list:
                algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
                self.securities_list.remove(security)
            for sector in self.sectors:
                if security in self.sectors[sector]:
                    self.sectors[sector].remove(security)

        for security in changes.AddedSecurities:
            
            sector = security.Fundamentals.AssetClassification.MorningstarSectorCode
            security_by_symbol[security.symbol] = security
            security.indicator = MomentumPercent(1)
            self._register_indicator(algorithm, security)
            self.securities_list.append(security)

            if sector not in self.sectors:
                self.sectors[sector] = set()
            self.sectors[sector].add(security)
            

            if security_by_symbol:
                history = algorithm.history[TradeBar](list(security_by_symbol.keys()), 7,
                                                  Resolution.DAILY,
                                                  data_normalization_mode=DataNormalizationMode.SCALED_RAW)
                for trade_bars in history:
                    for bar in trade_bars.values():
                        security_by_symbol[bar.symbol].consolidator.update(bar)

    def _register_indicator(self, algorithm, security):
        security.consolidator = TradeBarConsolidator(Calendar.WEEKLY)
        algorithm.subscription_manager.add_consolidator(security.symbol, security.consolidator)
        algorithm.register_indicator(security.symbol, security.indicator, security.consolidator)

# region imports
from AlgorithmImports import *
# endregion

# Your New Python File
# region imports
from AlgorithmImports import *

# endregion

class FredRate(PythonData):
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("https://fred.stlouisfed.org/graph/fredgraph.csv?id=DFF", SubscriptionTransportMedium.RemoteFile)

    def Reader(self, config, line, date, isLiveMode):
        data = line.split(',')
        if data[0] == 'DATE':
            return None
        rate = FredRate()
        rate.Symbol = config.Symbol
        rate.Time = datetime.strptime(data[0], "%Y-%m-%d")
        rate.Value = float(data[1])
        return rate

# Your New Python File

# region imports
from AlgorithmImports import *
class MyPcm(RiskParityPortfolioConstructionModel):
    def __init__(self, rebalance):
        super().__init__(rebalance)

    def CreateTargets(self, algorithm, insights):
        # Call the base method to get the targets
        targets = super().CreateTargets(algorithm, insights)
        
        # Adjust leverage for each target security
        for target in targets:
            if target.Quantity != 0:
                security = algorithm.Securities[target.Symbol]
                security.SetLeverage(2)
        
        return targets

# region imports
from AlgorithmImports import *
from DualMomentumAlphaModel import *
from MyPcm import *

from CustomImmediateExecutionModel import *
# endregion

class SectorDualMomentumStrategy(QCAlgorithm):
    undesired_symbols_from_previous_deployment = []
    checked_symbols_from_previous_deployment = False

    def initialize(self):
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2024, 7, 20)
        self.set_cash(100000)
        
        #self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.free_portfolio_value_percentage = 0.05

        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.asynchronous = True

        
        self.add_universe(self.universe.etf("SPY", self.universe_settings, self._etf_constituents_filter))





        self.add_alpha(DualMomentumAlphaModel())

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.day = -1
        self.set_portfolio_construction(RiskParityPortfolioConstructionModel(self._rebalance_func))    
        self.add_risk_management(TrailingStopRiskManagementModel())

        self.SetExecution(CustomImmediateExecutionModel(leverage=2.0))

        self.set_warm_up(timedelta(7))

       

    def _etf_constituents_filter(self, constituents: List[ETFConstituentUniverse]) -> List[Symbol]:
        selected = sorted([c for c in constituents if c.Weight],
                          key=lambda c: c.Weight, reverse=True)[:200]
        symbols = [c.Symbol for c in selected]



        return symbols

    def _rebalance_func(self, time):
        if self.day != self.time.day and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self.day = self.time.day
            return time
        return None



                
    # def on_data(self, data):
    #     # if not self.is_warming_up and not self.checked_symbols_from_previous_deployment:
    #     #     for symbol in self.undesired_symbols_from_previous_deployment:
    #     #         if self.is_market_open(symbol):
    #     #             self.liquidate(symbol, tag="Not backed up by current insights")
    #     #             self.undesired_symbols_from_previous_deployment.remove(symbol)

                
    #     #     for security_holding in self.portfolio.values():
    #     #         if not security_holding.invested:
    #     #             continue
    #     #         symbol = security_holding.symbol
    #     #         if not self.insights.has_active_insights(symbol, self.utc_time):
    #     #             self.undesired_symbols_from_previous_deployment.append(symbol)
    #     #     self.checked_symbols_from_previous_deployment = True

    #     # for symbol in self.undesired_symbols_from_previous_deployment:
    #     #     if self.is_market_open(symbol):
    #     #         self.liquidate(symbol, tag="Not backed up by current insights")
    #     #         self.undesired_symbols_from_previous_deployment.remove(symbol)