from Selection.OptionUniverseSelectionModel import OptionUniverseSelectionModel
from datetime import date, timedelta


class OptionsUniverseSelectionModel(OptionUniverseSelectionModel):
    def __init__(self, select_option_chain_symbols):
        super().__init__(timedelta(1), select_option_chain_symbols)

    def Filter(self, filter):
        # Define options filter -- strikes +/- 3 and expiry between 0 and 180 days away
        return (filteruniverse.IncludeWeeklys()
                .BackMonths()
                .PutsOnly()
                .Strikes(-40, 0)
                .Expiration(timedelta(self.filterStartDate), timedelta(self.filterEndDate))
                )

# ----------------------------------------------------------------------
#
# Custom Buying power model to solve insufficient funds problem. There is a fix coming in December/January
#
# ----------------------------------------------------------------------

class CustomBuyingPowerModel(BuyingPowerModel):
    def GetMaximumOrderQuantityForTargetBuyingPower(self, parameters):
        quantity = super().GetMaximumOrderQuantityForTargetBuyingPower(parameters).Quantity
        quantity = np.floor(quantity / 100) * 100
        return GetMaximumOrderQuantityResult(quantity)

    def HasSufficientBuyingPowerForOrder(self, parameters):
        return HasSufficientBuyingPowerForOrderResult(True)

# Your New Python File

from System import TimeSpan
from System.Drawing import Color
import numpy as np
import pandas as pd
from QuantConnect import Chart, DataNormalizationMode
from QuantConnect.Orders import OrderDirection
from QuantConnect.Securities import BuyingPowerModel
from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Securities.Option import OptionStrategies
from QuantConnect.Data.Custom.CBOE import CBOE
from datetime import timedelta

# lib
from lib import SelectionModel
from lib import CustomBuyingPowerModel
# ----------------------------------------------------------------------
#
# Bull Put Credit Spread Algorithm
#
# ----------------------------------------------------------------------


