from AlgorithmImports import *
import joblib
import io
import numpy as np
from collections import deque
from sklearn.preprocessing import StandardScaler
from datetime import datetime

class HMMNQRegimeStrategy(QCAlgorithm):
    def Initialize(self):
        # === SET BACKTEST PERIOD ===
        self.SetStartDate(2025, 3, 10)   # OUT-OF-SAMPLE (Week after training)
        self.SetEndDate(2025, 3, 14)
        self.SetCash(100000)

        # === TRAINING WINDOW for Model ID (must match QuantBook training period) ===
        train_start = datetime(2025, 3, 3)
        train_end   = datetime(2025, 3, 7)
        date_format = "%Y%m%d"
        self.model_id = f"hmm_qsl_nq_{train_start.strftime(date_format)}to{train_end.strftime(date_format)}_marco.pkl"

        # === SUBSCRIBE TO CONTINUOUS NQ FUTURES ===
        future = self.AddFuture(
            Futures.Indices.NASDAQ100EMini,
            Resolution.Minute,
            extendedMarketHours=True,
            dataMappingMode=DataMappingMode.OpenInterest,
            dataNormalizationMode=DataNormalizationMode.BackwardsRatio,
            contractDepthOffset=0
        )
        future.SetFilter(timedelta(0), timedelta(days=60))
        self.future_symbol = future.Symbol
        self.current_contract = None
        self.previous_contract = None

        # === LOAD HMM MODEL + SCALER ===
        model_bytes = self.ObjectStore.ReadBytes(self.model_id)
        self.model, self.scaler = joblib.load(io.BytesIO(model_bytes))
        self.Debug(f"✅ Loaded model: {self.model_id}")

        # === ROLLING WINDOWS FOR FEATURE ENGINEERING ===
        self.window = deque(maxlen=31)
        self.close_window = deque(maxlen=11)
        self.high_window = deque(maxlen=1)
        self.low_window = deque(maxlen=1)

        # Track predictions for diagnostics
        self.regime_counts = {0: 0, 1: 0, 2: 0}
        self.SetWarmUp(TimeSpan.FromMinutes(31))

    def OnData(self, slice: Slice):
        if self.IsWarmingUp:
            return

        # === HANDLE ROLLOVERS ===
        for evt in slice.SymbolChangedEvents.Values:
            if evt.Symbol == self.future_symbol:
                self.previous_contract = evt.OldSymbol
                self.current_contract = evt.NewSymbol
                self.Debug(f"{self.Time} - Rollover: {self.previous_contract} → {self.current_contract}")
                self.HandleRollover()

        # === FIRST TIME CONTRACT SELECTION ===
        if self.current_contract is None:
            if self.future_symbol not in slice.FutureChains:
                return
            chain = slice.FutureChains[self.future_symbol]
            contracts = sorted([c for c in chain if c.Expiry > self.Time], key=lambda c: c.Expiry)
            if contracts:
                self.current_contract = contracts[0].Symbol
                self.Debug(f"{self.Time} - Selected front-month: {self.current_contract}")
            else:
                self.Debug(f"{self.Time} - No valid contracts in chain")
            return

        # === GET BAR DATA ===
        if not slice.Bars.ContainsKey(self.current_contract):
            self.Debug(f"{self.Time} - No bar data for {self.current_contract}")
            return

        bar = slice.Bars[self.current_contract]
        close = bar.Close
        high = bar.High
        low = bar.Low

        # === UPDATE ROLLING WINDOWS ===
        self.window.append(close)
        self.close_window.append(close)
        self.high_window.append(high)
        self.low_window.append(low)

        if len(self.window) < 31:
            return

        # === FEATURE ENGINEERING ===
        log_ret = np.log(self.window[-1] / self.window[-2])
        log_returns = np.log(np.array(self.window)[1:] / np.array(self.window)[:-1])
        vol = np.std(log_returns)
        slope = (self.close_window[-1] - self.close_window[0]) / 10
        range_pct = (self.high_window[-1] - self.low_window[-1]) / self.close_window[-1]

        x = np.array([[log_ret, vol, slope, range_pct]])
        x_scaled = self.scaler.transform(x)
        regime = self.model.predict(x_scaled)[0]
        self.regime_counts[regime] += 1

        self.Debug(f"{self.Time} - Regime: {regime}")

        # === STRATEGY LOGIC: TRADE IN REGIME 0 ONLY ===
        if regime == 0:
            if not self.Portfolio[self.current_contract].Invested:
                self.MarketOrder(self.current_contract, 1)
                self.Debug(f"{self.Time} - Entered LONG (Regime 0)")
        else:
            if self.Portfolio[self.current_contract].Invested:
                self.Liquidate(self.current_contract)
                self.Debug(f"{self.Time} - Exited position (Regime {regime})")

    def HandleRollover(self):
        if self.Portfolio[self.previous_contract].Invested:
            qty = self.Portfolio[self.previous_contract].Quantity
            self.Liquidate(self.previous_contract)
            self.MarketOrder(self.current_contract, qty)
            self.Debug(f"{self.Time} - Rolled over: {qty} → {self.current_contract}")