#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:
            
            if self.EntriesToday(tm) < self.algo.entries_per_day:
                entered = self.EntryLogic(bar)
                if entered:
                    self.IncrementEntriesToday(tm)

        if self.Invested:
            self.ExitPackage()

            if self.algo.signal_exit:
                if self.IsLong and bar.Close < self.vwaps[0]:
                    self.algo.Liquidate(self.symbol, "VWAP LX")
            
                if self.IsShort and bar.Close > self.vwaps[0]:
                    self.algo.Liquidate(self.symbol, "VWAP SX")


    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 and not self.IsLong
        if long_ok:
            sma_ok = bar.Close > self.sma.Current.Value
            sma_trend = self.smas[0] > self.smas[self.algo.n]
            above_vwap = bar.Close > self.vwaps[0]

            # A bit hacked -- just using what we have available. (No vwap bands)
            # below_lower = bar.Close < self.vwaps[0]  - self.sigma.Current.Value * self.algo.nstd

            # Removed sma_ok and sma_trend
            if above_vwap:
                self.algo.SetHoldings(self.symbol, alloc_pct)
                return True 

        short_ok = self.algo.direction <= 0 and not self.IsShort
        if short_ok:
            below_vwap = bar.Close < self.vwaps[0]

            if below_vwap:
                self.algo.SetHoldings(self.symbol, -alloc_pct)
                return True

    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

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

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

    def EntriesToday(self, bar_end_time):
        date = bar_end_time.date()
        if date not in self.entries_by_date: 
            self.entries_by_date[date] = 0
        return self.entries_by_date[date]

    def IncrementEntriesToday(self, bar_end_time):
        date = bar_end_time.date()
        if date not in self.entries_by_date:
            self.entries_by_date[date] = 1
        else:
            self.entries_by_date[date] += 1

# region imports
from AlgorithmImports import *
# endregion

from Strategy import Strategy

class SwimmingRedTermite(QCAlgorithm):

    tickers = ['QQQ']

    length = 50
    nstd = 2

    # Direction -- 1 is long only, -1 is short only, 0 is both.
    direction = 1

    entries_per_day = 100

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

    # If 0, will not set. 
    # .01 == 1%
    tgt_pct = 0.0
    stop_pct = 0.0 
    signal_exit = True
    eodx = True

    

    symbols = []
    strats = {}

    def Initialize(self):
        self.SetStartDate(2022, 1, 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):
        if self.eodx:
            self.Liquidate(tag="EOD")

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