# region imports
from AlgorithmImports import *
from datetime import datetime, timedelta
from collections import deque
import math
# endregion

class ETFConstituentUniverseROCPAlphaModelAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2013, 1, 7)
        self.SetCash(100000)

        self.SetAlpha(ConstituentWeightedROCPAlphaModel())
        self.SetPortfolioConstruction(PercentagePortfolioConstructionModel(self, 0.04))  # 4% per trade
        self.SetExecution(ImmediateExecutionModel())

        iwb = self.AddEquity("IWB", Resolution.Daily).Symbol  # IWB is the iShares Russell 1000 ETF
        self.UniverseSettings.Resolution = Resolution.Daily
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.01
        self.AddUniverse(self.Universe.ETF(iwb, self.UniverseSettings, self.FilterETFConstituents))

    def FilterETFConstituents(self, constituents):
        return [i.Symbol for i in constituents if i.Weight is not None and i.Weight >= 0.001]

class ConstituentWeightedROCPAlphaModel(AlphaModel):
    def __init__(self):
        self.rocp_symbol_data = {}
        self.entry_signals = {}
        self.trade_start_time = {}

    def Update(self, algorithm: QCAlgorithm, data: Slice):
        insights = []

        for symbol, symbol_data in self.rocp_symbol_data.items():
            if symbol not in data.Bars:
                continue

            bar = data.Bars[symbol]
            symbol_data.Update(bar)

            if not symbol_data.rocp.IsReady or not symbol_data.atr.IsReady:
                continue

            current_roc = symbol_data.rocp.Current.Value
            yesterday_roc = symbol_data.previous_roc

            if symbol not in self.entry_signals:
                if current_roc < -20 and current_roc > yesterday_roc:
                    self.entry_signals[symbol] = True
                    algorithm.Debug(f"Entry signal for {symbol}: ROC = {current_roc:.2f}, Yesterday ROC = {yesterday_roc:.2f}")
            elif self.entry_signals[symbol]:
                entry_price = bar.Open
                profit_target = entry_price + symbol_data.atr.Current.Value
                stop_loss = entry_price - 2.5 * symbol_data.atr.Current.Value

                insights.append(Insight.Price(
                    symbol,
                    timedelta(days=20),
                    InsightDirection.Up,
                    abs(current_roc),
                    None
                ))
                self.entry_signals[symbol] = False
                self.trade_start_time[symbol] = algorithm.Time
                algorithm.Debug(f"Entering trade for {symbol} at {entry_price:.2f}, PT: {profit_target:.2f}, SL: {stop_loss:.2f}")
            elif symbol in self.trade_start_time:
                if (algorithm.Time - self.trade_start_time[symbol]).days >= 20:
                    insights.append(Insight.Price(
                        symbol,
                        timedelta(days=1),
                        InsightDirection.Flat,
                        0,
                        None
                    ))
                    del self.trade_start_time[symbol]
                    algorithm.Debug(f"Exiting trade for {symbol} due to 20-day hold period")

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            if added.Symbol not in self.rocp_symbol_data:
                self.rocp_symbol_data[added.Symbol] = SymbolData(added.Symbol, algorithm, 14)
        
        for removed in changes.RemovedSecurities:
            if removed.Symbol in self.rocp_symbol_data:
                del self.rocp_symbol_data[removed.Symbol]
            if removed.Symbol in self.entry_signals:
                del self.entry_signals[removed.Symbol]
            if removed.Symbol in self.trade_start_time:
                del self.trade_start_time[removed.Symbol]

class SymbolData:
    def __init__(self, symbol, algorithm, period):
        self.symbol = symbol
        self.rocp = algorithm.ROCP(symbol, period)
        self.atr = algorithm.ATR(symbol, period)
        self.previous_values = deque(maxlen=2)
        self.previous_roc = 0

    def Update(self, bar):
        self.rocp.Update(bar.EndTime, bar.Close)
        self.atr.Update(bar)
        self.previous_values.append(self.rocp.Current.Value)
        if len(self.previous_values) == 2:
            self.previous_roc = self.previous_values[0]

class PercentagePortfolioConstructionModel(PortfolioConstructionModel):
    def __init__(self, algorithm, percentage_per_trade):
        self.algorithm = algorithm
        self.percentage_per_trade = percentage_per_trade
        self.targets = {}

    def CreateTargets(self, algorithm, insights):
        targets = []

        for insight in insights:
            if insight.Direction == InsightDirection.Up:
                self.targets[insight.Symbol] = self.percentage_per_trade
            elif insight.Direction == InsightDirection.Flat:
                if insight.Symbol in self.targets:
                    del self.targets[insight.Symbol]

        portfolio_value = algorithm.Portfolio.TotalPortfolioValue
        for symbol, target_percentage in self.targets.items():
            target_value = portfolio_value * target_percentage
            quantity = math.floor(target_value / algorithm.Securities[symbol].Price)
            targets.append(PortfolioTarget(symbol, quantity))

        return targets