# ----------------------------------------------------------------------
#
# 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 datetime import timedelta
from QuantConnect.Orders import OrderDirection
import numpy as np
import re


class OrderManagement:

    def __init__(self, algorithm):
        self.algo = algorithm

    def open_hedge(self):

        slice = self.algo.CurrentSlice

        for i in slice.OptionChains:
            chain = i.Value
            # spot_price = chain.Underlying.Price
            contracts = [x for x in chain if x.Right == 1]
            contracts = sorted(contracts, key=lambda x: (x.Expiry, x.Greeks.Delta))
            self.algo.hedge["contract"] = min(contracts, key=lambda x: abs(x.Greeks.Delta-(-0.30)))

        if self.algo.hedge["contract"] is None:
            return

        qty = (self.algo.Portfolio.Cash * 0.60) / (self.algo.hedge["contract"].AskPrice * 100)

        if qty < 1:
            return
        else:
            self.algo.Debug(f"H - S {self.algo.hedge['contract'].Strike} Ex {self.algo.hedge['contract'].Expiry} D{round(self.algo.hedge['contract'].Greeks.Delta, 3)}")

            self.algo.longHedgeOrder = self.algo.Buy(self.algo.hedge["contract"].Symbol,  np.floor(qty))
            self.algo.hedge["position"] = True
            self.algo.hedge["quantity"] = self.algo.longHedgeOrder.Quantity

    def open_spread(self, shortContract, longContract):

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

        # Get the quantities
        qty = margin * self.algo.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.algo.Debug(f"S - Expiry: {shortContract.Expiry} Delta: {round(shortContract.Greeks.Delta, 6)} macd:{round(self.algo.macd.Signal.Current.Value, 2)}")
            self.algo.Debug(f"L - Expiry: {longContract.Expiry} Delta: {round(longContract.Greeks.Delta, 6)}")
            
            # Perform the order
            self.algo.longOrder = self.algo.Buy(longContract.Symbol,  np.floor(qty))
            self.algo.shortOrder = self.algo.Sell(shortContract.Symbol,  np.floor(qty))

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

            # Store the net Credit
            self.algo.netCredit = (
                np.abs(self.algo.shortOrder.AverageFillPrice) - np.abs(self.algo.longOrder.AverageFillPrice)) * np.abs(self.algo.longOrder.QuantityFilled) * 100

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

            # store expiry
            self.algo.spread_expiry = shortContract.Expiry

    def close_all(self):
        # kill all positions
        self.algo.Liquidate()
        self.algo.inPosition = False

class PositionManagement:

    def __init__(self, algorithm):
        self.algo = algorithm
        self.previousGainCategory = 0

    def hedge_handler(self):
        # helper variables
        unreal = self.algo.Portfolio[self.algo.hedge["contract"].Symbol].UnrealizedProfit
        cost = self.algo.Portfolio[self.algo.hedge["contract"].Symbol].HoldingsCost
        qty = int(round(self.algo.hedge["quantity"] / 4))
            
        # wait for the gain categories to be set
        if self.algo.hedge["gainCategory"] == None:
            return
        
        # its selling more than we have.
        if self.algo.Portfolio[self.algo.hedge["contract"].Symbol].Quantity > 0:

            if self.algo.hedge["gainCategory"] == -1:
                self.algo.Liquidate()
                reset(self)
                
            elif self.algo.hedge["gainCategory"] == self.previousGainCategory:
                return

            elif self.algo.hedge["gainCategory"] > self.previousGainCategory:
                
                # handler any remainder quantities.
                if self.algo.Portfolio[self.algo.hedge["contract"].Symbol].Quantity < qty:
                    qty = self.algo.Portfolio[self.algo.hedge["contract"].Symbol].Quantity
                    
                self.algo.MarketOrder(self.algo.hedge["contract"].Symbol, -qty)
                
                self.algo.Debug(f'Hedge Sell | gain%: {safe_div(unreal, cost)} | scale:{self.algo.hedge["gainCategory"]} @ {self.algo.Time}')
        
                # set the previous gain category to the current catehory
                self.previousGainCategory = self.algo.hedge["gainCategory"]
            
        else:
            self.algo.Liquidate()
            reset(self)


