# region imports
from AlgorithmImports import *
import numpy as np
from collections import deque
# endregion

# ======================================================================
# Curve Spread / Butterfly Mean-Reversion — OANDA bond CFDs on LEAN
# ----------------------------------------------------------------------
# Generalised, parameterised version. Same signal core (rolling OLS hedge
# ratio = empirical DV01, OU half-life -> z-window, z-score, regime gate),
# but:
#   * Targets OANDA CFDs (AddCfd / Market.Oanda), not CME futures.
#   * Tenor pair chosen by parameter over the full ladder 2/5/10/30.
#   * Optional THIRD leg turns the 2-leg curve spread into a butterfly.
#   * CFDs are continuous, so there is NO roll handling (unlike futures).
#
# Leg convention:
#   leg_a = the DEPENDENT leg (regressed on the others); leg_b (+ leg_c)
#   are the hedging legs. The residual of that regression is the spread.
#   - 2-leg curve spread:  leg_a/leg_b = the two tenors (e.g. 10 vs 30).
#   - butterfly (3 legs):  leg_a = belly, leg_b/leg_c = the two wings
#                          (e.g. leg_a=10, leg_b=5, leg_c=30 -> 5s10s30s fly).
#   Flip use_regime_filter 0/1 across two backtests for the filter on/off curves.
#
# READ BEFORE TRUSTING RESULTS:
#   1. CFD overnight financing IS modelled here via a daily scheduled charge
#      (LEAN's OANDA model itself uses a null interest-rate model). SET
#      financing_rate / financing_markup to the instrument's real OANDA rates;
#      the markup is a drag on gross notional, the rate acts on net notional.
#   2. Confirm the bond CFDs (USB02YUSD..USB30YUSD) actually have backtest
#      data depth in QC's OANDA CFD dataset for your window; the short tenors
#      may be thinner than the 10Y/30Y.
#   3. CFD prices are quotes; we read Securities[..].Price (mid).
#   4. Orders route as marketable LIMIT orders: a market order placed on a daily
#      bar (market closed) becomes a MarketOnOpen order, which OANDA rejects.
#      exec_buffer sets how far through the book they are priced.
#
# Not investment advice. Backtested results are not live results.
# ======================================================================

TENOR_TO_CFD = {"2": "USB02YUSD", "5": "USB05YUSD", "10": "USB10YUSD", "30": "USB30YUSD"}


