from AlgorithmImports import *
import pandas as pd

class HipsterFluorescentPinkRabbit(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2023, 3, 8)
        self.SetCash(100000)  

        self.DefaultOrderProperties.TimeInForce = TimeInForce.Day

        spy = self.AddEquity("SPY", Resolution.Daily)
        spy.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        self.spy = spy.Symbol

        self.ema8 = self.EMA(self.spy, 8, Resolution.Daily)
        self.ema20 = self.EMA(self.spy, 20, Resolution.Daily)
        self.sma200 = self.SMA(self.spy, 200, Resolution.Daily)
        self.atr = self.ATR(self.spy, 14, MovingAverageType.Wilders)
        self.bb = self.BB(self.spy, 20, 2)
        self.kch2 = self.KCH(self.spy, 20, 2)
        self.kch15 = self.KCH(self.spy, 20, 1.5)
        self.kch1 = self.KCH(self.spy, 20, 1)

    def OnData(self, data):
        if not data.ContainsKey(self.spy):
            return 
            
        emaVal = round(self.ema8.Current.Value, 2)
        smaVal = round(self.sma200.Current.Value, 2)
        
        hstv = self.History(self.spy, 500, Resolution.Daily)
        histClose = pd.DataFrame(hstv)["close"]
        tmp1 = pd.Series(np.where(histClose > histClose.shift(1), histClose - histClose.shift(1), 0))
        tmp2 = pd.Series(np.where(histClose < histClose.shift(1), histClose.shift(1) - histClose, 0))
        d2 = round(pd.Series(tmp1.rolling(21).sum()), 2)
        d4 = round(pd.Series(tmp2.rolling(21).sum()), 2)
        cond = d2 + d4 == 0
        ad3__ = (d2 - d4) / (d2 + d4) * 100
        ad3_ = ad3__.fillna(0)
        ad3 = pd.Series(np.where(cond, 0, ad3_))
        coeff = 2 / (21 + 1) * abs(ad3) / 100
        newAsd = histClose.copy()
        for i in range(1, len(histClose)):
            newAsd[i] = coeff[i] * histClose[i] + newAsd[i-1] * (1-coeff[i])
        vmaToday = round(newAsd[-1], 2)
        vmaYes = round(newAsd[-2], 2)
        atrVals = self.atr.Current.Value
        vmaColor = self.getVMAColor(vmaToday, vmaYes, atrVals)

        uband = self.bb.UpperBand.Current.Value
        lband = self.bb.LowerBand.Current.Value
        UKCWide = self.kch2.UpperBand.Current.Value
        UKCNormal = self.kch15.UpperBand.Current.Value
        UKCNarrow = self.kch1.UpperBand.Current.Value
        lowKCWide = self.kch2.LowerBand.Current.Value
        loeKCNormal = self.kch15.LowerBand.Current.Value
        lowKCNarrow = self.kch1.LowerBand.Current.Value
        inSqueeze = self.getSqueeze(uband, lband, UKCWide, UKCNormal, UKCNarrow, lowKCWide, loeKCNormal, lowKCNarrow)

        histColor = self.getHistogramColor(self.ema20.Current.Value)

        cond1 = emaVal > vmaToday > smaVal
        cond2 = data.Bars[self.spy].High > vmaToday and data.Bars[self.spy].Low < vmaToday
        
        if not self.Portfolio.Invested:
            if cond1 and cond2:
                self.entryTicket = self.StopLimitOrder(self.spy, 100, round(data.Bars[self.spy].High+.05, 2), round(data.Bars[self.spy].High+.05, 2))
                self.Log(self.entryTicket)


    def getVMAColor(self, vmaT, vmaY, atrVal):
        color = ""
        if abs(vmaT - vmaY) / atrVal <= .05:
            color = "Grey"
        elif vmaT > vmaY:
            color = "Green"
        else:
            color = "Red"
        return color
    
    def getSqueeze(self, uband2, lband2, upKCWide, upKCNormal, upKCNarrow, lowKCWide, lowKCNormal, lowKCNarrow):
        color = ""
        if (lband2 >= lowKCWide) and (uband2 <= upKCWide):
            color = "Black"
        if (lband2 >= lowKCNormal) and (uband2 <= upKCNormal):
            color = "Red"
        if (lband2 >= lowKCNarrow) and (uband2 <= upKCNarrow):
            color = "Orange"
        if (lband2 < lowKCWide) and (uband2 > upKCWide):
            color = "Green"
        return color
    
    def getHistogramColor(self, ema20):
        length_KC = 20
        hstv = self.History(self.spy, 40, Resolution.Daily)
        histClose = pd.DataFrame(hstv)["close"]
        histHigh = pd.DataFrame(hstv)["high"]
        histLow = pd.DataFrame(hstv)["low"]
        highest = histHigh.rolling(window=length_KC).max()
        lowest = histLow.rolling(window=length_KC).min()
        m1 = round((highest + lowest) / 2 + ema20, 2)
        hstv["value"] = histClose - m1 / 2
        fit_y = np.array(range(0, length_KC))
        hstv['value'] = round(hstv['value'].rolling(window=length_KC).apply(lambda x: np.polyfit(fit_y, x, 1)[0] * (length_KC - 1) + np.polyfit(fit_y, x, 1)[1], raw=True), 2)
        histVals = hstv["value"]
        myHist = pd.DataFrame(histVals)
        myHist = myHist.dropna()
        myHist["PV"] = myHist.shift(1)
        myHist = myHist.dropna()

        def determineHistColor(curr, prev):
            color = ""
            if curr > 0:
                if curr > prev:
                    color = "LB"
                else:
                    color = "DB"
            else:
                if curr < prev:
                    color = "Red"
                else:
                    color = "Yellow"
            return color
        myHist["HistColor"] = myHist.apply(lambda x: determineHistColor(x["value"], x["PV"]), axis=1)
        return myHist["HistColor"][-1]