import numpy as np
from time import sleep


class CalibratedParticleReplicator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)  # Set Start
        #self.SetEndDate(2010, 7, 1) # Set end
        self.SetCash(100000)  # Set Strategy Cash
        self.SetWarmUp(timedelta(2))
        #self.SetBrokerageModel(BrokerageName.Alpaca, AccountType.Margin)
        
        # Ensure that we have data available
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddUniverse(
            self.SelectCoarse, self.SelectFine
        )
        
        self.MRKT = self.AddEquity('SPY', Resolution.Minute).Symbol  # market
        #self.BND = self.AddEquity('TLT', Resolution.Minute).Symbol # bonds
        
        self.stop_limits = {}
        
        # Max leverage
        self.max_lever = 1.0
        self.max_lever_per_sec = 0.1
    
        # Market
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY", 5),
                 self.EveryDayAfterMarketOpen)
    
        # schedule an event to fire every trading day for a security the
        # time rule here tells it to fire 10 minutes before SPY's market close
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                         self.TimeRules.BeforeMarketClose("SPY", 10),
                         self.EveryDayBeforeMarketClose)
                         
        self.SetBenchmark("SPY")

    def GetCalls(self, symbol, start_days = 30, end_days = 60):

        # Get price
        hist = self.History(symbol, 1, Resolution.Minute)

        #pulls contract data for select equity at current time
        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)

        #sorts contracts by closet expiring date date and closest strike price (sorts in ascending order)
        calls = sorted(sorted([x for x in contracts if x.ID.OptionRight == OptionRight.Call], key = lambda x: x.ID.Date), key = lambda x: x.ID.StrikePrice)

        calls = [x for x in calls if x.ID.StrikePrice > hist.iloc[-1]['close']]

        #then selects all contracts that meet our expiration criteria
        calls = [x for x in calls if  start_days<(x.ID.Date - self.Time).days <= end_days]

        if not calls:
           return []

        #will eventually return array of optimal puts and calls
        return calls

    def add_overnight_returns(self, securities):
        new_securities = []
        for sec in securities:
            try:
                hist = self.History(sec.Symbol, 2, Resolution.Daily)
                #if len(hist.index) != 2:
                #    self.Debug("unable to get proper history for {}, got {}".format(sec, hist.index))
                overnight = (hist.iloc[-1]['open'] / hist.iloc[-2]['close']) - 1
                sec.overnight_returns = overnight
                new_securities.append(sec)
            except IndexError as ex:
                pass
                #self.Debug("could not analyze {}: {}".format(sec, ex))
        return new_securities

    def SelectCoarse(self, coarse):
        dollarvolume_array = np.array([ x.DollarVolume for x in coarse ])
        dollarvolume_percentile = np.percentile(dollarvolume_array, 90)
        # Select top volume
        coarse = [x for x in coarse if x.DollarVolume >= dollarvolume_percentile]
        # Low price
        price_array = np.array([ x.Price for x in coarse ])
        price_percentile = np.percentile(price_array, 10)
        coarse = [x.Symbol for x in coarse if x.Price <= price_percentile]
        return coarse

    def SelectFine(self, fine):
        # Overnight returns
        selected = self.add_overnight_returns(fine)
        overnight_returns = np.array([ s.overnight_returns for s in selected ])
        mask = np.all(np.isnan(overnight_returns), axis=0)
        if not len(overnight_returns[~mask] > 0):
            return []
        overnight_returns_pct_high = np.percentile(overnight_returns[~mask], 90)
        overnight_returns_pct_low = np.percentile(overnight_returns[~mask], 10)
        if not selected:
            return selected
        self.shorts = [f.Symbol for f in sorted(selected, key=lambda x: x.overnight_returns, reverse=True) if f.overnight_returns >= overnight_returns_pct_high][:5]
        self.longs = [f.Symbol for f in sorted(selected, key=lambda x: x.overnight_returns) if f.overnight_returns <= overnight_returns_pct_low][:5]
        self.symbols = self.longs+self.shorts
        return self.symbols

    def EveryDayAfterMarketOpen(self):
        if self.IsWarmingUp:
            return
        if not self.symbols:
            return
        #self.Debug("got {} from fine filter".format(len(self.symbols)))
        self.lever_per_sec = min(self.max_lever / len(self.symbols), self.max_lever_per_sec)
        # Set excess to SPY/TLT
        self.leftover = 1.0 - len(self.symbols)*self.lever_per_sec
        # SPY
        self.SetHoldings(self.MRKT, self.leftover)
        # BNDs
        #self.SetHoldings(self.BND, self.leftover)
        # Shorts
        for symbol in self.shorts:
            # Get call options
            #calls = self.GetCalls(symbol)
            self.SetHoldings(symbol, -self.lever_per_sec)
            self.stop_limits[symbol] = {
                "stop": self.Securities[symbol].Price*(1.2),
                "limit": self.Securities[symbol].Price*(0.5),
            }
        # Longs
        for symbol in self.longs:
            # Get call options
            #calls = self.GetCalls(symbol)
            self.SetHoldings(symbol, self.lever_per_sec)
            self.stop_limits[symbol] = {
                "limit": self.Securities[symbol].Price*(1.5),
                "stop": self.Securities[symbol].Price*(0.8),
            }
            #if len(calls) > 0:
            #    try:
            #        call_option = calls[0]
            #        self.AddOptionContract(call_option, Resolution.Minute)
                    #self.Debug("buying call contract: {}".format(calls[0]))
            #        call = self.MarketOrder(call_option, 1, True)
            #        self.stop_limits[symbol]["option"] = call
            #    except Exception as ex:
            #        pass
                    #self.Debug("error trying to order call contract for {}".format(symbol))
                    #self.Debug(ex)
            #else:
            #    pass
                #self.Debug("no call contracts found for {}".format(symbol))
        self.Plot("Holdings", "number", len(self.Securities.keys()))
        self.Plot("Allocation", "mrkt", self.leftover)
        self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
        self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)

    def EveryDayBeforeMarketClose(self):
        if self.IsWarmingUp:
            return
        self.Plot("Holdings", "number", len(self.Securities.keys()))
        self.Plot("Allocation", "mrkt", self.leftover)
        self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
        self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)
        self.Liquidate()

    def OnSecuritiesChanged(self, changes):
        if self.IsWarmingUp:
            return
        # Subscribe to minute bars for added equities
        for instrument in changes.AddedSecurities:
            if instrument.Symbol.SecurityType == SecurityType.Equity:
                self.AddEquity(instrument.Symbol, Resolution.Minute)
                self.Securities[instrument.Symbol].FeeModel = ConstantFeeModel(0.0)
                self.Securities[instrument.Symbol].MarginModel = PatternDayTradingMarginModel()
                #.Securities[instrument.Symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
        # Remove equities
        for instrument in changes.RemovedSecurities:
            if instrument.Symbol.SecurityType == SecurityType.Equity:
                self.RemoveSecurity(instrument.Symbol)
                #if "option" in self.stop_limits[instrument.Symbol].keys():
                #    call = self.stop_limits[instrument.Symbol]["option"]
                #    self.RemoveSecurity(call.Symbol)
                
    
    def OnData(self, data):
        if self.IsWarmingUp:
            return
        # Limits and stops
        for symbol, security in self.Portfolio.items():
            if symbol == self.MRKT:
                continue
            if not security.Invested:
                continue
            if not data.Bars.ContainsKey(symbol):
                continue
            if not symbol.SecurityType == SecurityType.Equity:
                continue
            last_price = data.Bars[symbol].Close
            stop_limit = self.stop_limits[symbol]
            limit_price = stop_limit["limit"]
            stop_price = stop_limit["stop"]
            # Handle shorts
            if symbol in self.shorts:
                if last_price >= stop_price:
                    #self.Debug("{} triggered stop at {}, liquidating".format(symbol, last_price))
                    self.Liquidate(symbol)
                    self.Plot("Holdings", "number", len(self.Securities.keys()))
                    self.Plot("Allocation", "mrkt", self.leftover)
                    self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
                    self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)
                elif last_price <= limit_price:
                    #self.Debug("{} triggered limit at {}, liquidating".format(symbol, last_price))
                    self.Liquidate(symbol)
                    self.Plot("Holdings", "number", len(self.Securities.keys()))
                    self.Plot("Allocation", "mrkt", self.leftover)
                    self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
                    self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)
                    #if "option" in stop_limit.keys():
                    #    call = stop_limit["option"]
                    #    #self.Debug("liquidating call contract {}".format(call.Symbol, last_price))
                    #    self.Liquidate(call.Symbol)
            # Longs
            elif symbol in self.longs:
                if last_price >= limit_price:
                    #self.Debug("{} triggered stop at {}, liquidating".format(symbol, last_price))
                    self.Liquidate(symbol)
                    self.Plot("Holdings", "number", len(self.Securities.keys()))
                    self.Plot("Allocation", "mrkt", self.leftover)
                    self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
                    self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)
                elif last_price <= stop_price:
                    #self.Debug("{} triggered limit at {}, liquidating".format(symbol, last_price))
                    self.Liquidate(symbol)
                    self.Plot("Holdings", "number", len(self.Securities.keys()))
                    self.Plot("Allocation", "mrkt", self.leftover)
                    self.Plot("Allocation", "long", len(self.longs)*self.lever_per_sec)
                    self.Plot("Allocation", "short", len(self.shorts)*self.lever_per_sec)
                    #if "option" in stop_limit.keys():
                    #    call = stop_limit["option"]
                    #    #self.Debug("liquidating call contract {}".format(call.Symbol, last_price))
                    #    self.Liquidate(call.Symbol)