class CurveSpreadOanda(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(100_000)
        self.SetBrokerageModel(BrokerageName.OandaBrokerage, AccountType.Margin)

        # ---- legs (parameterised over the ladder) ----
        a = self.GetParameter("leg_a") or "10"
        b = self.GetParameter("leg_b") or "30"
        c = self.GetParameter("leg_c") or ""          # "" -> 2-leg spread; set for a fly
        tenors = [a, b] + ([c] if c else [])
        self.legs = [self.AddCfd(TENOR_TO_CFD[t], Resolution.Daily, Market.Oanda) for t in tenors]
        self.Log(f"Trading legs (dependent first): {[TENOR_TO_CFD[t] for t in tenors]}")

        # ---- signal / risk parameters ----
        self.use_regime_filter = int(self.GetParameter("use_regime_filter", 1)) == 1
        self.entry_z   = float(self.GetParameter("entry_z", 2.0))
        self.exit_z    = float(self.GetParameter("exit_z", 0.5))
        self.stop_z    = float(self.GetParameter("stop_z", 3.8))
        self.reg_window= int(self.GetParameter("reg_window", 60))
        self.hl_mult   = float(self.GetParameter("hl_window_mult", 3.0))
        self.min_win   = int(self.GetParameter("min_window", 15))
        self.base_units= int(self.GetParameter("base_units", 10))   # units of the dependent leg
        self.exec_buf  = float(self.GetParameter("exec_buffer", 0.01))  # marketable-limit buffer
        self.trend_lb  = int(self.GetParameter("trend_lookback", 20))
        self.trend_th  = float(self.GetParameter("trend_strength", 1.0))

        # CFD overnight financing — SET to the instrument's real OANDA rates
        self.fin_rate   = float(self.GetParameter("financing_rate", 0.)) # 0.045))
        self.fin_markup = float(self.GetParameter("financing_markup", 0.)) # 0.010))
        self.last_fin = None

        maxlen = self.reg_window + 10
        self.px = [deque(maxlen=maxlen) for _ in self.legs]
        self.position = 0          # +1 long residual, -1 short, 0 flat
        self.SetWarmUp(maxlen, Resolution.Daily)

        # charge financing once per active day; the calendar-day delta picks up weekends
        self.Schedule.On(self.DateRules.EveryDay(self.legs[0].Symbol),
                         self.TimeRules.At(17, 0), self._apply_financing)

    def OnData(self, slice: Slice):
        syms = [lg.Symbol for lg in self.legs]
        if not all(slice.ContainsKey(s) for s in syms):
            return
        for i, s in enumerate(syms):
            self.px[i].append(float(self.Securities[s].Price))
        if self.IsWarmingUp or len(self.px[0]) < self.reg_window:
            return

        # 1) regress dependent leg on the hedging leg(s); residual = spread
        cols = [np.array(p, dtype=float)[-self.reg_window:] for p in self.px]
        coefs, _, resid = self._ols_multi(cols[0], cols[1:])
        if coefs is None:
            return

        # 2) OU half-life -> z-window
        hl = self._ou_half_life(resid)
        zwin = int(np.clip(self.hl_mult * hl, self.min_win, self.reg_window)) \
            if (hl is not None and np.isfinite(hl) and hl > 1) else self.reg_window
        seg = resid[-zwin:]
        sd = seg.std()
        if sd == 0:
            return
        z = (resid[-1] - seg.mean()) / sd

        # 3) regime gate — slope normalised by vol of DAILY CHANGES
        in_trend = False
        if self.use_regime_filter:
            slope = self._ema_slope(resid, self.trend_lb, self.trend_lb)
            dres = np.diff(resid)
            dvol = dres[-zwin:].std() if len(dres) >= 2 else sd
            denom = dvol * np.sqrt(self.trend_lb)
            in_trend = (abs(slope) / denom > self.trend_th) if denom > 0 else False

        # 4) state machine
        if self.position == 0:
            if not in_trend:
                if z > self.entry_z:
                    self._open(-1, coefs)        # residual rich -> short it
                elif z < -self.entry_z:
                    self._open(+1, coefs)        # residual cheap -> long it
        else:
            if abs(z) < self.exit_z or abs(z) > self.stop_z or (self.use_regime_filter and in_trend):
                self._close()

        self.Plot("Spread", "residual", float(resid[-1]))
        self.Plot("Signal", "z", float(z))

    # direction +1 = long residual (long dependent, short hedges); -1 = opposite.
    # Hedge units scale by the regression coefficients (empirical DV01 weights).
    def _open(self, direction, coefs):
        if any(self.Securities[lg.Symbol].Price == 0 for lg in self.legs):
            return
        self._limit(self.legs[0].Symbol, int(round(direction * self.base_units)))
        for j, co in enumerate(coefs):
            q = int(round(-direction * co * self.base_units))
            if q != 0:
                self._limit(self.legs[j + 1].Symbol, q)
        self.position = direction

    def _close(self):
        for lg in self.legs:
            h = self.Portfolio[lg.Symbol]
            if h.Invested:
                self._limit(lg.Symbol, -h.Quantity)
        self.position = 0

    # Marketable LIMIT: OANDA rejects MarketOnOpen (what a market order becomes
    # when the market is closed at a daily bar), so price through the book.
    def _limit(self, symbol, qty):
        if qty == 0:
            return
        px = float(self.Securities[symbol].Price)
        if px <= 0:
            return
        buf = 1.0 + (self.exec_buf if qty > 0 else -self.exec_buf)
        self.LimitOrder(symbol, qty, px * buf)

    def _apply_financing(self):
        # daily CFD financing: long pays (rate+markup), short receives (rate-markup),
        # both on notional => rate acts on NET notional, markup drags on GROSS notional.
        if self.last_fin is None:
            self.last_fin = self.Time
            return
        days = (self.Time - self.last_fin).total_seconds() / 86400.0
        self.last_fin = self.Time
        if days <= 0 or self.IsWarmingUp:
            return
        cash = 0.0
        for lg in self.legs:
            h = self.Portfolio[lg.Symbol]
            if not h.Invested:
                continue
            notional = float(h.HoldingsValue)        # signed, account currency
            cash += -(self.fin_rate * notional + self.fin_markup * abs(notional)) * days / 365.0
        if cash != 0.0:
            self.Portfolio.CashBook[self.AccountCurrency].AddAmount(cash)

    # ---- numerics ----
    @staticmethod
    def _ols_multi(y, regressors):
        """Regress y on the regressor columns + intercept. Returns (coefs, intercept, resid)."""
        n = len(y)
        A = np.column_stack(list(regressors) + [np.ones(n)])
        try:
            sol, *_ = np.linalg.lstsq(A, y, rcond=None)
        except Exception:
            return None, None, None
        return sol[:-1], sol[-1], y - A @ sol

    @staticmethod
    def _ou_half_life(s):
        if len(s) < 10:
            return None
        ds, lag = np.diff(s), s[:-1]
        A = np.vstack([lag, np.ones_like(lag)]).T
        try:
            theta = np.linalg.lstsq(A, ds, rcond=None)[0][0]
        except Exception:
            return None
        if theta >= 0:
            return None
        phi = 1.0 + theta
        return -np.log(2) / np.log(phi) if phi > 0 else None

    @staticmethod
    def _ema_slope(s, span, lookback):
        a = 2.0 / (span + 1.0)
        e = s[0]
        out = np.empty_like(s)
        for i, v in enumerate(s):
            e = a * v + (1 - a) * e
            out[i] = e
        return 0.0 if len(out) <= lookback else out[-1] - out[-1 - lookback]