# region imports
from AlgorithmImports import *

from sklearn.linear_model import LinearRegression
# endregion

class FuturesIntradayTrendFollowingWithRealizedGammaAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2011, 1, 1)
        self.set_end_date(2024, 10, 1)
        self._future = self.add_future(
            Futures.Indices.SP_500_E_MINI,
            data_normalization_mode=DataNormalizationMode.BACKWARDS_RATIO,
            data_mapping_mode=DataMappingMode.OPEN_INTEREST,
            contract_depth_offset=0
        )
        self._future.set_filter(lambda universe: universe.front_month())
        self._future.realized_gamma_by_time = {}
        self._future.yesterdays_close = None
        self._future.previous_interval_close = None
        self._trading_interval = timedelta(minutes=60)
        self._lookback = 120 # trading days
        date_rule = self.date_rules.every_day(self._future.symbol)
        self.schedule.on(date_rule, self.time_rules.midnight, self._record_close_price)
        self.schedule.on(date_rule, self.time_rules.every(self._trading_interval), self._rebalance)
        self.set_warm_up(timedelta(int(1.5*self._lookback)))

    def _record_close_price(self):
        self._future.yesterdays_close = self._future.price

    def _rebalance(self):
        # Wait until the market is open.
        t = self.time
        if (not self._future.yesterdays_close or
            not self._future.exchange.hours.is_open(t - self._trading_interval, False)):
            return
        # Create a realized Gamma indicator for this time interval if it doesn't already exist.
        time_period = (t.hour, t.minute)
        if time_period not in self._future.realized_gamma_by_time:
            self._future.realized_gamma_by_time[time_period] = RealizedGamma(time_period, self._lookback)
        # Get the indicator value for this time interval.
        realized_gamma = self._future.realized_gamma_by_time[time_period]
        return_since_last_close = self._future.price / self._future.yesterdays_close - 1
        if realized_gamma.update(IndicatorDataPoint(t, return_since_last_close)):
            self.plot('Realized Gamma', str(time_period), realized_gamma.value)
        # Update the training data of the previous interval's indicator.
        if self._future.previous_interval_close:
            previous_t = t - self._trading_interval
            previous_time_period = (previous_t.hour, previous_t.minute)
            if previous_time_period in self._future.realized_gamma_by_time:
                self._future.realized_gamma_by_time[previous_time_period].add_label(
                    self._future.price / self._future.previous_interval_close - 1
                )
        # Record the interval close price.
        self._future.previous_interval_close = self._future.price
        # Check if we can rebalance.
        if not self._future.exchange.hours.is_open(t + self._trading_interval - timedelta(seconds=1), False):
            return
        # Rebalance


class RealizedGamma(PythonIndicator):

    def __init__(self, time_period, lookback):
        self.name = f'RealizedGamma{time_period}'
        self.time = datetime.min
        self.value = 0
        self._X = np.array([]) # Return from previous close to t.
        self._y = np.array([]) # Return from t to t+trading_interval.
        self._lookback = lookback
        self._model = LinearRegression(fit_intercept=True)
    
    def update(self, input):
        # Check if there is sufficient training data.
        is_ready = len(self._y) == self._lookback
        if is_ready:
            # Fit model.
            self._model.fit(self._X.reshape(-1, 1), self._y.reshape(-1, 1))
            # Set the value to the predicted the return from t to t+trading_interval.
            # `input.value` is the return from previous close to t.
            self.value = self._model.predict([[input.value]])[0][0]
        # Add the sample of the independent variable to the training data. 
        self._X = np.append(self._X, input.value)[-self._lookback:] 
        self.time = input.time
        return is_ready
    
    def add_label(self, label):
        self._y = np.append(self._y, label)[-self._lookback:]