#region imports
from AlgorithmImports import *
#endregion

# Your New Python File

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, 50, 2, Resolution.Minute)
        self.atr = self.algo.ATR(self.symbol, 20, resolution=Resolution.Minute)

    
    def OnData(self, data: Slice):

        data = data.Bars
        if not data.ContainsKey(self.symbol): return 
        bar = data[self.symbol]
        tm = bar.EndTime


        # Update things needed for strategy
        self.UpdateHighLowDay(bar)
        self.UpdateGaps(bar)


        # Entry Logic 
        if not self.Invested:
            if self.EntriesToday(tm) < self.algo.entries_per_day:
                entered = self.EntryLogic(bar)
                if entered:
                    self.IncrementEntriesToday(tm)
            
        # Exit Logic (Stops)
        if self.Invested:
            self.ExitPackage()

    #region Entry Logic 


    def EntryLogic(self, bar):
        # TODO: consider using atr, a minimal range, to improve signal for GC
        # if self.bb.IsReady:
        #     self.algo.Log(f'{self.symbol} -- bb: {self.bb.MiddleBand.Current.Value}')
        #     self.algo.Log(f'{self.symbol} -- atr: {self.atr.Current.Value}')

        alloc_pct = 1.0 / len(self.algo.strats) 
        
        # if self.algo.gap_fill and abs(self.GapPct) >= 2.0: return False

        # GAP UP 
        if self.GapPct > self.algo.thres:
            # Gap up, we are selling if we can with dir.
            if self.algo.gap_fill:
                if self.algo.dir <= 0:
                    self.algo.SetHoldings(self.symbol, -1 * alloc_pct)
                    return True
            # Gap up, we are buying if we can w dir.
            else:
                # tried testing this high day thing... #  and bar.High == self.high_of_day:
                if self.algo.dir >= 0: 
                    self.algo.SetHoldings(self.symbol, alloc_pct)
                    return True
        
        # GAP DN 
        if self.GapPct < -1 * self.algo.thres:
            # Buy, on gap down
            if self.algo.gap_fill:
                # and bar.High >= ((self.high_of_day - self.low_of_day)*.5) + self.low_of_day (kinda works..)
                if self.algo.dir >= 0:
                    self.algo.SetHoldings(self.symbol, alloc_pct)
                    return True
            else:
                if self.algo.dir <= 0:
                    self.algo.SetHoldings(self.symbol, -1 * alloc_pct)
                    return True
        
        return False

    #endregion 


    # region Exit Logic 

    def ExitPackage(self):
        self.StopPct()
        self.TgtPct()

    def StopPct(self, pct=None):
        urpct = self.algo.Portfolio[self.symbol].UnrealizedProfitPercent
        ur = self.algo.Portfolio[self.symbol].UnrealizedProfit
        # self.algo.Log(f'{self.symbol} -- {urpct} -- {ur}')

        if not pct:
            pct = self.algo.stop_pct

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

        ur = self.algo.Portfolio[self.symbol].UnrealizedProfit
        if self.algo.Portfolio[self.symbol].UnrealizedProfitPercent > pct:
            self.algo.Liquidate(self.symbol, tag=f"Tgt -- {ur}")

    # endregion 

    # region Properties 
        
    @property
    def Gap(self):
        if self.day_close and self.day_open:
            return self.day_open - self.day_close
        return 0 # 0 or None?


    @property
    def GapPct(self):
        if self.day_close and self.day_open:
            return ((self.day_open - self.day_close) / self.day_close) * 100
        return 0 # or none?

    @property
    def IsReady(self):
        return self.day_close and self.day_open

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

    # endregion

    # region State Mgmt

    def UpdateHighLowDay(self, bar):
        tm = bar.EndTime
        h,l = bar.High, bar.Low
        if tm.hour == 9 and tm.minute == 31:
            self.high_of_day = bar.High 
            self.low_of_day = bar.Low 
        elif tm.hour >= 16:
            self.high_of_day = None
            self.low_of_day = None 
        else:
            self.high_of_day = max(h, self.high_of_day)
            self.low_of_day = min(l, self.low_of_day)

    def UpdateGaps(self, bar):
        tm = bar.EndTime 
        # self.algo.Log(f'{self.symbol} -- tm: {tm} -- type: {type(tm)}')
        if tm.hour == self.algo.bod_hr and tm.minute == self.algo.bod_mn:
            self.day_open = bar.Close
        
        if tm.hour == self.algo.closed_hr and tm.minute == self.algo.closed_mn:
            self.day_close = bar.Close
            self.day_open = None # Need to await the next open!


    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

    # endregion

# region imports
from AlgorithmImports import *
# endregion

from Strategy import Strategy

'''
'Best' GF Params:
stop_dec: .005, tgt_dec: .0275 - .0325, gap_thres_pct: 3; (fill=1, dir=1)

'Best' GC Params:
stop_dec: .055, tgt_dec: .035, gap_thres_pct: .75
'''

class GeekyVioletJellyfish(QCAlgorithm):

    tickers = ['NVDA','TSLA','AMD','AMZN','AAPL','MSFT','NFLX'] #,'MARA']

    # How should I format this...
    # config = []
    # want to give EACH symbol their own stop, target, etc?

    resolution = Resolution.Minute
    symbols = []
    strats = {}


    # -------------------- Parameters ------------------- 
    
    entries_per_day = 3
    # Per strategy. 

    # do 30 mins off open, and close.
    bod_hr = 9 
    bod_mn = 59 # GF does better with 31 time.
    closed_hr = 15
    closed_mn = 30

    eod_exit_offset = 30

    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        # self.SetEndDate(2023, 2, 1)
        self.AddEquity('SPY', self.resolution)

        # stop, and target, are in form of decimal -- .05 = 5%
        self.stop_pct = self.GetParameter("stop_dec", .05)
        self.tgt_pct = self.GetParameter("tgt_dec", .1)
        self.thres = self.GetParameter("gap_thres_pct", 0.0) # 1 = 1%, REAL percent.
        self.gap_fill = bool(self.GetParameter("fill", 1))
        self.dir = self.GetParameter("dir", 1)

        self.SetCash(100000)

        for t in self.tickers:
            s = self.AddEquity(t, self.resolution).Symbol
            self.symbols.append(s)
            self.strats[s] = Strategy(self, s)

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose("SPY", self.eod_exit_offset),
                 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)
            if strat.IsReady:
                gap_pcts.append((symbol, strat.GapPct))
        

        # Sort this... (Skip for now)
        # all_ready = len([i for i in self.strats if i.IsReady]) == len(self.strats)
        # desc_gap_pcts = [i for i in sorted(gap_pcts, key=lambda x: x[1], reverse=True)]
        # asc_gap_pcts = [i for i in sorted(gap_pcts, key=lambda x: x[1], reverse=False)]
        # biggest gaps first... 


        # Try 'biggest' gaps -- clone and try that?