from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Data import *
from QuantConnect.Algorithm import *
import numpy as np
from datetime import timedelta

from QuantConnect.Securities.Option import OptionPriceModels
import pandas as pd
import numpy as np

class BasicTemplateIndexOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2009, 1, 1)
        self.SetEndDate(2009, 12, 31)
        self.initial_cash = 1000000
        self.SetCash(self.initial_cash)
        self.benchmark_shares = None

        # Set securities and universe
        self.spy = self.AddEquity('SPY', Resolution.Minute).Symbol
        self.Securities['SPY'].SetDataNormalizationMode(DataNormalizationMode.Raw);
        self.spx = self.AddIndex('SPX', Resolution.Minute).Symbol
        self.underlying = self.spy

        # Set parameters
        self.spy_percentage = 1.0 # % invested in SPY
        self.max_leverage = 1.05
        self.min_leverage = 0.95
        self.target_expiry = 365
        self.multiplier = 100 # 100 for SPY and 1000 for SPX since SPX is 10 times largest

        # Set helper amounts
        self.opts_to_trade = pd.Series()
        self.opts_to_sell = pd.Series()
        self.opt_invested = False

        # Use SPY as the benchmark
        self.SetBenchmark('SPY')
        self.SetWarmUp(TimeSpan.FromDays(30))

        # Schedule OTM put hedge every week 30 mins prior to market close
        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.BeforeMarketClose(self.underlying,30),
                         self.opt_replication)

        self.Schedule.On(self.DateRules.EveryDay(),
                         self.TimeRules.BeforeMarketClose(self.underlying,15),
                         self.check_expiries)
                         
