#region imports
from AlgorithmImports import *
#endregion


class Strategy:

    def __init__(self, algo, symbol):
        self.algo = algo 
        self.symbol = symbol

        self.day_close = None
        self.day_open = None

        self.entries_by_date = {}

        self.high_of_day = None
        self.low_of_day = None

        # Just to see how it's done -- not too bad.
        self.bb = self.algo.BB(self.symbol, self.algo.length, self.algo.nstd, Resolution.Minute)
        self.vwap = self.algo.VWAP(self.symbol)
        self.sma = self.algo.SMA(self.symbol, self.algo.length)
        self.sigma = self.algo.STD(self.symbol, self.algo.length)
         
        self.smas = RollingWindow[float](self.algo.n + 1)
        self.vwaps = RollingWindow[float](self.algo.n + 1)

    # Decided on a delegate method, rather than a consolidator 
    # just simpler, less boilerplate.
    def OnData(self, data: Slice):
        data = data.Bars
        if not data.ContainsKey(self.symbol): return 
        bar = data[self.symbol]
        tm = bar.EndTime

        self.UpdateWindows(bar)
        if self.IsReady:
            
            self.EntryLogic(bar)

        if self.Invested:
            self.ExitPackage()


    def UpdateWindows(self, bar):
        if self.sma.IsReady:
            self.smas.Add(self.sma.Current.Value)

        if self.vwap.IsReady:
            self.vwaps.Add(self.vwap.Current.Value)

    def EntryLogic(self, bar):
        alloc_pct = 1.0 / len(self.algo.strats) 

        long_ok = self.algo.direction >= 0
        sma_ok = bar.Close > self.sma.Current.Value
        sma_trend = self.smas[0] > self.smas[self.algo.n]
        # A bit hacked -- just using what we have available.
        below_lower = bar.Close < self.vwaps[0] - self.sigma.Current.Value * self.algo.nstd

        if long_ok and sma_ok and sma_trend and below_lower:
            self.algo.SetHoldings(self.symbol, alloc_pct)


    def ExitPackage(self):
        self.StopPct(self.algo.stop_pct)
        self.TgtPct(self.algo.tgt_pct)

    def StopPct(self, pct):
        if pct == 0.0: return

        urpct = self.algo.Portfolio[self.symbol].UnrealizedProfitPercent
        ur = self.algo.Portfolio[self.symbol].UnrealizedProfit
        if self.algo.Portfolio[self.symbol].UnrealizedProfitPercent < -1 * pct:
            self.algo.Liquidate(self.symbol, tag=f"Stop -- {ur}")
    
    def TgtPct(self, pct):
        if pct == 0.0: return

        ur = self.algo.Portfolio[self.symbol].UnrealizedProfit
        if self.algo.Portfolio[self.symbol].UnrealizedProfitPercent > pct:
            self.algo.Liquidate(self.symbol, tag=f"Tgt -- {ur}")
    
    @property
    def Sigma(self):
        if self.IsReady:
            return self.bb.Upper.Current.Value - self.bb.Middle.Current.Value 

    @property
    def IsReady(self):
        return self.bb.IsReady and self.vwap.IsReady and self.sma.IsReady and self.vwaps.IsReady

    @property
    def Invested(self):
        return self.algo.Portfolio[self.symbol].Invested

# region imports
from AlgorithmImports import *
# endregion

from Strategy import Strategy

class SwimmingRedTermite(QCAlgorithm):

    tickers = ['QQQ']

    entry_style_1 = True

    length = 50
    nstd = 2
    direction = 1

    # Bars back, used for slope measure.
    n = 7

    # If 0, will not set. 
    # .01 == 1%
    tgt_pct = 0
    stop_pct = 0 

    symbols = []
    strats = {}

    def Initialize(self):
        self.SetStartDate(2022, 9, 1)
        self.SetCash(100000)

        # Set benchmark -- SPY (most common, need it unrelated to tickers)
        self.AddEquity('SPY', Resolution.Minute)
        self.SetBenchmark("SPY")

        # symbol = self.AddEquity("QQQ", Resolution.Minute).Symbol
        for t in self.tickers:
            s = self.AddEquity(t, Resolution.Minute).Symbol
            self.symbols.append(s)
            self.strats[s] = Strategy(self, s)

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY", 1),
                 self.EODx)

    def EODx(self):
        self.Liquidate(tag="EOD")

    def OnData(self, data: Slice):
        gap_pcts = []
        for symbol, strat in self.strats.items():
            strat.OnData(data)