from AlgorithmImports import *
import math
import pandas as pd
from cmath import sqrt
from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data.Custom import *
from QuantConnect.Python import PythonData

class IntelligentSkyRodent(QCAlgorithm):
    def Initialize(self):
        self.cash = 100000
        self.buffer_pct = 0.03
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 8, 14)
        self.SetCash(self.cash)
        self.equities = ['BULZ','FNGU','FNGS','SPLV','TMV', 'VCIT', 'XLY', 'HIBL', 'XLK', 'XLP', 'SVXY', 'QID', 'TBF', 'TSLA', 'LQD', 'VTIP', 'EDV', 'STIP', 'SPTL', 'IEI', 'USDU', 'SQQQ', 'VIXM', 'SPXU', 'QQQ', 'BSV', 'TQQQ', 'SPY', 'DBC', 'SHV', 'IAU', 'VEA', 'UTSL', 'UVXY', 'UPRO', 'EFA', 'EEM', 'TLT', 'SHY', 'GLD', 'SLV', 'USO', 'WEAT', 'CORN', 'SH', 'DRN', 'PDBC', 'COMT', 'KOLD', 'BOIL', 'ESPO', 'PEJ', 'UGL', 'URE', 'VXX', 'UUP', 'BND', 'DUST', 'JDST', 'JNUG', 'GUSH', 'DBA', 'DBB', 'COM', 'PALL', 'AGQ', 'BAL', 'WOOD', 'URA', 'SCO', 'UCO', 'DBO', 'TAGS', 'CANE', 'REMX', 'COPX', 'IEF', 'SPDN', 'CHAD', 'DRIP', 'SPUU', 'INDL', 'BRZU', 'ERX', 'ERY', 'CWEB', 'CHAU', 'KORU', 'MEXX', 'EDZ', 'EURL', 'YINN', 'YANG', 'TNA', 'TZA', 'SPXL', 'SPXS', 'MIDU', 'TYD', 'TYO', 'TMF', 'TECL', 'TECS', 'SOXL', 'SOXS', 'LABU', 'LABD', 'RETL', 'DPST', 'DRV', 'PILL', 'CURE', 'FAZ', 'FAS', 'EWA', 'EWGS', 'EWG', 'EWP', 'EWQ', 'EWU', 'EWJ', 'EWI', 'EWN', 'ECC', 'NURE', 'VNQI', 'VNQ', 'VDC', 'VIS', 'VGT', 'VAW', 'VPU', 'VOX', 'VFH', 'VHT', 'VDE', 'SMH', 'DIA', 'UDOW', 'PSQ', 'SOXX', 'VTI', 'COST', 'UNH', 'SPHB', 'BTAL', 'VIXY', 'WEBL', 'WEBS', 'UBT', 'PST', 'TLH', 'QLD', 'SQM', 'SSO', 'SD', 'DGRO', 'SCHD', 'SGOL', 'TIP', 'DUG', 'EWZ', 'TBX', 'VGI', 'XLU', 'XLV', 'EUO', 'YCS', 'MVV', 'USD', 'BIL', 'TMF', 'EPI', 'IYK', 'DIG', 'AGG', 'PUI', 'UDN', 'QQQE', 'VTV', 'VOOG', 'VOOV']

        self.MKT = self.AddEquity("SPY",Resolution.Daily).Symbol
        self.mkt = []
        for equity in self.equities:
            self.AddEquity(equity,Resolution.Minute)
            self.Securities[equity].SetDataNormalizationMode(DataNormalizationMode.Adjusted)
        self.AddEquity('BIL',Resolution.Minute)
        self.Securities['BIL'].SetDataNormalizationMode(DataNormalizationMode.TotalReturn)
       
        self.PT1 = 0.94

        self.HT1 = {str(i).zfill(2): 0 for i in range(1,10)}
        self.HTS1 = {str(i).zfill(2): [] for i in range(1,10)}
        self.HT2 = {str(i).zfill(2): 0 for i in range(1,10)}
        self.HTS2 = {str(i).zfill(2): [] for i in range(1,10)}

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.BeforeMarketClose("SPY", 2),
                self.FunctionBeforeMarketClose)
    def RSI(self,equity,period):
        extension = min(period*5,250)
        r_w = RollingWindow[float](extension)
        history = self.History(equity,extension - 1,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < extension:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)
        if r_w.IsReady:
            average_gain = 0
            average_loss = 0
            gain = 0
            loss = 0
            for i in range(extension - 1,extension - period -1,-1):
                gain += max(r_w[i-1] - r_w[i],0)
                loss += abs(min(r_w[i-1] - r_w[i],0))
            average_gain = gain/period
            average_loss = loss/period
            for i in range(extension - period - 1,0,-1):
                average_gain = (average_gain*(period-1) + max(r_w[i-1] - r_w[i],0))/period
                average_loss = (average_loss*(period-1) + abs(min(r_w[i-1] - r_w[i],0)))/period
            if average_loss == 0:
                return 100
            else:
                rsi = 100 - (100/(1 + average_gain / average_loss))
                return rsi
        else:
            return None
    def CumReturn(self,equity,period):
        history = self.History(equity,period,Resolution.Daily)
        closing_prices = pd.Series([bar.Close for bar in history])
        current_price = self.Securities[equity].Price
        closing_prices = closing_prices.append(pd.Series([current_price]))
        first_price = closing_prices.iloc[0]
        if first_price == 0:
            return None
        else:
            return_val = (current_price / first_price) - 1
            return return_val
    def STD(self,equity,period):
        r_w = RollingWindow[float](period + 1)
        r_w_return = RollingWindow[float](period)
        history = self.History(equity,period,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < period + 1:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)
        for i in range (period,0,-1):
            daily_return = (r_w[i-1]/r_w[i] - 1)
            r_w_return.Add(daily_return)
        dfstd = pd.DataFrame({'r_w_return':r_w_return})
        if r_w.IsReady:
            std = dfstd['r_w_return'].std()
            if std == 0:
                return 0
            else:
                return std
        else:
            return 0
    def MaxDD(self,equity,period):
        history = self.History(equity,period - 1,Resolution.Daily)
        closing_prices = pd.Series([bar.Close for bar in history])
        current_price = self.Securities[equity].Price
        closing_prices = closing_prices.append(pd.Series([current_price]))
        rolling_max = closing_prices.cummax()
        drawdowns = (rolling_max - closing_prices) / rolling_max
        max_dd = drawdowns.min()
        return max_dd
    def SMA(self,equity,period):
        r_w = RollingWindow[float](period)
        history = self.History(equity,period - 1,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < period:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)        
        if r_w.IsReady:
            sma = sum(r_w) / period
            return sma
        else:
            return 0
    def IV(self,equity,period):
        r_w = RollingWindow[float](period + 1)
        r_w_return = RollingWindow[float](period)
        history = self.History(equity,period,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < period + 1:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)
        for i in range (period,0,-1):
            if r_w[i] == 0:
                return 0
            else:
                daily_return = (r_w[i-1]/r_w[i] - 1)
                r_w_return.Add(daily_return)
        dfinverse = pd.DataFrame({'r_w_return':r_w_return})       
        if r_w.IsReady:
            std = dfinverse['r_w_return'].std()
            if std == 0:
                return 0
            else:
                inv_vol = 1 / std
                return inv_vol
        else:
            return 0
    def SMADayRet(self,equity,period):
        r_w = RollingWindow[float](period + 1)
        r_w_return = RollingWindow[float](period)
        history = self.History(equity,period,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < period + 1:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)
        for i in range (period,0,-1):
            if r_w[i] == 0:
                return None
            daily_return = (r_w[i-1]/r_w[i] - 1)
            r_w_return.Add(daily_return)
        if r_w.IsReady:
            smareturn = sum(r_w_return) / period
            return smareturn
        else:
            return 0
    def EMA(self,equity,period):
        extension = period + 50
        r_w = RollingWindow[float](extension)
        history = self.History(equity,extension - 1,Resolution.Daily)
        for historical_bar in history:
            r_w.Add(historical_bar.Close)
        while r_w.Count < extension:
            current_price = self.Securities[equity].Price
            r_w.Add(current_price)
        if r_w.IsReady:
            total_price = 0
            for i in range(extension - 1,extension - period - 2,-1):
                total_price += r_w[i]
            average_price = total_price/period
            for i in range(extension - period - 2,-1,-1):
                average_price = r_w[i]*2/(period+1) + average_price*(1-2/(period+1))
            return average_price
        else:
            return None
    def Sort(self,sort_type,equities,period,reverse,number,multiplier):
        self.PT = getattr(self,f"PT{number}") * multiplier
        returns = {}
        for equity in equities:
            returns[equity] = getattr(self,sort_type)(equity,period)
        s_e = sorted([item for item in returns.items() if item[1] is not None],key = lambda x: x[1],reverse = reverse)
        t3e = s_e[:1]
        ht = getattr(self,f"HT{number}")
        hts = getattr(self,f"HTS{number}")
        for i in ht.keys():
            if ht[i] == 0:
                ht[i] = self.PT
                hts[i].append(t3e[0][0])
                break
        setattr(self,f"HT{number}",ht)
        setattr(self,f"HTS{number}",hts)
    def AH(self, equities, PTnumber, multiplier): #AppendHolding
        if not isinstance(equities, list):
            equities = [equities]
        
        HT = getattr(self, f"HT{PTnumber}")
        HTS = getattr(self, f"HTS{PTnumber}")
        PT = getattr(self, f"PT{PTnumber}") * multiplier
        
        for equity in equities:
            for i in HT.keys():
                if HT[i] == 0:
                    HT[i] = PT
                    HTS[i].append(equity)
                    break

    def OnData (self,data):
        pass
    def FunctionBeforeMarketClose(self):
        mkt_price = self.History(self.MKT,2,Resolution.Daily)['close'].unstack(level= 0).iloc[-1]
        self.mkt.append(mkt_price)
        mkt_perf = self.cash * self.mkt[-1] / self.mkt[0]
        self.Plot('Strategy Equity',self.MKT,mkt_perf)
        self.tqqq_ftlt()
        self.ExecuteTrade()

    
    def sideways_market_deleverage(self):
            if self.Securities['SPY'].Price > self.SMA('SPY', 20):
                self.AH('SPY', 1, 1.0)
            else:
                if self.RSI('TLT', 10) > self.RSI('SQQQ', 10):
                    self.AH('QQQ', 1, 1.0)
                else:
                    self.AH('PSQ', 1, 1.0)
            if self.RSI('TLT', 10) > self.RSI('SQQQ', 10):
                self.AH('QQQ', 1, 1.0)
            else:
                self.AH('PSQ', 1, 1.0)
    
    def nasdaq_in_crash_territory_time_to_deleverage(self):
        if self.Securities['QQQ'].Price < self.SMA('QQQ', 20):
            if self.CumReturn('QQQ', 60) <= -12:
                self.sideways_market_deleverage()
            else:
                if self.RSI('TLT', 10) > self.RSI('SQQQ', 10):
                    self.AH('TQQQ', 1, 1.0)
                else:
                    self.AH('SQQQ', 1, 1.0)
        else:
            if self.RSI('SQQQ', 10) < 31:
                self.AH('PSQ', 1, 1.0)
            else:
                if self.CumReturn('QQQ', 10) > 5.5:
                    self.AH('PSQ', 1, 1.0)
                else:
                    self.Sort('RSI',('QQQ','SMH'),10,True,1,1.0)
    
    def bear_market_sideways_protection(self):
        if self.CumReturn('QQQ', 252) < -20:
            self.nasdaq_in_crash_territory_time_to_deleverage()
        else:
            if self.Securities['QQQ'].Price < self.SMA('QQQ', 20):
                if self.RSI('TLT', 10) > self.RSI('SQQQ', 10):
                    self.AH('TQQQ', 1, 1.0)
                else:
                    self.AH('SQQQ', 1, 1.0)
            else:
                if self.RSI('SQQQ', 10) < 31:
                    self.AH('SQQQ', 1, 1.0)
                else:
                    if self.CumReturn('QQQ', 10) > 5.5:
                        self.AH('SQQQ', 1, 1.0)
                    else:
                        self.Sort('RSI',('TQQQ','SOXL'),10,True,1,1.0)
        if self.Securities['QQQ'].Price < self.SMA('QQQ', 20):
            if self.CumReturn('QQQ', 60) <= -12:
                self.sideways_market_deleverage()
            else:
                if self.RSI('TLT', 10) > self.RSI('SQQQ', 10):
                    self.AH('TQQQ', 1, 1.0)
                else:
                    self.AH('SQQQ', 1, 1.0)
        else:
            if self.RSI('SQQQ', 10) < 31:
                self.AH('SQQQ', 1, 1.0)
            else:
                if self.CumReturn('QQQ', 70) < -15:
                    self.Sort('RSI',('TQQQ','SOXL'),10,True,1,1.0)
                else:
                    self.Sort('CumReturn',('SPY','QQQ','DIA','XLP'),15,True,1,1.0)
    
    def dip_buy_strategy(self):
        if self.RSI('TQQQ', 10) < 31:
            self.AH('TECL', 1, 1.0)
        else:
            if self.RSI('SMH', 10) < 30:
                self.AH('SOXL', 1, 1.0)
            else:
                if self.RSI('SPY', 10) < 30:
                    self.AH('UPRO', 1, 1.0)
                else:
                    self.bear_market_sideways_protection()
    
    def svxy_ftlt(self):
        if self.RSI('QQQ', 10) > 80:
            self.AH('VIXY', 1, 0.33)
        else:
            if self.RSI('SPY', 10) > 80:
                self.AH('VIXY', 1, 0.33)
            else:
                if self.RSI('QQQ', 10) < 30:
                    self.AH('XLK', 1, 0.33)
                else:
                    if self.Securities['SVXY'].Price > self.SMA('SVXY', 24):
                        self.Sort('SMADayRet',('SVXY','SPY'),20,True,1,0.33)
                    else:
                        self.AH('BTAL', 1, 0.33)
    
    def spy_cooloff_period(self):
        self.AH('SHV', 1, 1.0)
    
    def bull_market(self):
            if self.RSI('QQQ', 10) > 80:
                self.AH('UVXY', 1, 1.0)
            else:
                if self.RSI('SPY', 10) > 80:
                    self.AH('UVXY', 1, 1.0)
                else:
                    if self.RSI('SPY', 60) > 60:
                        self.spy_cooloff_period()
                    else:
                        self.Sort('SMADayRet',('TQQQ','TECL','UDOW','UPRO','BULZ'),21,True,1,0.67)
                        self.svxy_ftlt()
    
    def tqqq_ftlt(self):
        if self.Securities['SPY'].Price > self.SMA('SPY', 200):
            self.bull_market()
        else:
            self.dip_buy_strategy()

    def ExecuteTrade(self):
        group1 = {
            'HTS': [self.HTS1[i][0] if len(self.HTS1[i]) == 1 else self.HTS1[i] for i in self.HTS1],
            'HT': [self.HT1[i] for i in self.HT1]
        }
        df1 = pd.DataFrame(group1)
        group2 = {
            'HTS': [self.HTS2[i][0] if len(self.HTS2[i]) == 1 else self.HTS2[i] for i in self.HTS2],
            'HT': [self.HT2[i] for i in self.HT2]
        }
        df2 = pd.DataFrame(group2)
        df = pd.concat([df1, df2])
        df['HTS'] = df['HTS'].astype(str)
        result = df.groupby(['HTS']).sum().reset_index()
        for equity in self.equities:
            if all(not pd.isnull(result.iloc[i,0]) and not equity == result.iloc[i,0] for i in range(len(result))):
                if self.Portfolio[equity].HoldStock:
                    self.Liquidate(equity)
        output = "*****"
        for i in range(len(result)):
            if result.iloc[i,0]:
                percentage = round(result.iloc[i,1] * 100,2)
                output += "{}: {}% - ".format(result.iloc[i,0],percentage)
        output = output.rstrip(" - ")
        self.Log(output)
        for i in range(len(result)):
            if not result.iloc[i,1] == 0 and not result.iloc[i,0] == 'BIL':
                percentage_equity = self.Portfolio[result.iloc[i,0]].HoldingsValue / self.Portfolio.TotalPortfolioValue
                if result.iloc[i,1] < percentage_equity and abs(result.iloc[i,1] / percentage_equity - 1) > self.buffer_pct:
                    self.SetHoldings(result.iloc[i,0],result.iloc[i,1])
                else:
                    pass
        for i in range(len(result)):
            if not result.iloc[i,1] == 0 and not result.iloc[i,0] == 'BIL':
                percentage_equity = self.Portfolio[result.iloc[i,0]].HoldingsValue / self.Portfolio.TotalPortfolioValue
                if result.iloc[i,1] > percentage_equity and abs(percentage_equity / result.iloc[i,1] - 1) > self.buffer_pct:
                    self.SetHoldings(result.iloc[i,0],result.iloc[i,1])
                else:
                    pass

        self.HT1 = {str(i).zfill(2): 0 for i in range(1,10)}
        self.HTS1 = {str(i).zfill(2): [] for i in range(1,10)}
        self.HT2 = {str(i).zfill(2): 0 for i in range(1,10)}
        self.HTS2 = {str(i).zfill(2): [] for i in range(1,10)}