# 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)
        self._weight_scaler = 5 # To utilize more cash.
        # Add the E-mini.
        self._future = self.add_future(
            Futures.Indices.SP_500_E_MINI,
            fill_forward=False,
            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.indicators_by_time = {}
        self._future.yesterdays_close = None
        self._future.previous_interval_close = None
        # Create some Scheduled Events.
        date_rule = self.date_rules.every_day(self._future.symbol)
        self.schedule.on(date_rule, self.time_rules.midnight, self._on_midnight)
        self.schedule.on(date_rule, self.time_rules.every(self._trading_interval_length), self._rebalance)
        # 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(self._future.symbol, 16), 
            self._close_position
        )
        # 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):
        self._future.yesterdays_close = self._future.price
        self._future.daily_volume = 0

    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_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 self._future.indicators_by_time:
            self._future.indicators_by_time[trading_interval] = {
                'RealizedGamma' : RealizedGamma(trading_interval, self._realized_gamma_period),
                'ZScore' : ZScore(trading_interval, self._liquidity_zscore_period)
            }
        indicators = self._future.indicators_by_time[trading_interval]
        
        # MANAGE INDICATOR 1: Realized Gamma
        #  Update the indicator.
        realized_gamma = indicators['RealizedGamma']
        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(trading_interval), realized_gamma.value)
        #  Update the training data of the previous interval's realized Gamma indicator.
        if self._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 self._future.indicators_by_time:
                self._future.indicators_by_time[previous_time_period]['RealizedGamma'].add_label(
                    self._future.price / self._future.previous_interval_close - 1
                )
        #  Record the interval close price.
        self._future.previous_interval_close = self._future.price

        # MANAGE INDICATOR 2: Liquidity Z-Score
        #  Update the liquidity indicator.
        liquidity_z_score = indicators['ZScore']
        if (self._future.exchange.hours.is_open(t + self._trading_interval_length - timedelta(seconds=1), False) and
            self._future.daily_volume and 
            liquidity_z_score.update(IndicatorDataPoint(t, self._future.liquidity / self._future.daily_volume))):
            self.plot("Liquidity Z-Score", str(trading_interval), liquidity_z_score.value)
            self.plot('Liquidity', str(trading_interval), self._future.liquidity)
            self.plot("Daily Volume", str(trading_interval), self._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(
            self._future.mapped, 
            int(liquidity_z_score.value < 0) * int(realized_gamma.value < 0) * self._weight_scaler * return_since_last_close
        )

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

    def on_data(self, data):
        # Track volume and liquidity for the z-score indicator.
        if self._future.symbol in data.bars:
            self._future.daily_volume += self._future.volume
            self._future.liquidity = (self._future.bid_size + self._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