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

# ------------- NEW: 5-percent-per-signal portfolio model -------------
class FixedFractionPortfolioConstructionModel(PortfolioConstructionModel):
    """Always allocate fixed fraction of portfolio value per insight."""
    def __init__(self, fraction: float = 0.05):
        super().__init__()
        self.fraction = fraction

    def CreateTargets(self, algorithm, insights):
        targets = []
        for ins in insights:
            if ins.Direction == InsightDirection.Flat:
                continue
            direction = 1 if ins.Direction == InsightDirection.Up else -1
            weight = direction * self.fraction
            targets.append(PortfolioTarget.Percent(algorithm, ins.Symbol, weight))
        return targets
# --------------------------------------------------------------------

class MatureBondAndIndexETFsWithStops(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2024, 2, 12)
        self.SetCash(100_000)
        self.SetTimeZone(TimeZones.NewYork)

        # ---------------- Universe ------------------------------------
        self.tickers = [
            # Bond / fixed-income ETFs
            "SHY","IEI","IEF","TLT","TIP","AGG","BND","LQD","HYG","IBND",
            # Major equity-index ETFs
            "SPY","QQQ","DIA","IWM","VTI"
        ]

        self.SL_mult = 0.5
        self.TP_mult = 2.0
        self.insightAtr = {}

        for tkr in self.tickers:
            self.AddEquity(tkr, Resolution.Daily)
        self.anchor = "SPY"

        # --------- << use the new 5 % model here >> -------------------
        self.SetPortfolioConstruction(FixedFractionPortfolioConstructionModel(0.05))
        self.SetExecution(ImmediateExecutionModel())
        # --------------------------------------------------------------

        self.Schedule.On(
            self.DateRules.EveryDay(self.anchor),
            self.TimeRules.AfterMarketClose(self.anchor, 1),
            self.GenerateInsights
        )

    # ================================================================
    def GenerateInsights(self):
        hist = self.History(self.tickers, 16, Resolution.Daily)
        if hist.empty:
            return

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

            # two up-days momentum signal
            r1 = df["close"].pct_change().iloc[-1]
            r2 = df["close"].pct_change().iloc[-2]
            if not (r1 > 0 and r2 > 0):
                continue

            prev_close = df["close"].shift(1)
            tr = pd.concat([
                    df["high"] - df["low"],
                    (df["high"] - prev_close).abs(),
                    (df["low"]  - prev_close).abs()
                 ], axis=1).max(axis=1)
            atr = tr.rolling(14).mean().shift(1).iloc[-1]
            if pd.isna(atr):
                continue

            self.insightAtr[sym] = atr
            insights.append(
                Insight.Price(sym, timedelta(days=1), InsightDirection.Up)
            )
            self.Log(f"{sym} insight @ {self.Time}  ATR14={atr:.4f}")

        if insights:
            self.EmitInsights(insights)

    # ================================================================
    def OnOrderEvent(self, oe):
        if oe.Status != OrderStatus.Filled or oe.FillQuantity <= 0:
            return
        sym = oe.Symbol
        atr = self.insightAtr.get(sym)
        if atr is None:
            return

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

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

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