class OptionsAlgorithm(QCAlgorithm):

    # ----------------------------------------------------------------------
    # Initialize QuantConnect Algorithm
    # ----------------------------------------------------------------------
    def Initialize(self):
        # Base QuantConnect Parameters
        self.SetStartDate(2017, 11, 1)
        self.SetEndDate(2019, 2, 1)
        self.SetCash(100000)

        # Base Algorithm Paramters
        self.investPercent = 0.9
        self.filterStartDate = 25
        self.filterEndDate = 45
        self.shortDelta = 0.25
        self.longDelta = 0.15

        # Helper Variables
        self.netCredit = None
        self.shortOrder = None
        self.longOrder = None
        self.expiry = None
        self.exitDate = None
        self.inPosition = False
        self.openPortfolioValue = None

        # Set Instruments
        option = self.AddOption("SPY")
        option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        option.SetBuyingPowerModel(CustomBuyingPowerModel.CustomBuyingPowerModel())
        self.optionSymbol = option.Symbol
        self.SetUniverseSelection(SelectionModel.OptionsUniverseSelectionModel(self.SelectOptionsSymbols))
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))
        
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.vix = self.AddData(CBOE, "VIX").Symbol
        self.rsi = self.RSI("SPY", 10, MovingAverageType.Simple, Resolution.Daily, Field.Close)

        # Charting
        overlayPlot = Chart("Overlay Plot")
        overlayPlot.AddSeries(Series("RSI", SeriesType.Line, "", Color.Aqua))
        overlayPlot.AddSeries(Series("Over Bought", SeriesType.Line, "", Color.Navy))
        overlayPlot.AddSeries(Series("Over Sold", SeriesType.Line, "", Color.Navy))
        overlayPlot.AddSeries(Series("Mid", SeriesType.Line, "", Color.Navy))
        overlayPlot.AddSeries(Series("Sell", SeriesType.Line, "", Color.Red))
        overlayPlot.AddSeries(Series("Buy", SeriesType.Line, "", Color.Green))
        self.AddChart(overlayPlot)

        # Set warmup for Greeks and RSI
        self.SetWarmUp(TimeSpan.FromDays(30))

        # Check exits everyday
        self.Schedule.On(self.DateRules.EveryDay(
            "SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), self.checkExit)

    # ----------------------------------------------------------------------
    # Primary Data function
    # ----------------------------------------------------------------------
    def OnData(self, slice):
        
        if self.IsWarmingUp:
            return

        if self.Time.hour == 9 and self.Time.minute == 31:

            if self.IsWarmingUp:
                return

            #
            # ISSUE: RSI IS NOT PLOTTING
            #
            self.Plot("Overlay Plot", "RSI", self.rsi.Current.Value)
            self.Plot("Overlay Plot", "Over Bought", 80)
            self.Plot("Overlay Plot", "Over Sold", 20)
            self.Plot("Overlay Plot", "Mid", 50)

            # if self.rsi.Current.Value > 50:
            self.getContracts(slice)

            #if self.inPosition:
            #    self.checkExit
    # ----------------------------------------------------------------------
    # Get Short and Long Put Contracts
    # ----------------------------------------------------------------------

    def getContracts(self, slice):
        #self.Debug('getting Contracts')
        # Define the option Chain
        chain = slice.OptionChains[self.optionSymbol]

        # set float format so delta displays correctly
        pd.set_option('display.float_format', lambda x: '%.5f' % x)

        # Set relevant information in the Dataframe
        df = pd.DataFrame([[x.Right, float(x.Strike), x.Expiry, float(x.BidPrice), float(x.AskPrice), x.Greeks.Delta, x.UnderlyingLastPrice] for x in chain],
                          index=[x.Symbol.Value for x in chain],
                          columns=['type', 'strike', 'expiry', 'askPrice', 'bidPrice', 'delta', 'underlyingLast'])

        # Set the Dataframe to the option contract data frame
        self.dfOptionsContracts = df

        # Create a new column set to absolute value of the values of the delta column - the short delta to determine closest available contract to our specified detla
        self.dfOptionsContracts["shortDeltaDiff"] = np.abs(
            self.dfOptionsContracts["delta"] - self.shortDelta)

        # Create a new column set to absolute value of the values of the delta column - the long delta to determine closest available contract to our specified detla
        self.dfOptionsContracts["longDeltaDiff"] = np.abs(
            self.dfOptionsContracts["delta"] - self.longDelta)

        # Create two separate dateframes, one including the expiry date and the (delta - self.shortDelta ) values, and another including the expiry date and the (delta - self.longDelta) values
        shortDeltaExpiryDf = self.dfOptionsContracts.filter(
            items=["expiry", "shortDeltaDiff"])
        longDeltaExpiryDf = self.dfOptionsContracts.filter(
            items=["expiry", "longDeltaDiff"])

        # Create three different dataframes from the shortDeltaExpiryDf and longDeltaExpiryDf:
        # 1. A dateframe that combines the (delta - self.shortDelta/self.longDelta) values on the expiry dates so that they are matched only on the same date
        # to prevent use of contracts expiring on different dates.
        # 2. A dateframe from the dateframe in Step 1 that only contains the rows that contain the minimum value of (delta - self.shortDelta).
        # 3. A dateframe that sorts the values of dataframe in Step 2 so that the data is sorted with the (delta - self.longDelta) values from least to greatest.
        ##
        shortLongDeltaCombinedOnExpiry = pd.merge(left=shortDeltaExpiryDf, right=longDeltaExpiryDf, left_on='expiry', right_on='expiry')
        combinedExpiryOnlyShortDeltaMinRows = shortLongDeltaCombinedOnExpiry[shortLongDeltaCombinedOnExpiry.shortDeltaDiff == shortLongDeltaCombinedOnExpiry.shortDeltaDiff.min()]
        combinedExpirySortedByLongAndShortMin = combinedExpiryOnlyShortDeltaMinRows.sort_values(by=['shortDeltaDiff', 'longDeltaDiff'])

        # Get the contract ids for short and long contracts (eg. SPY   190227P00261000)
        shortContract = self.dfOptionsContracts[(self.dfOptionsContracts.expiry == combinedExpirySortedByLongAndShortMin.expiry.iloc[0]) & (
            self.dfOptionsContracts.shortDeltaDiff == combinedExpirySortedByLongAndShortMin.shortDeltaDiff.iloc[0])].index[0]
        longContract = self.dfOptionsContracts[(self.dfOptionsContracts.expiry == combinedExpirySortedByLongAndShortMin.expiry.iloc[0]) & (
            self.dfOptionsContracts.longDeltaDiff == combinedExpirySortedByLongAndShortMin.longDeltaDiff.iloc[0])].index[0]

        # Create a variable that stores a boolean for whether the two  contracts are equal.
        areTheySame = shortContract == longContract

        # Create an iterator value for the longContract
        longContractIterator = 0

        # Iterate over the long contracts until it is not equal to the short contract. Usually just the next row of dateframe: combinedExpirySortedByLongAndShortMin
        while areTheySame:
            if longContractIterator not in combinedExpirySortedByLongAndShortMin.expiry.index:
                self.Log(f"longContractIterator: {longContractIterator} combinedExpirySortedByLongAndShortMin.expiry:\n{combinedExpirySortedByLongAndShortMin.expiry}")
                self.Quit("Out of bounds")
            
            longContract = self.dfOptionsContracts[(self.dfOptionsContracts.expiry == combinedExpirySortedByLongAndShortMin.expiry.iloc[longContractIterator]) & (
                self.dfOptionsContracts.longDeltaDiff == combinedExpirySortedByLongAndShortMin.longDeltaDiff.iloc[longContractIterator])].index[0]
            longContractIterator += 1
            areTheySame = shortContract == longContract

        # self.Debug(f'SHORT: {shortContract.index[0]} and LONG: {longContract.index[0]} and are they EQUAL: {areTheySame} SHORT EXPIRY:{combinedExpirySortedByLongAndShortMin.expiry.iloc[0]} LONG EXPIRY: {combinedExpirySortedByLongAndShortMin.expiry.iloc[longContractIterator]}')

        # Store the additonal contract information (eg. bidPrice, askPrice, delta, underlyingPrice)
        # shortContractInfo = self.dfOptionsContracts.loc[shortContract]
        # longContractInfo = self.dfOptionsContracts.loc[longContract]

        # Log out what our contracts are:
        #self.Debug(
        #    f"Underlying: {shortContractInfo.underlyingLast} Strike: {shortContractInfo.strike} Expiry: {shortContractInfo.expiry} Delta: {shortContractInfo.delta}")
        #self.Debug(
        #    f"Underlying: {longContractInfo.underlyingLast} Strike: {longContractInfo.strike} Expiry: {longContractInfo.expiry} Delta: {longContractInfo.delta}")

        # self.placeOrder(shortContract, shortContractInfo,
        #                 longContract, longContractInfo)

    # ----------------------------------------------------------------------
    # Order Functions
    # ----------------------------------------------------------------------

    def placeOrder(self, shortContract, shortContractInfo, longContract, longContractInfo):

        # Get our margin
        margin = self.Portfolio.GetBuyingPower(
            shortContract, OrderDirection.Sell)

        # Get the quantities
        qty = margin * self.investPercent / \
            ((shortContractInfo.bidPrice + longContractInfo.askPrice) * 100)

        # Check that contracts are not the same
        if qty < 1:
            return
        else:

            self.Sell(OptionStrategies.BullPutSpread(self.optionSymbol,
                                                     shortContractInfo.strike, longContractInfo.strike, shortContractInfo.expiry), np.floor(qty))

            self.Plot("Overlay Plot", "Buy", self.rsi.Current.Value)

            # Set in position as true so we don't continue buying
            self.inPosition = True

            # Store the net Credit
            self.netCredit = (np.abs(shortContractInfo.askPrice) -
                              np.abs(longContractInfo.bidPrice)) * np.abs(qty) * 100

            # Generate last trading days
            self.expiry = shortContractInfo.expiry
            startDate = self.expiry + timedelta(days=-7)
            endDate = self.expiry + timedelta(days=-1)
            self.exitDate = self.TradingCalendar.GetTradingDays(
                startDate, endDate)

            # Set the openPortfolioValue for Profit Calculations
            self.openPortfolioValue = self.Portfolio.TotalPortfolioValue

    # ----------------------------------------------------------------------
    # Check Exit
    # ----------------------------------------------------------------------
    def checkExit(self):

        # store portfolio change
        if self.openPortfolioValue is not None:

            change = self.Portfolio.TotalPortfolioValue - self.openPortfolioValue

            # Exit at 70% Profit
            if (change / self.netCredit > .7):
                self.liquidate()
                #self.Debug('Liquidating postion because we reached 70% profit')

            # Exit at 20% loss
            # if (change / self.netCredit < -.2):
            #     self.liquidate()
            #     self.Debug('Liquidating postion 20 percent stop loss')

    # ----------------------------------------------------------------------
    # Liquidate
    # ----------------------------------------------------------------------
    def liquidate(self):
        self.Liquidate()
        self.Plot("Overlay Plot", "Sell", self.rsi.Current.Value)
        self.inPosition = False

    # ----------------------------------------------------------------------
    # Define Options universe
    # ----------------------------------------------------------------------
    def SelectOptionsSymbols(self, utcTime):
        ticker = self.optionSymbol
        return [Symbol.Create(ticker, SecurityType.Option, Market.USA, f"?{ticker}")]

    # ----------------------------------------------------------------------
    # Helper Functions
    # ----------------------------------------------------------------------
    def nearest(self, array, value):
        array = np.asarray(array)
        idx = (np.abs(array - value)).argmin()
        return array[idx]