def safe_div(x, y):
    if x == 0:
        return 0
    return round(x / y, 2)


def reset(self):
    # reset our hedge as we have none left
    self.previousGainCategory = 0
    self.algo.hedge["position"] = False
    self.algo.hedge["profile"] = 1
    self.algo.hedge["gainCategory"] = None
    self.algo.hedge["contract"] = None
    self.algo.hedge["scale"] = 0
    self.algo.hedge["stopLoss"] = 0
    self.algo.hedge["quantity"] = 0

from signals import *
from position import *
from order import *
from lib import CustomBuyingPowerModel
from datetime import timedelta
from QuantConnect.Data.Custom.CBOE import CBOE
from QuantConnect.Securities.Option import OptionPriceModels
from QuantConnect.Indicators import RollingWindow
from QuantConnect.Data.Market import TradeBar
from QuantConnect.Algorithm import *
from QuantConnect import *
from System import *
from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")


class OptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        """
        Initializes the algorithm

        Args:
            self: write your description
        """
        # Base QuantConnect Parameters
        self.SetStartDate(2018, 1, 1)
        self.SetEndDate(2019, 1, 1)
        self.SetCash(150000)

        # Base Algorithm Paramters                                        # needs to be set in Quantconnect GUI
        self.investPercent = float(self.GetParameter("investPercent"))    # default 0.9
        self.shortDelta = float(self.GetParameter("shortDelta"))/-100     # default -0.25
        self.longDelta = float(self.GetParameter("longDelta"))/-100       # default -0.15
        self.t1DTE = int(self.GetParameter("t1DTE"))                      # default 25
        self.t2DTE = int(self.GetParameter("t2DTE"))                      # default 45

        # Handlers:
        self.order = OrderManagement(self)
        self.position = PositionManagement(self)
        self.signal = SignalManagement(self)

        # Helper Variables
        self.netCredit = None
        self.spread_expiry = None
        self.inPosition = False
        self.hedge = {
            "position": False,
            "quantity": None,
            "contract": None,
            "gainCategory": None,
            "profile": 1,
            "stopLoss": 0,
            "scale": 0,
        }
        self.stoppedOut = False
        self.openPortfolioValue = None

        # Set Option Instruments
        self.symbol = "SPY"
        self.option = self.AddOption(self.symbol, Resolution.Minute)
        self.option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        self.option.SetFilter(-40, -10, timedelta(self.t1DTE), timedelta(self.t2DTE))
        self.option.SetBuyingPowerModel(CustomBuyingPowerModel.CustomBuyingPowerModel())
        self.vix = self.AddData(CBOE, "VIX", Resolution.Daily)

        # Set Other Securities
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))
        self.equity = self.AddEquity(self.symbol, Resolution.Daily)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)

        # Initialize TA Parameters
        self.EMA10 = self.EMA(self.symbol, 10, Resolution.Daily)
        self.EMA25 = self.EMA(self.symbol, 25, Resolution.Daily)
        self.SMA50 = self.SMA(self.symbol, 50, Resolution.Daily)
        self.SMA125 = self.SMA(self.symbol, 125, Resolution.Daily)
        self.macd = self.MACD(self.symbol, 12, 26, 9, resolution = Resolution.Daily)
        self.windowSPX = RollingWindow[TradeBar](2)
        self.windowCross = RollingWindow[float](2)

        # Set warmup for Greeks and Indicators
        self.SetWarmUp(timedelta(days=125))

        # Check exits everyday
        self.Schedule.On(self.DateRules.EveryDay(
            self.symbol), self.TimeRules.AfterMarketOpen(self.symbol, 5), self.balance)

    def OnData(self, slice):
        # do some warmup checks
        if self.Time.hour == 9 and self.Time.minute == 31:

            if self.IsWarmingUp:
                return

            # add slice
            self.windowSPX.Add(slice[self.symbol])
            # add our ema to a window
            self.windowCross.Add(self.EMA10.Current.Value)
            # wait for windows to be ready
            if not self.windowSPX.IsReady and not self.windowCross.IsReady:
                return
            self.Log(f"o: {self.Securities['SPY'].Open} h: {self.Securities['SPY'].High} l: {self.Securities['SPY'].Low} c: {self.Securities['SPY'].Close}")
            # check for buy signals
            signal = self.signal.spread_buy()

            if self.inPosition:
                self.balance
            else:
                if signal:
                    # reset our stopped out variable so we can get into a hedge if we stop out.
                    self.stoppedOut = False
                    self.get_contracts(slice)

    def get_contracts(self, slice):
        # find contracts
        shortContract = None
        longContract = None

        for i in slice.OptionChains:
            chain = i.Value
            contracts = [x for x in chain if x.Right == 1]
            contracts = sorted(contracts, key=lambda x: (x.Expiry, x.Greeks.Delta))

            shortContract = min(contracts, key=lambda x: abs(x.Greeks.Delta-self.shortDelta))
            longContract = min(contracts, key=lambda x: abs(x.Greeks.Delta-self.longDelta))

        # Check that contracts are not the same
        if shortContract != longContract:
            self.order.open_spread(shortContract, longContract)

    def balance(self):

        # Determine if we need to hedge
        if not self.hedge["position"] and self.stoppedOut:
            signal = self.signal.hedge_buy()

            if signal:
                self.order.open_hedge()

        if self.hedge["position"]:

            signal = self.signal.hedge_sell()

            if signal:
                self.position.hedge_handler()

        # check that we have a spread open
        if self.openPortfolioValue is not None and self.inPosition:

            # Generate signals to sell
            signal = self.signal.spread_sell()

            if signal:
                self.order.close_all()

