#region imports
from AlgorithmImports import *
import numpy as np
from math import log, sqrt, exp
from scipy.stats import norm
from QuantConnect.DataSource import *

#endregion
Balance = 10000
Minimum_Expiry = 720
Maximum_Expiry = 540
Min_Option_Price = 0
Profit_Taking = 2

class NOKLeapsBreakout(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2023, 3, 1)
        self.SetCash(Balance)
        equity = self.AddEquity("SPY", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.equity = equity.Symbol
        self.SetBenchmark(self.equity)
        option = self.AddOption("SPY", Resolution.Minute)
        option.SetFilter(0, 5, timedelta(Minimum_Expiry), timedelta(99999))
        self.yieldCurveTwo = self.AddData(USTreasuryYieldCurveRate, "USTYCR").Symbol
        self.SetWarmUp(timedelta(200))
        self.sigma = 0.25  # Set predetermined volatility
        self.T = 720/365  # Set predetermined time to expiration
        self.r = None
        self.Schedule.On(self.DateRules.EveryDay(self.equity), \
                        self.TimeRules.AfterMarketOpen(self.equity, 30), \
                        self.Plotting)                


    def OnData(self,data):
        if self.IsWarmingUp: return
        if data.ContainsKey(self.yieldCurveTwo):
            rates = data[self.yieldCurveTwo]
            self.r = ((rates.TwoYear) / 100)

        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        
        if option_invested:
            if self.Time + timedelta(Maximum_Expiry) > option_invested[0].ID.Date:
                self.Liquidate(option_invested[0], "Too close to expiration")
                last_o_id = self.Transactions.LastOrderId
                self.Debug (f"| {self.Time}   [+]---  Liquidate @ {str(self.Transactions.GetOrderTicket(last_o_id).AverageFillPrice)} || Stock @ {str(self.Portfolio[self.equity].Price)}|| Profit/Loss @ {str(self.Portfolio[option_invested[0]].LastTradeProfit)}")
                self.Debug (f"| {self.Time}   [-]--- REASON: ||  <{(Maximum_Expiry)} DTE  | {(option_invested[0].ID.Date - self.Time).days} DTE")
            return
        

        if not option_invested:
            for i in data.OptionChains:
                chains = i.Value
                self.BuyCall(chains)
                             
        # Buy options - Call
    def BuyCall(self,chains):
        self.S = self.Securities[self.equity].Close
        self.K = self.S * 1.05
        if self.r is None: return
        self.option_price = self.BlackScholesCall(self.S, self.K, self.T, self.r, self.sigma)
        expiry = sorted(chains,key = lambda x: x.Expiry, reverse=True)[0].Expiry
        calls = [i for i in chains if i.Expiry == expiry and i.Right == OptionRight.Call and i.AskPrice < self.option_price and i.AskPrice > Min_Option_Price]
        call_contracts = sorted(calls,key = lambda x: abs(x.Strike - x.UnderlyingLastPrice))

        if len(call_contracts) == 0: 
            return
        self.call = call_contracts[0]
        self.SetHoldings(self.call.Symbol, 0.8)
        quantity = self.Portfolio[self.call.Symbol].Quantity   # <-- quantity is zero and the following orders will be invalid
        self.LimitOrder(self.call.Symbol, -quantity, (self.call.AskPrice * Profit_Taking))
        self.Debug ("\r+-------------")
        self.Debug (f"| {self.Time}   [+]---  BUY  {str(self.call)} || Stock @ {str(self.call.UnderlyingLastPrice)} Option Price BS @ {(self.option_price)} rates: {self.r}")
        self.Debug(f"Order Quantity filled: {self.Portfolio[self.call.Symbol].Quantity}; Fill price: {self.Portfolio[self.call.Symbol].AveragePrice}")
        
        self.Log("Bought NOK Options")
        

    def Plotting(self):
        # plot underlying's price
        self.Plot("Data Chart", self.equity, self.Securities[self.equity].Close)
        # plot strike of call option
        option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        if option_invested:
                self.Plot("Data Chart", "strike", option_invested[0].ID.StrikePrice)

                self.Plot("Option Pricing", "BS - Pricing", self.option_price)
                self.Plot("Option Pricing", "Market - Pricing", self.Portfolio[self.call.Symbol].AveragePrice)


    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        # Cancel remaining order if limit order or stop loss order is executed
        if order.Status == OrderStatus.Filled:
            if order.Type == OrderType.Limit or OrderType.StopMarket:
                self.Transactions.CancelOpenOrders(order.Symbol)
        if order.Status == OrderStatus.Filled and order.Type == OrderType.Limit:
            self.Debug(f"{self.Time}   [+]---  SELL  {str(order.Quantity)} || Price @ {str(order.LimitPrice)}|| Profit @ {str(self.Portfolio[self.call.Symbol].LastTradeProfit)}")
            if order.Status == OrderStatus.Canceled:
                self.Log(str(orderEvent))
        # Liquidate before options are exercised
        if order.Type == OrderType.OptionExercise:
            self.Liquidate()

    def BlackScholesCall(self, S, K, T, r, sigma):
        d1 = (log(S / K) + (r + sigma ** 2 / 2) * T) / (sigma * sqrt(T))
        d2 = d1 - sigma * sqrt(T)
        return S * norm.cdf(d1) - K * exp(-r * T) * norm.cdf(d2)