# Jingxiang Zou 
# Quantitative Research Analyst

# region imports
from datetime import timedelta
from numpy import array, exp, log, sqrt, arange, var
from AlgorithmImports import *
from math import floor
import numpy as np 
from pandas import DataFrame
from datetime import datetime 
import pandas as pd
# endregion

class HedgedFutureStrategy(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2023, 4, 1)
        self.set_cash(1000000)
        self.set_security_initializer(CustomSecurityInitializer(self))
        self.last_roll_date_1 = None  # Variable to track the last roll date
        self.last_roll_date_2 = None

        # leverage and protection level 
        # self.leverage = 5.0
        self.protection_level = 1.0
        self.margin_pct = 0.25

        # the index 
        self.nasdaq100 = self.AddIndex("NDX", Resolution.MINUTE).Symbol
        self.cumulative_return = 1  # start with 1 for multiplicative return calculation
        self.previous_price = None

        # our strategy 
        self.previous_port_value = None 
        self.normalized_port_value = 1

        # add the futures to the universe 
        self.universe_settings.asynchronous = True
        self.future = self.add_future(Futures.Indices.NASDAQ_100_E_MINI,
            extended_market_hours=False,
            data_mapping_mode=DataMappingMode.LAST_TRADING_DAY, # time to roll : on the last trading day 
            data_normalization_mode=DataNormalizationMode.RAW, # adjust for jumps, backward in time 
            contract_depth_offset=0,  # getting the front contract 
            resolution=Resolution.MINUTE
            )
        self.future.set_filter(0, 180)

        # add the option to the universe
        self.add_future_option(self.future.symbol, self._contract_selector)

        self.lev = None
        self.rets = []



    def _contract_selector(self, universe):
            return universe.Expiration(0, 365)

    def on_end_of_day(self, symbol):
        self.plot(f'Security Type {symbol.id.security_type}', 'EOD', self.securities[symbol].price)

    def closest_number(self, target, sorted_list):
        # initialize variables to keep track of the closest number and the smallest difference
        closest = sorted_list[0]
        min_difference = abs(target - closest)

        # iterate through the list to find the number with the smallest difference
        for number in sorted_list:
            difference = abs(target - number)
            if difference < min_difference:
                closest = number
                min_difference = difference

        return closest

    def on_data(self, data: Slice) -> None:
  
        if self.Time.date() != self.last_roll_date_2 and self.Time.date() > datetime(2021, 11, 12).date():  # Check if the date has changed
            self.last_roll_date_2 = self.Time.date()  # Update the last roll date

            if self.future.Symbol in data.FutureChains:
                future_chain = data.FutureChains[self.future.Symbol]
                valid_contracts = [contract for contract in future_chain if contract.Expiry > self.Time]
            
                if valid_contracts:
                    self.front_contract = sorted(valid_contracts, key=lambda x: x.Expiry)[0]
                    self.debug(f"what's in your portfolio ??? ")
                    for symbol in self.Portfolio.Keys:
                        holding = self.Portfolio[symbol]
                        if holding.Invested:
                            self.Debug(f"{symbol}: {holding.Quantity}")

            if not self.portfolio.invested:

                watcher = 0

                if self.future.Symbol in data.FutureChains:
                    future_chain = data.FutureChains[self.future.Symbol]
                    valid_contracts = [contract for contract in future_chain if contract.Expiry > self.Time]
            
                    if valid_contracts:
                        front_contract = sorted(valid_contracts, key=lambda x: x.Expiry)[0]

                    kks = [option_chain for _, option_chain in data.option_chains.items()]      
                    kss = sorted(kks, key= lambda x : list(x.contracts)[0].value.expiry, reverse=False)
                    sss1 = [x for x in kss if list(x.contracts)[0].value.expiry > front_contract.expiry - timedelta(30)]

                    for option_chain in sss1:
                
                        watcher += 1
                        put_options = [contract for contract in option_chain if contract.Right == OptionRight.PUT]
                        if put_options:

                            cleveland = sorted([int(x.strike) / 10 for x in put_options])
                            fp = self.Securities[self.front_contract.Symbol].Price
                            self.debug(f"this is the current front NQ price = {fp}")
                            self.debug(f"these are the striks")
                            self.debug(cleveland[100:105])

                            gs = self.closest_number(fp * self.protection_level, cleveland)
                            self.debug(f"here is the nearest strike = {gs}")
                            opt_contract = [x for x in put_options if int(x.strike) == gs * 10][0]
                            self.market_order(opt_contract, 1)

                            if watcher == 1:
                                break

    def on_end_of_algorithm(self):
        for symbol, security in self.securities.items():
            if security.is_tradable:
                self.log(f'{self.time} :: {symbol} :: {security.price}')

# Outside of the algorithm class
class CustomSecurityInitializer(BrokerageModelSecurityInitializer):

    def __init__(self, algorithm: QCAlgorithm) -> None:
        security_seeder = SecuritySeeder.NULL
        super().__init__(algorithm.brokerage_model, security_seeder)    
    def initialize(self, security: Security) -> None:
        # First, call the superclass definition.
        # This method sets the reality models of each security using the default reality models of the brokerage model.
        super().initialize(security)

        # Next, overwrite the assignment model
        if security.Type == SecurityType.FUTURE_OPTION: # Option type
            security.set_option_exercise_model(MyOptionExerciseModel())

class MyOptionExerciseModel(DefaultExerciseModel):
    def option_exercise(self, option: Option, order: OptionExerciseOrder) -> List[OrderEvent]:
        underlying = option.underlying
        utc_time = Extensions.convert_to_utc(option.local_time, option.exchange.time_zone)
        underlying_price = underlying.price
        strike = option.scaled_strike_price / 10  # NQ hack
        intrinsic_value = max(0, underlying_price - strike if option.right == OptionRight.Call else strike - underlying_price)
        in_the_money = intrinsic_value >= 0.01
        is_assignment = in_the_money and option.holdings.is_short

        order_event = OrderEvent(order.id, option.symbol, utc_time, OrderStatus.FILLED, 
            Extensions.get_order_direction(order.quantity), 0, order.quantity, OrderFee.ZERO,
            Messages.DefaultExerciseModel.ContractHoldingsAdjustmentFillTag(in_the_money, is_assignment, option)
        )
        order_event.is_assignment = is_assignment
        order_event.is_in_the_money = in_the_money

        order_events = [order_event]

        if in_the_money and option.exercise_settlement == SettlementType.PHYSICAL_DELIVERY:
            exercise_quantity = option.get_exercise_quantity(order.quantity)
            delivery_order_event = OrderEvent(order.id, underlying.symbol, utc_time, OrderStatus.FILLED,
                Extensions.get_order_direction(exercise_quantity), strike, exercise_quantity, OrderFee.ZERO,
                Messages.DefaultExerciseModel.OptionAssignment if is_assignment else Messages.DefaultExerciseModel.OptionExercise
            )
            delivery_order_event.is_in_the_money = True
            order_events.append(delivery_order_event)

        return order_events