# region imports
from AlgorithmImports import *
# endregion
from System.Drawing import Color
from datetime import timedelta
from collections import deque

class DancingBrownCow(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2021, 12, 18)
        self.set_end_date(2022, 1, 1)
        self.set_cash(100000)
        self.qqq = self.add_equity("QQQ", Resolution.Minute)
        self.qqq.set_data_normalization_mode(mode = DataNormalizationMode.RAW) # Options only support raw data
        self.symbol = self.qqq.symbol

        # Consolidate data
        self.consolidate(self.symbol, timedelta(days=1), self.consolidation_handler)

        # Set brokerage model
        self.set_brokerage_model(BrokerageName.QuantConnectBrokerage, AccountType.MARGIN)
        #self.SetBrokerageModel(BrokerageName.TradierBrokerage, AccountType.Margin)
        self.set_benchmark(self.symbol)
    
        self.filtered_contracts = []
        self.implied_volatiliy = []
        self.contract_added = set()
        self.put_contract = None
        self.call_contract = None
        self.call_contract_order = None
        self.put_contract_order = None
        self.days_before_expiration = datetime.min
        self.iv_scaling = 1.5 # How much we need to scale IV
        
        # Schedule events
        self.schedule.on(date_rule = self.schedule.date_rules.every_day(self.symbol), 
                        time_rule = self.schedule.time_rules.before_market_open(self.symbol, 30), 
                        callback = self.add_contracts)
        
        self.schedule.on(date_rule = self.schedule.date_rules.every_day(self.symbol), 
                        time_rule = self.schedule.time_rules.after_market_close(self.symbol, 30), 
                        callback = self.clean_contracts)

        
        # Import the necessary module before using Custom color
        stockPlot = Chart('Trade Plot')
        stockPlot.add_series(CandlestickSeries('Price', index = 0, unit = '$'))
        #stockPlot.add_series(Series('Close Price', SeriesType.Line, '$', Color.Orange))
        stockPlot.add_series(Series('Upper-Band', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.Triangle))
        stockPlot.add_series(Series('Lower-Band', SeriesType.Scatter, '$', Color.Blue, ScatterMarkerSymbol.TriangleDown))
        self.AddChart(stockPlot)

    # Define the consolidation handler
    def consolidation_handler(self, consolidated_bar: TradeBar):

        # Plot candlestick and plot debugging data
        price = consolidated_bar
        self.Plot('Trade Plot', 'Price', open = price.open, high = price.high, low = price.low, close = price.close)

        # Print whether close price within the bar
        if self.call_contract and self.put_contract:
            closed_inside = (self.call_contract[0].ID.strike_price > price.close and self.put_contract[0].ID.strike_price < price.close)
            self.log(f"Open: {price.open}, Close: {price.close}, Upper: {self.call_contract[0].ID.strike_price}, Lower: {self.put_contract[0].ID.strike_price}, Success: {closed_inside}")

        # Populate upper and lower bands
        if self.call_contract:
            self.plot('Trade Plot', 'Upper-Band', self.call_contract[0].ID.strike_price)

        if self.put_contract:
            self.plot('Trade Plot', 'Lower-Band', self.put_contract[0].ID.strike_price)

    def clean_contracts(self):
        # Close all pending contracts
        if self.call_contract:
            self.Transactions.CancelOpenOrders(self.call_contract[0])
        if self.put_contract:
            self.Transactions.CancelOpenOrders(self.put_contract[0])

        # First clean contracts. We are populating it every day
        self.filtered_contracts = []
        self.implied_volatiliy = []
        self.contract_added = set()
        self.put_contract_order = None
        self.call_contract_order = None

    def add_contracts(self):
        # Rest contracts here, since we would like to use them in consolidated data
        self.put_contract = None
        self.call_contract = None

        # Add option chain contracts using todays date
        # One disadvantage is option chain providers don't have greeks
        contracts = self.option_chain_provider.get_option_contract_list(self.symbol, date = self.time)

        # Get last known price
        self.underlying_price = self.get_last_known_price(self.securities[self.symbol]).price

        # Get contracts that expire same day
        closest_contracts = sorted([p for p in contracts if (p.ID.date - self.time).days == 0], key = lambda x: x.ID.date)
        
        # Filter contract based on price closenest
        closest_contracts = [(p, abs(self.underlying_price - p.ID.strike_price)) 
                                    for p in closest_contracts if abs(self.underlying_price - p.ID.strike_price) < self.underlying_price * 0.05]

        # Sort based on closenest
        closest_contracts = sorted(closest_contracts, key = lambda x: x[1])

        # Get contract values onlu
        self.filtered_contracts = [c[0] for c in closest_contracts]
        
        # Now add this data to contract
        for contract in self.filtered_contracts:
            if contract not in self.contract_added:
                self.contract_added.add(contract)
                option = self.add_option_contract(contract)
                
                # Set early assignment model to Null
                #option.set_option_assignment_model(NullOptionAssignmentModel())

    def on_data(self, data: Slice):

        # This has to be here otherwise data is not available
        if self.symbol not in data or (self.symbol in data and data[self.symbol] is None):
            #self.log(self.Time.hour)
            #self.log(data[self.symbol])
            return

        # Sell underlying stock if it has been assigned
        if self.portfolio[self.symbol].invested:
            self.liquidate(self.symbol)

        # First get undelying security price. For each bar
        curr_price = data[self.symbol].price

        # Check if we need to exit position
        if self.call_contract and self.portfolio[self.call_contract[0]].invested:
            # Check PL of this contract
            # if self.Securities[self.call_contract[0]].Holdings.UnrealizedProfitPercent < -1: #If 100% loss then exit
            if curr_price > self.call_contract[0].ID.strike_price:
                self.liquidate(self.call_contract[0])
                self.log(f"Close position on {self.call_contract[0].value}")

        if self.put_contract and self.portfolio[self.put_contract[0]].invested:
            #if self.securities[self.put_contract[0]].Holdings.UnrealizedProfitPercent < -1:
            if curr_price < self.put_contract[0].ID.strike_price:
                self.liquidate(self.put_contract[0])
                self.log(f"Close position on {self.put_contract[0].value}")


        # We found contracts on this day
        if len(self.filtered_contracts) > 0:

            if self.time.hour * 60 + self.time.minute >= 10 * 60:
                average_IV = 8.8827

                # Sell put contract
                if self.put_contract is None:
                    # Sort put options that are closest to curr_price + iv
                    put_contracts = sorted([x for x in self.filtered_contracts 
                                                        if (curr_price - average_IV - x.ID.strike_price) > 0 
                                                        and x.ID.option_right == OptionRight.PUT], 
                                                        key = lambda x: abs(curr_price - average_IV - x.ID.strike_price))

                    if len(put_contracts) > 0:
                        self.put_contract = put_contracts

                else:
                    
                    if self.put_contract and self.put_contract_order is None and not self.portfolio[self.put_contract[0]].invested:
                        # Get last ask price and set limit order
                        put_ask = self.Securities[self.put_contract[0]].AskPrice
                        put_bid = self.Securities[self.put_contract[0]].BidPrice
                        put_price = np.round((put_ask + put_bid) * 0.5, 2)

                        # Set limit order
                        self.put_contract_order = self.limit_order(symbol = self.put_contract[0], 
                                                                    quantity = -1, 
                                                                    limit_price = put_price)


                # Check if we have not sold any options yet
                if self.call_contract is None:
                    
                    # Sort call options that are closest to current price
                    call_contracts = sorted([x for x in self.filtered_contracts 
                                                        if (x.ID.strike_price - curr_price - average_IV) > 0 
                                                        and x.ID.option_right == OptionRight.CALL], 
                                                        key = lambda x: (x.ID.strike_price - curr_price - average_IV))

                    if len(call_contracts) > 0:
                        self.call_contract = call_contracts


                # Now sell option contract
                else:
                    
                    if self.call_contract and self.call_contract_order is None and not self.portfolio[self.call_contract[0]].invested:
                        # Get last ask price and set limit order
                        call_ask = self.Securities[self.call_contract[0]].AskPrice
                        call_bid = self.Securities[self.call_contract[0]].BidPrice
                        call_price = np.round((call_ask + call_bid) * 0.5, 2)

                        # Set limit sell
                        self.call_contract_order = self.limit_order(symbol = self.call_contract[0], 
                                                                    quantity = -1, 
                                                                    limit_price = call_price)