#region imports
from AlgorithmImports import *
#endregion

class SparseOptimizationIndexTrackingDemo(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(2017, 1, 1)
        self.set_start_date(2022, 1, 1)
        self.set_brokerage_model(BrokerageName.ALPHA_STREAMS)
        self.set_cash(1000000)
        
        # Add our ETF constituents of the index that we would like to track.
        self.QQQ = self.add_equity("QQQ", Resolution.MINUTE).symbol
        self.universe_settings.resolution = Resolution.MINUTE
        self.add_universe(self.universe.etf(self.QQQ, self.universe_settings, self.etf_selection))
        
        self.set_benchmark("QQQ")
        
        # Set up varaibles to flag the time to recalibrate and hold the constituents.
        self.last_time = datetime.min
        self.assets = []
        
    def etf_selection(self, constituents):
        # We want all constituents to be considered.
        self.assets = [x.symbol for x in constituents]
        return self.assets
            
    def on_data(self, data):
        qb = self
        if self.last_time > self.time:
            return
        
        # Prepare the historical return data of the constituents and the ETF (as index to track).
        history = qb.history(self.assets, 252, Resolution.DAILY)
        if history.empty: return
    
        history_portfolio = history.close.unstack(0)
        pct_change_portfolio = np.log(history_portfolio/history_portfolio.shift(1)).dropna()
        
        history_q_q_q = qb.history(self.add_equity("QQQ").symbol, 252, Resolution.DAILY)
        history_q_q_q = history_q_q_q.close.unstack(0)
        pct_change_q_q_q = np.log(history_q_q_q/history_q_q_q.shift(1)).loc[pct_change_portfolio.index]
        
        m = pct_change_portfolio.shape[0]; n = pct_change_portfolio.shape[1]
        
        # Set up optimization parameters.
        p = 0.5; M = 0.0001; l = 0.01
        
        # Set up convergence tolerance, maximum iteration of optimization, iteration counter and Huber downward risk as minimization indicator.
        tol = 0.001; maxIter = 20; iters = 1; hdr = 10000
        
        # Initial weightings and placeholders.
        w_ = np.array([1/n] * n).reshape(n, 1)
        self.weights = pd.Series()
        a = np.array([None] * m).reshape(m, 1)
        c = np.array([None] * m).reshape(m, 1)
        d = np.array([None] * n).reshape(n, 1)
        
        # Iterate to minimize the HDR.
        while iters < maxIter:
            x_k = (pct_change_q_q_q.values - pct_change_portfolio.values @ w_)
            for i in range(n):
                w = w_[i]
                d[i] = d_ = 1/(np.log(1+l/p)*(p+w))
            for i in range(m):
                xk = float(x_k[i])
                if xk < 0:
                    a[i] = M / (M - 2*xk)
                    c[i] = xk
                else:
                    c[i] = 0
                    if 0 <= xk <= M:
                        a[i] = 1
                    else:
                        a[i] = M/abs(xk)
                        
            L3 = 1/m * pct_change_portfolio.T.values @ np.diagflat(a.T) @ pct_change_portfolio.values
            eig_val, eigVec = np.linalg.eig(L3.astype(float))
            eig_val = np.real(eig_val); eigVec = np.real(eigVec)
            q3 = 1/max(eig_val) * (2 * (L3 - max(eig_val) * np.eye(n)) @ w_ + eigVec @ d - 2/m * pct_change_portfolio.T.values @ np.diagflat(a.T) @ (c - pct_change_q_q_q.values))
            
            # We want to keep the upper bound of each asset to be 0.1
            mu = float(-(np.sum(q3) + 2)/n); mu_ = 0
            while mu > mu_:
                mu = mu_
                index1 = [i for i, q in enumerate(q3) if mu + q < -0.1*2]
                index2 = [i for i, q in enumerate(q3) if -0.1*2 < mu + q < 0]
                mu_ = float(-(np.sum([q3[i] for i in index2]) + 2 - len(index1)*0.1*2)/len(index2))
        
            # Obtain the weights and HDR of this iteration.
            w_ = np.amax(np.concatenate((-(mu + q3)/2, 0.1*np.ones((n, 1))), axis=1), axis=1).reshape(-1, 1)
            w_ = w_/np.sum(abs(w_))
            hdr_ = float(w_.T @ w_ + q3.T @ w_)
            
            # If the HDR converges, we take the current weights
            if abs(hdr - hdr_) < tol:
                break
            
            # Else, we would increase the iteration count and use the current weights for the next iteration.
            iters += 1
            hdr = hdr_
        
        # -----------------------------------------------------------------------------------------
        orders = []
        for i in range(n):
            orders.append(PortfolioTarget(pct_change_portfolio.columns[i], float(w_[i])))
        self.set_holdings(orders)
        
        # Recalibrate on quarter end.
        self.last_time = Expiry.end_of_quarter(self.time)