import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from scipy.stats import linregress
import decimal as d

class  MyFrameworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 9, 1)  # Set Start Date
        self.SetCash(10000)  # Set Strategy Cash
        #self.SetEndDate(2018, 1, 1)
        
        # create a dictionary to store momentum indicators for all symbols 
        self.back_period = 21     #  months
        self.vol_period = 21    # days for calc vol
        self.x = np.asarray(range(self.vol_period))
        self.Log(self.x) 
        self.ematwenty = {}
        
        #Warmup period 
        period = 200
        #MOM period
        period1 = 20
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.20))
        
        # warm up the MOM indicator
        self.SetWarmUp(period)
        self.symbols = []
        self.tickers = ["SPY", 
                        "TLT"]
                        
        for ticker in self.tickers:
            symbol = self.AddEquity(ticker, Resolution.Hour).Symbol
            self.symbols.append(symbol)
            self.ematwenty[symbol] = self.EMA(symbol, 91, Resolution.Hour)
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), \
                 self.TimeRules.At(10, 31), \
                 self.Rebalance)
    
    
     ###################################### 
    def rsquared(self, x, y):
        # slope, intercept, r_value, p_value, std_err
        _, _, r_value, _, _ = linregress(x, y)
        return r_value**2
    ######################################
    
    def Rebalance(self):
        if self.IsWarmingUp: return
    
        allPrices = self.History(self.symbols, self.back_period, Resolution.Daily).close.unstack(level=0)
        Alldata = pd.DataFrame(allPrices[symbol] for symbol in self.symbols)
        arr = Alldata.to_numpy()
        rsq = np.apply_along_axis(self.rsquared, 1, arr, self.x)
    
        Enterdata = pd.DataFrame(list(zip(self.tickers)), columns =['tickers'])
        Enterdata.insert(1, "Price", [self.Portfolio[symbol].Price for symbol in self.symbols], True)
        Enterdata.insert(2, "ema", [self.ematwenty[symbol].Current.Value for symbol in self.symbols], True)
        Enterdata['rsq'] = rsq
        self.Log(Enterdata)
    
        ema20 = Enterdata[Enterdata['Price'] > Enterdata['ema']]
    
        top3 = ema20.sort_values(by='rsq', ascending = False)[:3]
        top3_symbols = [self.symbols[i] for i in top3.index]
        
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and (security_hold.Symbol not in top3_symbols):
                self.Liquidate(security_hold.Symbol)
        
        added_symbols = []        
        for symbol in top3_symbols:
            if not self.Portfolio[symbol].Invested:
                added_symbols.append(symbol)
        for added in added_symbols:
            self.SetHoldings(added, 1.0/len(added_symbols))
            self.Log(added_symbols)
            
        return