from clr import AddReference
AddReference("System.Core")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")

from QuantConnect.Securities.Option import OptionPriceModels
import QuantConnect
import cloudpickle
import json
import math
import numpy as np
import pandas as pd
from scipy.optimize import minimize_scalar
from scipy.stats import linregress

class AutomatonTrade(QCAlgorithm):

    # Constants
    NUMBER_SYMBOLS = 200
    ATR_PERIOD = 15
    SHORT_ATR_PERIOD = 5
    ATR_MULTIPLIER = 1
    BAD_SYMBOLS = [ "|", " ", "ZVZZT", "ZXZZT", "ZWZZT", "TBLT", "MTECW", "SSLJ", "YERR", "CYTXZ", "BRACW", "NFINW", "MDR", "AIKI", "CEI" ]
    WINDOW_SIZE = 15
    BARS_TO_HOLD_ORDERS_FOR = 5
    BID_ADJUSTMENT = 0.01

    def Initialize(self):
        
        # Init
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 1, 3)
        self.SetCash(5000)
        
        self.SetTimeZone(TimeZones.NewYork)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
                
        # Live universe with selection.
        self.AddUniverse(self.CoarseSelectionFunction)

    def CustomSecurityInitializer(self, security):
        '''Initialize the security with raw prices'''
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)

    def CoarseSelectionFunction(self, coarse):

        # Sort by price, dollar volume and volume
        sortedByPrice = sorted(coarse, key = lambda x : (x.Price, x.DollarVolume, x.Volume))
    
        # Volume > 500k so we have liquidity and don't get shitty stocks
        # Any ticker with more than 4 chars is likely a warrant and will crash
        # Eliminate bad symbols
        filtered = [ x.Symbol for x in sortedByPrice if x.Volume > 500000 and len(x.Symbol.Value) < 5 and not any(y in str(x.Symbol.Value) for y in AutomatonTrade.BAD_SYMBOLS) ]

        top = filtered[:AutomatonTrade.NUMBER_SYMBOLS]
        
        self.currentActiveSecurities = top
        self.lastSelectedSymbolsAsCommaSeperatedString = ', '.join(list(map(lambda x : '"' + x.Value + '"', top)))
        self.printLog("------ CoarseSelectionFunction symbols (" + str(len(top)) + "):")
        #self.printLog(self.lastSelectedSymbolsAsCommaSeperatedString)

        # Reset vars
        self.prices = np.zeros([AutomatonTrade.WINDOW_SIZE, AutomatonTrade.NUMBER_SYMBOLS])
        self.lastSelectedSymbolsAsCommaSeperatedString = None
        self.ignoreDaySymbols = []
        self.barCount = 0

        # Store the list as the tickers
        self.lastSelectedSymbols = list(map(lambda x : x.Value, top))
        
        return top
        
    def OnData(self, dataSlice):
        
        # We use the bars (stock) for everything    
        data = dataSlice.Bars

        # Check if we have open orders
        openOrders = self.Transactions.GetOpenOrders()
        if len(openOrders) == 0:
            
            # If nothing is invested, look to buy
            if not self.Portfolio.Invested:

                # See if we can invest in anything
                for i, stock in enumerate(self.currentActiveSecurities):
                    
                    stockSymbol = stock
                    contracts = self.OptionChainProvider.GetOptionContractList(stockSymbol, data.Time)
                    for contract in contracts:
                        option = self.AddOptionContract(contract, Resolution.Minute)
                        option.PriceModel = OptionPriceModels.BlackScholes()
                    
                    optionchain = None
                    for kvp in dataSlice.OptionChains:
                        if str(stockSymbol) in str(kvp.Key):
                            optionchain = kvp.Value
                            break
                    
                    # Ensure we have data for it
                    if data.ContainsKey(stockSymbol) and data[stockSymbol] is not None and optionchain is not None:
                        
                        for contract in optionchain.Contracts:
                            if contract.Value.Greeks.Delta != 0:
                                self.printLog(str(contract.Value.Greeks.Delta))
                                
    def printLog(self, message, symbol = None):
        if symbol is not None:
            message = self.prependMessage(message, symbol)

        message = str(self.Time) + " -- " + message

        if not self.LiveMode:
            self.Debug(message) # This ends up in the log too
        else:
            self.Log(message)