'''
Random Forest classifier to predict the direction of SPY's price movement and trades options accordingly.
'''

from AlgorithmImports import *
# Import necessary libraries for machine learning
from sklearn.ensemble import RandomForestClassifier
import numpy as np

class SPYOptionMLTradingAlgorithm(QCAlgorithm):
    '''
    Initializes the algorithm with necessary parameters and settings.
    '''
    def Initialize(self):
        self.SetStartDate(2023, 1, 2)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100_000)

        self.spy = self.AddEquity("SPY", Resolution.Minute)
        self.lookback = 30
        self.training_data = []
        self.model = RandomForestClassifier(n_estimators=10)

        # Schedules the model training to occur at the start of each month.
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.after_market_open("SPY", 30), self.TrainModel)
        # Schedule trading logic to run before market close every day
        self.Schedule.On(self.date_rules.every_day("SPY"), self.time_rules.before_market_close("SPY", 10), self.TradeOptions)
        # Schedule liquidation of all holdings at 3:55 PM
        self.Schedule.On(self.date_rules.every_day(), self.time_rules.at(15, 55), lambda: self.Liquidate())

    '''
    Trains the Random Forest model using historical SPY price data.
    '''
    def TrainModel(self) -> None:
        history = self.History(self.spy.Symbol, self.lookback + 1, Resolution.Daily)
        if history.empty or history.isnull().values.any():
            return

        # Calculate features and labels for the model
        features = np.diff(history['close'])
        labels = np.where(features >= 0, 1, 0)[:-1]
        features = features.reshape(-1, 1)[:-1]

        # Fit the model to the training data
        self.model.fit(features, labels)

    '''
    Predicts the direction of SPY's price movement and trades options accordingly.
    '''
    def TradeOptions(self):
        # Get the current SPY price and calculate the feature for prediction
        today_features = np.array([self.spy.Price - self.spy.Close])

        # Make a prediction using the trained model
        prediction = self.model.predict(today_features.reshape(1, -1))

        # Buy a call option if the prediction is positive, otherwise buy a put option
        if prediction[0] == 1:
            option = self.BuyCallOption()
        else:
            option = self.BuyPutOption()

        # If an option contract is found, place a market order to buy it
        if option is not None:
            # Check if the option contract exists in the security list before placing an order
            if self.Securities.ContainsKey(option):
                self.MarketOrder(option, 1)
            else:
                self.Log("Option contract not found in security list: " + str(option))

    '''
    Returns a call option contract with the closest expiration date.
    '''
    def BuyCallOption(self):
        return self.GetOptionContract(True)

    '''
    Returns a put option contract with the closest expiration date.
    '''
    def BuyPutOption(self):
        return self.GetOptionContract(False)

    '''
    Returns an option contract based on the specified option type (call or put) and closest expiration date.
    '''
    def GetOptionContract(self, is_call):
        contracts = self.option_chain_provider.get_option_contract_list(self.spy.Symbol, self.Time)
        if is_call:
            contracts = [i for i in contracts if i.ID.OptionRight == OptionRight.Call]
        else:
            contracts = [i for i in contracts if i.ID.OptionRight == OptionRight.Put]

        # Sort contracts by expiration date and select the closest one
        contracts = sorted(contracts, key=lambda x: abs((x.ID.Date - self.Time).days))
        zero_dte_contracts = [i for i in contracts if (i.ID.Date - self.Time).days == 0]
        if zero_dte_contracts:
            return zero_dte_contracts[0]
        else:
            return None