from AlgorithmImports import *

class BrokerageModelAlgorithm(QCAlgorithm):
    def Initialize(self):

        self.SetCash(10000000)            # Set Strategy Cash
        self.SetStartDate(2013,10,20)    # Set Start Date
        self.SetEndDate(2013,11,10)     # Set End Date
        self.AddEquity("SPY", Resolution.Daily)
        
        self.SetBrokerageModel(RebateBalanceBrokerageModel(self, 1000000, 100))
        self.last = 1   
        
    def OnData(self, slice):
        # Simple buy and hold template
        if not self.Portfolio.Invested:
            self.SetHoldings("SPY", self.last)
        
        # should see a rebate on first trade day of November for testing
        if self.Time.month == 11:
            self.Liquidate()

# example: if certain trading volume is reached, rebate set amount of USD reward
class RebateBalanceBrokerageModel(DefaultBrokerageModel):
    def __init__(self, algorithm, rebateThreshold, rebateAmount):
        self.algorithm = algorithm
        self.rebateThreshold = rebateThreshold
        self.rebateAmount = rebateAmount
        
        # say the rebate volume is counted every month, so we set a timestamp for taring
        self.monthEnd = Expiry.EndOfMonth(algorithm.Time)
        # volume counter
        self.volumeTraded = 0
        # check if rebated yet
        self.rebated = False
    
    def CanSubmitOrder(self, security, order, message):
        # check if need to rebalance volume count and move timestamp
        if self.algorithm.Time > self.monthEnd:
            self.volumeTraded = 0
            self.monthEnd = Expiry.EndOfMonth(self.algorithm.Time)
            self.rebated = False
        
        # add up the trading volume
        self.volumeTraded += order.GetValue(security)
        
        # if threhold reached, rebate the rebate the reward amount
        if not self.rebated and self.volumeTraded >= self.rebateThreshold:
            self.algorithm.Portfolio.CashBook["USD"].AddAmount(self.rebateAmount)
            self.rebated = True
            # message
            self.algorithm.Debug(f"{self.algorithm.Time.date()} Reaching the rebate threshold, rebate ${self.rebateAmount}")
            
        return True