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
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(2018, 1, 1)
        self.SetEndDate(2018, 5, 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.option_symbol = 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.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):
        shortContract = None
        longContract = None

        #
        # Get Greeks --> Should be moved to a helper function
        # ISSUE: GREEKS ARE ALWAYS 0.40+ NOT OUR GREEKS. WE NEED TO FIND A BETTER WAY TO FILTER.
        #
        for i in slice.OptionChains:
            chain = i.Value
            contracts = [x for x in chain]
            contracts = sorted(contracts, key=lambda x: (x.Expiry, x.Strike, x.Greeks.Delta))
            shortContract = min(contracts, key=lambda x: abs(x.Greeks.Delta-self.shortDelta))
            longContract = min(contracts, key=lambda x: abs(x.Strike))
        self.Debug(f'Underlying: {shortContract.UnderlyingLastPrice}')
        self.Debug(f"shortcontract: Strike: {shortContract.Strike} Expiry: {shortContract.Expiry} Delta: {shortContract.Greeks.Delta}")
        self.Debug(f"longContract: Strike: {longContract.Strike} Expiry: {longContract.Expiry} Delta: {longContract.Greeks.Delta}")
        # Check that contracts are not the same
        if shortContract.Strike != longContract.Strike:
            self.placeOrder(shortContract, longContract)

    # ----------------------------------------------------------------------
    # Order Functions
    # ----------------------------------------------------------------------
    def placeOrder(self, shortContract, longContract):

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

        # Get the quantities
        qty = margin * self.investPercent / \
            ((shortContract.BidPrice + longContract.AskPrice) * 100)

        # Check that contracts are not the same
        if qty < 1:
            return
        else:
            # Log out what our contracts are:
            self.Debug(f"shortcontract: Strike: {shortContract.Strike} Expiry: {shortContract.Expiry} Delta: {shortContract.Greeks.Delta}")
            self.Debug(f"longContract: Strike: {longContract.Strike} Expiry: {longContract.Expiry} Delta: {longContract.Greeks.Delta}")

            self.Sell(OptionStrategies.BullPutSpread(self.option_symbol,
                                                     shortContract.Strike, longContract.Strike, shortContract.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(shortContract.AskPrice) -
                              np.abs(longContract.BidPrice)) * np.abs(qty) * 100

            # Generate last trading days
            self.expiry = shortContract.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.option_symbol
        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]