from QuantConnect import TradingDayType


class SignalManagement:

    def __init__(self, algorithm):
        self.algo = algorithm
        self.profiles = [
            {"hedge": [{0, 1}, {1, 2}, {2, 3}, {3, 4}], "range": [1, 2, 3, 4]},
            {"hedge": [{0, 1}, {0.3, 2}, {0.6, 3}, {0.9, 4}], "range": [0.3, 0.6, 0.9, 1.2]}
        ]
        self.tolerance = 0.1

    def spread_buy(self):
        # helper variables
        currEma10 = self.algo.windowCross[0]
        pastEma10 = self.algo.windowCross[1]

        # Conditions
        c1 = self.algo.SMA125.Current.Value < self.algo.Securities[self.algo.symbol].Open
        c2 = self.algo.EMA25.Current.Value < self.algo.Securities[self.algo.symbol].Open
        c3 = self.algo.EMA10.Current.Value > self.algo.EMA25.Current.Value
        c4 = pastEma10 < self.algo.EMA25.Current.Value and currEma10 > self.algo.EMA25.Current.Value
        c5 = self.algo.macd.Signal.Current.Value <= 2.1

        if c1 and c2 and c3 and c5:
            return True

        # if c4:
        #     self.algo.Debug('we got a cross')
        #     return True

    def spread_sell(self):
        # helper variables
        currClose = self.algo.windowSPX[0].Close
        oldClose = self.algo.windowSPX[1].Close
        days_till_expiry = list(self.algo.TradingCalendar.GetDaysByType(TradingDayType.BusinessDay, self.algo.Time, self.algo.Time + (self.algo.spread_expiry - self.algo.Time)))

        # Conditions
        c1 = safe_div(self.algo.Portfolio.TotalUnrealizedProfit, self.algo.netCredit) > 0.7
        c2 = safe_div(self.algo.Portfolio.TotalUnrealisedProfit, self.algo.Portfolio.TotalPortfolioValue) < -0.5
        c3 = self.algo.EMA10.Current.Value < self.algo.EMA25.Current.Value and self.algo.Portfolio.TotalUnrealisedProfit < 0
        c4 = ((currClose - oldClose) / oldClose) < -0.03
        c5 = len(days_till_expiry) > 7
        c6 = self.algo.macd.Signal.Current.Value > 2.1

        if c1:
            # self.algo.Debug('Sell Signal | 80% max profit')
            return True

        if c2:
            self.algo.Debug(f"Sell | -50 percent stop loss @ {self.algo.Time}")
            self.algo.stoppedOut = True
            return True

        if c3:
            self.algo.Debug('Sell | CROSS')
            # self.algo.stoppedOut = False
            return True

        # if c4:
        #     self.algo.Debug(f'Sell Signal | SPX falling {round(((currClose - oldClose) / oldClose), 2) * 100}% old close: {oldClose}. curr close: {currClose} @ {self.algo.Time}')
        #     self.algo.stoppedOut = True
        #     return True

        if c5 and c6:
            self.algo.Debug(f"Sell | macd: {round(self.algo.macd.Signal.Current.Value, 2)} days left: {len(days_till_expiry)}")
            self.algo.stoppedOut = True
            return True

    def hedge_buy(self):
        # helper variables
        currEma10 = self.algo.windowCross[0]
        pastEma10 = self.algo.windowCross[1]
        c1 = pastEma10 > self.algo.EMA25.Current.Value and currEma10 < self.algo.EMA25.Current.Value

        # Conditions
        if c1:
            self.algo.Debug(f"Hedge Buy Signal: 10EMA is {round(self.algo.EMA10.Current.Value, 2)} below 25 EMA {round(self.algo.EMA25.Current.Value, 2)} date: {self.algo.Time.date}")
            return True

    def hedge_sell(self):
        # helper variables
        unreal = self.algo.Portfolio[self.algo.hedge['contract'].Symbol].UnrealizedProfit
        cost = self.algo.Portfolio[self.algo.hedge['contract'].Symbol].HoldingsCost
        profitPercent = safe_div(unreal, cost)

        if self.algo.vix.Close > 35:
            if self.algo.hedge["profile"] == 1:
                # if a hedge went from non-volatile to volatile you need to restart the range checks
                # relative to where we are. so if you sold at 30% in non-V and it is volatile you're
                # hedge startegy needs to reset to the lowest bound in the gainCategory.
                self.algo.hedge["gainCategory"] = 0

            self.algo.hedge["profile"] = 0

        closestTuple = min(enumerate(self.profiles[self.algo.hedge["profile"]]["range"]), key=lambda x: abs(x[1]-profitPercent))
        self.algo.Debug(f"profit {profitPercent} profile: {self.algo.hedge['profile']} gc: {self.algo.hedge['gainCategory']}")
        if ((closestTuple[1] - self.tolerance) <= profitPercent):

            if closestTuple[0] == len(self.profiles[0]["range"])-1:
                self.algo.hedge["gainCategory"] = -1
                return True

            if profitPercent < (closestTuple[1] + self.tolerance):

                self.algo.hedge["gainCategory"] = closestTuple[0]

                hedgeProfile = self.profiles[self.algo.hedge["profile"]]["hedge"]

                if closestTuple[0] == 0 or self.algo.hedge["stopLoss"] >= list(hedgeProfile[closestTuple[0]-1])[0]:

                    self.algo.hedge["stopLoss"] = list(hedgeProfile[closestTuple[0]])[0]
                    self.algo.hedge["scale"] = list(hedgeProfile[closestTuple[0]])[1]

                return True


def safe_div(x, y):
    if x == 0:
        return 0
    return round(x / y, 2)