from AlgorithmImports import *
import joblib
import os
import numpy as np
import io # Import the io module

class MLBasedTradingStrategy(QCAlgorithm):
    def Initialize(self):
        # Set a reasonable backtest period for daily resolution and indicator warm-up
        self.SetStartDate(2022, 1, 1)
        # ✅ FIXED: Extend the end date to ensure enough time for warm-up and trading
        self.SetEndDate(2024, 12, 31) # Example: Backtest until end of 2024
        # Or, to run until the latest available data for a backtest:
        # self.SetEndDate(self.Time.date().replace(year=self.Time.date().year - 1)) # Example: Up to last year's end
        # Or simply remove SetEndDate() for backtests to run until current date based on data availability,
        # but be mindful if it hits "current date" before warm-up completes.
        # For reliable backtest duration, a fixed future date is often best.

        self.SetCash(100000)

        # Add the equity subscription for SPY
        self.symbol = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # Initialize model variables
        self.model = None
        self.model_loaded = False
        self.last_prediction = None # To store the last prediction for debugging or later use

        # === Indicators ===
        self.sma50 = self.SMA(self.symbol, 50)
        self.sma200 = self.SMA(self.symbol, 200)
        self.ema50 = self.EMA(self.symbol, 50)
        self.macd = self.MACD(self.symbol, 12, 26, 9, MovingAverageType.Exponential, Resolution.Daily)
        self.rsi14 = self.RSI(self.symbol, 14)
        self.stoch = self.STO(self.symbol, 14, 3, 3) 
        self.bb = self.BB(self.symbol, 20, 2) 
        self.atr = self.ATR(self.symbol, 14)
        self.obv = self.OBV(self.symbol, Resolution.Daily)
        self.volume_sma20 = self.SMA(self.symbol, 20, Resolution.Daily, Field.Volume)

        # Initialize RollingWindow for CMF calculation
        self.cmf_window = RollingWindow[TradeBar](20)

        # Set a warm-up period to ensure all indicators have enough data before the strategy starts trading
        self.SetWarmUp(self.sma200.Period + 5) # Warm-up for 205 daily bars

        # --- Model Loading ---
        model_name = "rf_trading_model_SPY.pkl" 
        
        if self.ObjectStore.ContainsKey(model_name):
            try:
                model_bytes = self.ObjectStore.ReadBytes(model_name)
                self.model = joblib.load(io.BytesIO(model_bytes)) 
                self.model_loaded = True
                self.Debug(f"Model '{model_name}' loaded successfully from ObjectStore.")
            except Exception as e:
                self.Error(f"Error loading model '{model_name}' from ObjectStore: {e}")
        else:
            self.Error(f"Model file '{model_name}' not found in ObjectStore. Please ensure it's uploaded with this exact name.")

        # Schedule the prediction and trading logic
        # This will run daily 5 minutes after market open, but only AFTER warm-up is finished.
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.symbol, 5), self.PredictAndTradeIfReady)

        self.Debug("Initialization complete.")

    def OnData(self, data):
        # This method is called whenever new data arrives.
        # Add the current TradeBar to the RollingWindow for CMF calculation.
        if data.Bars.ContainsKey(self.symbol):
            self.cmf_window.Add(data.Bars[self.symbol])

    def PredictAndTradeIfReady(self):
        # Added a check to ensure warm-up is complete
        if self.IsWarmingUp:
            return

        self.PredictAndTrade() # Call the actual prediction and trade logic

    def PredictAndTrade(self):
        # 1. Check if the model is loaded
        if not self.model_loaded or self.model is None:
            self.Debug("Model not loaded yet or failed to load. Skipping prediction and trade.")
            return

        # 2. Check if all required indicators and rolling window are ready
        indicators_ready = all([
            self.sma50.IsReady,
            self.sma200.IsReady,
            self.ema50.IsReady,
            self.macd.IsReady,
            self.rsi14.IsReady,
            self.stoch.IsReady,
            self.bb.IsReady,
            self.atr.IsReady,
            self.obv.IsReady,
            self.volume_sma20.IsReady,
            self.cmf_window.IsReady 
        ])

        if not indicators_ready:
            self.Debug("Indicators or CMF window not ready. Skipping prediction.")
            return

        # 3. Get current security data
        security = self.Securities[self.symbol]
        
        if not security.HasData:
            self.Debug("No current data for SPY. Skipping prediction.")
            return

        price = security.Price
        volume = security.Volume

        # 4. Calculate Chaikin Money Flow (CMF)
        cmf = self.CalculateCMF()

        # 5. Prepare features for the model
        features = [
            security.Open,
            security.High,
            security.Low,
            price,
            volume,
            self.sma50.Current.Value,
            self.sma200.Current.Value,
            self.ema50.Current.Value,
            self.macd.Current.Value,
            self.rsi14.Current.Value,
            self.stoch.StochK.Current.Value,
            self.stoch.StochD.Current.Value,
            self.bb.UpperBand.Current.Value,
            self.bb.LowerBand.Current.Value,
            self.atr.Current.Value,
            self.obv.Current.Value,
            cmf,
            0.0  # Market sentiment (Placeholder)
        ]

        features_array = np.array(features).reshape(1, -1)

        try:
            # 6. Make prediction
            prediction = self.model.predict(features_array)[0]
            self.last_prediction = prediction 
            self.Debug(f"Prediction: {prediction}")

            # 7. Execute trades based on prediction
            if prediction == 1 and not self.Portfolio[self.symbol].Invested:
                self.SetHoldings(self.symbol, 1) 
                self.Debug("BUY signal executed: Going long.")
            elif prediction == -1 and self.Portfolio[self.symbol].Invested:
                self.Liquidate(self.symbol) 
                self.Debug("SELL signal executed: Liquidating position.")
            elif prediction == 0:
                self.Debug("HOLD signal: No action taken.")
            
        except Exception as e:
            self.Error(f"Error during prediction or trading logic: {e}")

    def CalculateCMF(self):
        if not self.cmf_window.IsReady:
            return 0.0 

        money_flow_volume = 0.0
        total_volume = 0.0

        for bar in self.cmf_window:
            typical_price = (bar.High + bar.Low + bar.Close) / 3
            
            if (bar.High - bar.Low) != 0: 
                mfm = ((bar.Close - bar.Low) - (bar.High - bar.Close)) / (bar.High - bar.Low)
            else:
                mfm = 0.0
            
            mfv = mfm * bar.Volume
            
            money_flow_volume += mfv
            total_volume += bar.Volume

        return money_flow_volume / total_volume if total_volume > 0 else 0.0