# region imports
from AlgorithmImports import *
from scipy.optimize import minimize
from hmmlearn.hmm import GMMHMM
# endregion

class DrawdownRegimeGoldHedgeAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_end_date(2025, 1, 1)
        self.set_start_date(self.end_date - timedelta(6*365))
        self.set_cash(1000000)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        # Determine the lookback window (in weeks).
        self.history_lookback = self.get_parameter("history_lookback", 50)
        self.drawdown_lookback = self.get_parameter("drawdown_lookback", 50)

        # Request SPY as market representative for trading and HMM fitting.
        self.spy = self.add_equity("SPY", Resolution.MINUTE).symbol
        # Request GLD as hedge asset for trading.
        self.gold = self.add_equity("GLD", Resolution.MINUTE).symbol
        self.set_benchmark(self.spy)

        # Schdeuled a weekly rebalance.
        self.schedule.on(self.date_rules.week_start(self.spy), self.time_rules.after_market_open(self.spy, 1), self.rebalance)
    
    def rebalance(self) -> None:
        # Get the drawdown as the input to the drawdown regime. Since we're rebalancing weekly, we resample to study weekly drawdown.
        history = self.history(self.spy, self.history_lookback*5, Resolution.DAILY).unstack(0).close.resample('W').last()
        drawdown = history.rolling(self.drawdown_lookback).apply(lambda a: (a.iloc[-1] - a.max()) / a.max()).dropna()
        try:
            # Initialize the HMM, then fit by the drawdown data, as we're interested in the downside risk regime.
            # McLachlan & Peel (2000) suggested 2-3 components are used in GMMs to capture the main distribution and the tail to balance between complexity and characteristics capture.
            # By studying the ACF and PACF plots, the 1-lag drawdown series is suitable to supplement as exogenous variable.
            inputs = np.concatenate([drawdown[[self.spy]].iloc[1:].values, drawdown[[self.spy]].diff().iloc[1:].values], axis=1)
            model = GMMHMM(n_components=2, n_mix=3, covariance_type='tied', n_iter=100, random_state=10).fit(inputs)
            # Obtain the current market regime.
            regime_probs = model.predict_proba(inputs)
            current_regime_prob = regime_probs[-1]
            regime = 0 if current_regime_prob[0] > current_regime_prob[1] else 1

            # Determine the regime number: the higher the coefficient, the larger the drawdown in this state.
            high_regime = 1 if model.means_[0][1][0] < model.means_[1][1][0] else 0
            # Check the transitional probability of the next regime being the high volatility regime.
            # Calculated by the probability of the current regime being 1/0, then multiplied by the posterior probabilities of each scenario.
            next_prob_zero = current_regime_prob @ model.transmat_[:, 0]
            next_prob_high = next_prob_zero if high_regime == 0 else 1 - next_prob_zero

            # Buy more Gold and less SPY if the current regime is easier to have large drawdown.
            # Fund will shift to hedge asset like gold to drive up its price.
            # Weighted by the posterior probabilities.
            self.set_holdings([PortfolioTarget(self.gold, next_prob_high),
                                PortfolioTarget(self.spy, 1 - next_prob_high)],
                                liquidate_existing_holdings=True)
        except:
            pass