# region imports
from AlgorithmImports import *

from realized_gamma import RealizedGamma
from z_score import ZScore
# endregion


class FuturesIntradayTrendFollowingWithRealizedGammaAndLiquidityFilterAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2011, 1, 1)
        self.set_end_date(2024, 10, 1)
        self.set_cash(100_000_000)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        # Set some parameters.
        self._trading_interval_length = timedelta(minutes=60)
        self._realized_gamma_period = 20 # trading days (values in paper: 5, 20, 60, 120)
        self._liquidity_zscore_period = 252 # trading days (1 year for most Futures)
        self._weight_scaler = 5 # To utilize more cash.
        # Add the Futures.
        self._futures = []
        tickers = [
            Futures.Indices.SP_500_E_MINI,
            Futures.Indices.NASDAQ_100_E_MINI
        ]
        for ticker in tickers:
            future = self.add_future(
                ticker,
                fill_forward=False,
                data_normalization_mode=DataNormalizationMode.BACKWARDS_RATIO,
                data_mapping_mode=DataMappingMode.OPEN_INTEREST,
                contract_depth_offset=0
            )
            future.set_filter(lambda universe: universe.front_month())
            future.indicators_by_time = {}
            future.yesterdays_close = None
            future.previous_interval_close = None
            self._futures.append(future)
            # Create some Scheduled Events.
            date_rule = self.date_rules.every_day(future.symbol)
            self.schedule.on(date_rule, self.time_rules.midnight, lambda f=future: self._on_midnight(f))
            self.schedule.on(date_rule, self.time_rules.every(self._trading_interval_length), lambda f=future: self._rebalance(f))
            # Liquidate everything at the market close.
            self.schedule.on(
                date_rule, 
                # By default, you must place MOC orders at least 15.5 minutes before the close.
                self.time_rules.before_market_close(future.symbol, 16), 
                lambda f=future: self._close_position(f)
            )
        # Add a warm-up period to warm-up the indicators.
        self.set_warm_up(timedelta(int(1.5*max(self._realized_gamma_period, self._liquidity_zscore_period))))

    def _on_midnight(self, future):
        future.yesterdays_close = future.price
        future.daily_volume = 0

    def _rebalance(self, future):
        # Wait until the market is open.
        t = self.time
        if (not future.yesterdays_close or
            not future.exchange.hours.is_open(t - self._trading_interval_length, False)):
            return
        # Create indicators for this time interval if they don't already exist.
        trading_interval = (t.hour, t.minute)
        if trading_interval not in future.indicators_by_time:
            future.indicators_by_time[trading_interval] = {
                'RealizedGamma' : RealizedGamma(trading_interval, self._realized_gamma_period),
                'ZScore' : ZScore(trading_interval, self._liquidity_zscore_period)
            }
        indicators = future.indicators_by_time[trading_interval]
        
        # MANAGE INDICATOR 1: Realized Gamma
        #  Update the indicator.
        realized_gamma = indicators['RealizedGamma']
        return_since_last_close = future.price / future.yesterdays_close - 1
        if realized_gamma.update(IndicatorDataPoint(t, return_since_last_close)):
            self.plot(f'Realized Gamma ({future.symbol})', str(trading_interval), realized_gamma.value)
        #  Update the training data of the previous interval's realized Gamma indicator.
        if future.previous_interval_close:
            previous_t = t - self._trading_interval_length
            previous_time_period = (previous_t.hour, previous_t.minute)
            if previous_time_period in future.indicators_by_time:
                future.indicators_by_time[previous_time_period]['RealizedGamma'].add_label(
                    future.price / future.previous_interval_close - 1
                )
        #  Record the interval close price.
        future.previous_interval_close = future.price

        # MANAGE INDICATOR 2: Liquidity Z-Score
        #  Update the liquidity indicator.
        liquidity_z_score = indicators['ZScore']
        if (future.exchange.hours.is_open(t + self._trading_interval_length - timedelta(seconds=1), False) and
            future.daily_volume and 
            liquidity_z_score.update(IndicatorDataPoint(t, future.liquidity / future.daily_volume))):
            self.plot(f"Liquidity Z-Score ({future.symbol})", str(trading_interval), liquidity_z_score.value)
            self.plot(f'Liquidity ({future.symbol})', str(trading_interval), future.liquidity)
            self.plot(f"Daily Volume ({future.symbol})", str(trading_interval), future.daily_volume)
        else:
            return
        
        # Check if we can rebalance.
        if self.is_warming_up or not realized_gamma.ready or not liquidity_z_score.ready:
            return
        # Place trades to rebalance the portfolio.
        # Have exposure only when the realized gamma is negative (trending market) and liquidity is low.
        # Set the position proportional to the return since yesterday's close.
        self.set_holdings(
            future.mapped, 
            max(0, int(liquidity_z_score.value < 0) * int(realized_gamma.value < 0) * self._weight_scaler * return_since_last_close / len(self._futures))
        )

    def _close_position(self, future):
        quantity = self.portfolio[future.mapped].quantity
        if quantity:
            self.market_on_close_order(future.mapped, -quantity)

    def on_data(self, data):
        # Track volume and liquidity for the z-score indicator.
        for future in self._futures:
            if future.symbol in data.bars:
                future.daily_volume += future.volume
                future.liquidity = (future.bid_size + future.ask_size) / 2

# region imports
from AlgorithmImports import *

from sklearn.linear_model import LinearRegression
# endregion


class RealizedGamma(PythonIndicator):

    def __init__(self, trading_interval, period, fit_intercept=True):
        self.name = f'RealizedGamma({trading_interval}, {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._period = period
        self._model = LinearRegression(fit_intercept=fit_intercept)
    
    def update(self, input):
        # Check if there is sufficient training data.
        self.ready = len(self._y) == self._period
        if self.ready:
            # Fit model.
            self._model.fit(self._X.reshape(-1, 1), self._y.reshape(-1, 1))
            # Set the value to the opposite (negative) of 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._period:] 
        self.time = input.time
        return self.ready
    
    def add_label(self, label):
        self._y = np.append(self._y, label)[-self._period:]

# region imports
from AlgorithmImports import *
# endregion


class ZScore(PythonIndicator):

    def __init__(self, trading_interval, period):
        self.name = f'ZScore({trading_interval}, {period})'
        self.time = datetime.min
        self.value = 0
        self._mean = SimpleMovingAverage(period)
        self._std = StandardDeviation(period)

    def update(self, input):
        self._mean.update(input.time, input.value)
        self._std.update(input.time, input.value)
        self.ready = self._mean.is_ready and self._std.is_ready
        if self.ready:
            self.value = (input.value - self._mean.current.value) / self._std.current.value
            if self.value < 0:
                a = 1
        self.time = input.time
        return self.ready