# region imports
from AlgorithmImports import *
# endregion
from QuantConnect import Resolution, TimeZones
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Algorithm.Framework.Portfolio import EqualWeightingPortfolioConstructionModel
from QuantConnect.Algorithm.Framework.Execution import ImmediateExecutionModel
from QuantConnect.Algorithm.Framework.Alphas import Insight, InsightDirection
from QuantConnect.Orders import OrderStatus
from datetime import timedelta

class MatureSmallCapsWithStopsStrategy(QCAlgorithm):

    def Initialize(self):
        # Backtest window & cash
        self.SetStartDate(2022,1,1)
        self.SetEndDate(2024,2,12)
        self.SetCash(100000)
        self.SetTimeZone(TimeZones.NewYork)

        # Mature small-cap universe
        self.tickers = [
                # Technology
                "AAPL", "MSFT", "NVDA",
                # Consumer Discretionary
                "AMZN", "TSLA", "NKE",
                # Consumer Staples
                "KO", "PEP", "PG",
                # Healthcare
                "UNH", "JNJ", "PFE",
                # Financials
                "JPM", "V", "MA",
                # Energy
                "XOM", "CVX", "COP",
                # Industrials
                "HON", "BA", "CAT",
                # Utilities
                "NEE", "DUK", "SO",
                # Materials
                "LIN", "APD", "SHW"
                ]
        # Strategy params
        self.SL_mult = 0.5
        self.TP_mult = 2.0

        # Will hold ATR per symbol at insight time
        self.insightAtr = {}

        # 1) Subscriptions
        for sym in self.tickers:
            self.AddEquity(sym, Resolution.Daily)
        self.AddEquity("SPY", Resolution.Daily)  # scheduler anchor

        # 2) Framework Setup
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())

        # 3) Schedule daily insight generation 1′ after SPY close (~16:01 ET)
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.AfterMarketClose("SPY", 1),
            self.GenerateInsights
        )

    def GenerateInsights(self):
        """Compute 2-day up signal + lagged ATR, emit Up Insights and store ATR."""
        history = self.History(self.tickers, 16, Resolution.Daily)
        ts      = self.Time
        insights = []

        for sym in self.tickers:
            if sym not in history.index.get_level_values(0): continue
            df = history.loc[sym].copy()
            if df.shape[0] < 16: continue

            df["ret"]    = df["close"].pct_change()
            signal      = (df["ret"].shift(1)>0) & (df["ret"].shift(2)>0)
            # True Range + lagged ATR14
            df["prev_close"] = df["close"].shift(1)
            df["tr"] = df[["high","low","prev_close"]].apply(
                lambda row: max(row["high"]-row["low"],
                                abs(row["high"]-row["prev_close"]),
                                abs(row["low"] -row["prev_close"])),
                axis=1
            )
            df["atr"] = df["tr"].rolling(14).mean().shift(1)

            if signal.iloc[-1] and not df["atr"].iloc[-1] != df["atr"].iloc[-1]:
                atr = df["atr"].iloc[-1]
                self.insightAtr[sym] = atr
                insights.append(Insight.Price(sym, timedelta(days=1), InsightDirection.Up))
                self.Log(f"{sym} INSIGHT @ {ts:yyyy-MM-dd HH:mm} ET  ATR={atr:.2f}")

        if insights:
            self.EmitInsights(insights)

    def OnOrderEvent(self, orderEvent):
        """Place ATR-based SL/TP immediately after our market entry fills."""
        if orderEvent.Status != OrderStatus.Filled:
            return

        symbol = orderEvent.Symbol
        # only act on our entry orders (Insight-triggered market orders)
        if orderEvent.FillQuantity <= 0:
            return

        # get entry price & ATR
        price = orderEvent.FillPrice
        atr   = self.insightAtr.get(symbol)
        if atr is None:
            return

        qty = orderEvent.FillQuantity
        sl  = price - self.SL_mult * atr
        tp  = price + self.TP_mult * atr

        self.StopMarketOrder(symbol, -qty, sl)
        self.LimitOrder    (symbol, -qty, tp)
        self.Log(f"{symbol} SL@{sl:.2f}  TP@{tp:.2f}  qty={qty}")

    def OnEndOfAlgorithm(self):
        self.Log(f"Final Portfolio Value: ${self.Portfolio.TotalPortfolioValue:0.2f}")