# region imports
from AlgorithmImports import *
from statsmodels.tsa.vector_ar.vecm import VECM
# endregion

class CreativeRedBadger(QCAlgorithm):
    tickers = ['SPY','QQQ','IWM','DIA','EZU', 'EWJ','FXI','EWT','EWG','EWC','EWI','EWW','INDA']

    length = 10
    sigma = 1.5
    res = Resolution.Hour # Daily, Hour, Minute

    time_exit_bars = 10 # Same as length? (HL?)
    extra_edge = 0.05

    # ---------------------------------------------------

    wt_arr = None # Numpy array.
    wt_vec = []
    wts_by_symbol = {}

    series_ask_terms_by_symbol = {}
    series_bid_terms_by_symbol = {}

    position = 0
    bar_ct = 0

    def Initialize(self):
        self.SetStartDate(2020, 4, 19)
        # self.SetEndDate(2022,5, 20)
        self.SetCash(100000)

        self.BuyBasis = RollingWindow[float](self.length)
        self.SellBasis = RollingWindow[float](self.length)
        self.Basis = RollingWindow[float](self.length)

        self.bb = BollingerBands(self.length, self.sigma)

        symbols = []
        for t in self.tickers:
            try:
                symbol = self.AddEquity(t, self.res).Symbol
                symbols.append(symbol)
            except:
                pass 

        self.SetSecurityInitializer(lambda x: x.SetFeeModel(ConstantFeeModel(0)))  
        # # 0 fees, for now.
        # for security in self.Securities:
        #     security.SetFeeModel(ConstantFeeModel(0))

        # Only run this weekly?
        self.Schedule.On(self.DateRules.WeekStart(0),
                 self.TimeRules.AfterMarketOpen("SPY", -10),
                 self.BeforeMarketOpen)

        

    def BeforeMarketOpen(self):
        ## TODO: when we are dealing with quoting, and bid/ask spreads -- use this.
        # quote_bars_df = self.History(QuoteBar, self.Securities.Keys, 5, Resolution.Minute)
        df = self.History(self.Securities.Keys, 10000, self.res) # IF fails, use DAILY, or HOURLY
        closes = df.unstack(level=0).close
        closes.dropna(inplace=True, how='any')
        self.Log(f'Closes -- head: {closes.head()}')

        closes = closes[[i for i in self.Securities.Keys]] # Addit, for order safety.

        vecm = VECM(closes, deterministic='n', k_ar_diff=1)
        vecm_fit = vecm.fit()

        norm_wts = vecm_fit.beta / np.abs(vecm_fit.beta).sum()

        self.wt_vec = norm_wts[:,0] # (n,) vec
        self.wt_arr = norm_wts


        # THIS might have broken something... this was using Securities.Keys
        for n, k in enumerate(self.Securities.Keys): # Ensure columns line up w this order.
            self.wts_by_symbol[k] = self.wt_vec[n]

        self.Log(f'Weights: {self.wts_by_symbol.items()}')
        self.basis = np.dot(closes, vecm_fit.beta)


        # Valid...
        # self.long_targets = [PortfolioTarget(self.s1, per_symbol), PortfolioTarget(self.s2, -per_symbol)]
        # self.short_targets = [PortfolioTarget(self.s1, -per_symbol), PortfolioTarget(self.s2, per_symbol)]
        # self.flat_targets = [PortfolioTarget(self.s1, 0.0), PortfolioTarget(self.s2, 0.0)]


        self.long_targets = []
        self.flat_targets = []
        self.short_targets = []
        for symbol, wt in self.wts_by_symbol.items():
            print(type(symbol), symbol)
            # self.Log(f'type: {type(symbol)}, symbol: {symbol} -- wt: {wt}, type: {type(wt)}')

            pft = PortfolioTarget(symbol, wt)
            self.long_targets.append(pft)

            pft = PortfolioTarget(symbol, 0.0)
            self.flat_targets.append(pft)

            pft = PortfolioTarget(symbol, -1.0 * wt)
            self.short_targets.append(pft)

        # WARMUP? 
        basis = pd.DataFrame(self.basis[:,0], index=closes.index, columns=['Series'])
        for dt, row in basis.iterrows():
            # self.Log(f'dt: {dt}, Basis: {row}, {row.Series}')
            b = row.Series
            self.bb.Update(dt, b)
            self.Basis.Add(b)


    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            security.SetFeeModel(ConstantFeeModel(0.0))

    def OnData(self, data: Slice):

        # Check data valid
        # data = data.QuoteBars
        data = data.Bars
        if not data.ContainsKey('SPY'):
            self.Log('No Data -- return out')
            return 

        # Check weights valid (Model Fit)
        if len(self.wt_vec) == 0: return 


        # # How do we price this out, to quote them iteratively? 
        # # i.e. pull a wrt b-f, etc.
        # # Price each 'term', maybe?
        curr_basis = 0
        for symbol, wt in self.wts_by_symbol.items():
            
            try:
                # If wt is positive, we are BUYING this, so price it wrt ASK (for pos spread)
                buy_term = data[symbol].Ask.Close * wt 
                sell_term = data[symbol].Bid.Close * wt
            except:
                buy_term = data[symbol].Close 
                sell_term = buy_term 

            # Assuming this is a buy, our bid term would be pricing the series as if buying
            # thus, for bid, we are pricing at ask of positive weighted symbols (BUY Spread)
            #       for ask, we are pricing at bid of negative weighted symbols.(SELL Spread)
            term = buy_term if wt > 0 else sell_term 
            self.series_bid_terms_by_symbol[symbol] = buy_term if wt > 0 else buy_term 
            self.series_ask_terms_by_symbol[symbol] = sell_term if wt > 0 else sell_term 
            curr_basis += term
        
        self.Basis.Add(curr_basis)
        self.bb.Update(self.Time, curr_basis)

        # # This is a bit muddled, but it should work.
        # if not self.bb.IsReady: return 

        # ## TODO: price out each leg, wrt the rest...
        # # Prototype elsewhere, then bring into this.

        upper, lower, mid = self.bb.UpperBand.Current.Value, self.bb.LowerBand.Current.Value, self.bb.MiddleBand.Current.Value
        self.Log(f'Basis BB: \
                    Upper: {self.bb.UpperBand.Current.Value},\
                    Lower: {self.bb.LowerBand.Current.Value}\
                    Middle: {self.bb.MiddleBand.Current.Value}')

        
        curr = curr_basis
        if self.position == 0:
            # could sum the bid_terms for bid basis. 
            if curr < lower - self.extra_edge:
                self.position = 1
                self.SetHoldings(self.short_targets, False, f"LE -- Series: {curr}")
            
            if curr > upper + self.extra_edge:
                self.position = -1
                self.SetHoldings(self.long_targets, False, f"SE -- Series: {curr}")


        # Increment bar count
        if self.position != 0:
            self.bar_ct += 1
            if self.bar_ct > self.time_exit_bars:
                self.Liquidate()
        
        opnl = self.Portfolio.TotalUnrealisedProfit
        # Long -- Looking to exit.
        if self.position > 0:
            if curr > mid:
                self.position = 0
                self.bar_ct = 0
                self.SetHoldings(self.flat_targets, False, f"LX -- Series: {curr}, PNL: {opnl}")
                return 
            

        # Short -- Looking to exit.
        if self.position < 0:
            if curr < mid:
                self.position = 0
                self.bar_ct = 0
                self.SetHoldings(self.flat_targets, False, f"SX -- Series: {curr}, PNL: {opnl}")
                return