#        self.Schedule.On(self.DateRules.EveryDay(),
#                         self.TimeRules.BeforeMarketClose(self.underlying,5),
#                         self.check_leverage)

    def OnAssignmentOrderEvent(self, assignmentEvent):
        # If we get assigned, close out all positions
        self.Liquidate()
        self.Log(str(assignmentEvent))
        self.opt_invested = False

    def OnData(self, data):
        if self.IsWarmingUp:
            return

        if self.underlying not in data.Bars:
            return

        if not self.opts_to_trade.empty:
            remaining_opts = self.opts_to_trade.copy(True)
            for opt, amount in self.opts_to_trade.iteritems():
                if data.ContainsKey(opt):
                    self.MarketOrder(opt, amount)
                    del remaining_opts[opt]
                    self.opt_invested = True
            self.opts_to_trade = remaining_opts

    def get_benchmark_performance(self): 
        price = self.Securities[self.underlying].Close

        if not self.benchmark_shares:
            self.benchmark_shares = self.initial_cash / price
            
        return self.benchmark_shares * price
        
    def log_holdings(self):
        # Log holdings, amount of portfolio hedged, etc.
        return

    def check_expiries(self):
        opt_pos = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        if opt_pos:
            # Get current time
            curr_dt = self.Time
            close_out_opts = [x for x in opt_pos if (x.ID.Date - curr_dt) < timedelta(0)]
            for opt in close_out_opts:
                # Get SPY level to see if option is ITM
                underlying_px = self.Securities[self.underlying].Price
                if opt.ID.OptionRight == 1:
                    if opt.ID.StrikePrice >= underlying_px:
                        self.Liquidate(opt)
                        self.opt_invested = False
                else:
                    if opt.ID.StrikePrice <= underlying_px:
                        self.Liquidate(opt)
                        self.opt_invested = False
                

    def opt_replication(self):
        if self.opt_invested:
            return
        
        # Get contract list and convert to DF
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlying, self.Time)
        if len(contracts) == 0:
            return
        opts = contract_list_to_df(contracts)
        puts = opts[opts['Right'] == 1]
        calls = opts[opts['Right'] == 0]

        # Get index level
        curr_idx_lvl = self.Securities[self.underlying].Price

        # Get puts 1 year out
        filtered_puts_expiry = self.select_closest_expiry(puts, self.target_expiry)
        filtered_puts_expiry = puts.loc[filtered_puts_expiry.index]

        # Get calls 1 year out
        filtered_calls_expiry = self.select_closest_expiry(calls, self.target_expiry)
        filtered_calls_expiry = calls.loc[filtered_calls_expiry.index]

        # Get ATM put
        puts_to_sell = self.select_closest_strike(filtered_puts_expiry, curr_idx_lvl)
        puts_to_sell_info = puts.loc[puts_to_sell]
        puts_to_sell_symbol = puts_to_sell_info['Symbol']

        # Get ATM call
        calls_to_buy = self.select_closest_strike(filtered_calls_expiry, curr_idx_lvl)
        calls_to_buy_info = calls.loc[calls_to_buy]
        calls_to_buy_symbol = calls_to_buy_info['Symbol']

        # Calculate size to buy
        puts_to_sell_amount = -self.Portfolio.TotalPortfolioValue / (self.multiplier * curr_idx_lvl)
        calls_to_buy_amount = self.Portfolio.TotalPortfolioValue / (self.multiplier * curr_idx_lvl)
        
        # Create series of weights
        opts_to_trade = pd.Series(index = puts_to_sell_symbol, data = puts_to_sell_amount)
        opts_to_trade = opts_to_trade.append(pd.Series(index = calls_to_buy_symbol, data = calls_to_buy_amount))
        
        # Calculate final amounts
        # opts_to_buy *= opts_to_buy_amount
        opts_to_trade = opts_to_trade.round(0)

        # Get rid of zero quantity
        opts_to_trade = opts_to_trade[opts_to_trade.abs() > 0]

        # Add security
        [self.AddOptionContract(sym, Resolution.Minute) for sym in opts_to_trade.index]
        self.opts_to_trade = opts_to_trade

    def select_closest_strike(self, options, strike):
        strike_distance = (options['Strike'] - strike).abs()
        opts_to_trade = strike_distance.groupby(options['Expiry']).idxmin().values
        return opts_to_trade

    def select_closest_expiry(self, options, expiry_time):
        """Select option closest to desired expiration

        Args:
            options (pd.DataFrame): df with options information
            expiry_time (int): number of calendar days until expiry desired

        Returns:
            pd.Series: series with days to expiration as values and option symbols as index
        """
        # Get current time
        curr_dt = self.Time

        # Get desired expiration date
        desired_expiry = curr_dt + pd.Timedelta(days = expiry_time)
        desired_expiry = desired_expiry.replace(hour = 0, minute = 0)

        # Find distance to expiration
        expiry_distances = options['Expiry'] - desired_expiry

        # Get only expiries 1 year out
        #expiry_distances = expiry_distances[expiry_distances > pd.Timedelta(0)]

        # Get lowest expiration distance and options expiring on that date
        min_expiry_days = expiry_distances.abs().min()
        min_expiry = expiry_distances.abs().idxmin()
        min_expiry_date = expiry_distances.loc[min_expiry]
        
        # If min expiry is after desired time, then get expiry before desired time
        if min_expiry_days.days > 0:
            second_min_expiry_days = expiry_distances[expiry_distances < min_expiry_days].max()
            second_min_expiry = expiry_distances[expiry_distances < min_expiry_days].idxmax()
            second_min_expiry_date = expiry_distances.loc[second_min_expiry]
        else:
            second_min_expiry_days = expiry_distances[expiry_distances > min_expiry_days].min()
            second_min_expiry = expiry_distances[expiry_distances > min_expiry_days].idxmin()
            second_min_expiry_date = expiry_distances.loc[second_min_expiry]

        # If the expiry is more than 5 days out, then create custom expiry
        # by trading two options of different expiries
        if min_expiry_days.days > 5:
            opts_to_trade = expiry_distances[expiry_distances.isin([min_expiry_date, second_min_expiry_date])]
        else:
            opts_to_trade = expiry_distances[expiry_distances == min_expiry_date]

        return opts_to_trade

    def OnEndOfAlgorithm(self) -> None:
        spy_pos = self.Portfolio['SPY'].Quantity
        self.Log('SPY quantity: {0}'.format(spy_pos))

def contract_list_to_df(contract_list):
    return pd.DataFrame([[x.ID.OptionRight,
                          float(x.ID.StrikePrice),
                          x.ID.Date,
                          x.ID.ToString()
                          ] for x in contract_list],
                        index=[x.ID for x in contract_list],
                        columns=['Right', 'Strike', 'Expiry',
                                 'Symbol'])