#region imports
from AlgorithmImports import *
from abc import ABC, abstractmethod
#endregion

class Charts(ABC):
    def __init__(
        self, 
        algo: QCAlgorithm, 
        chart_name: str, 
        series_list: List[Dict]
    ) -> None:

        self._algo: QCAlgorithm = algo
        self._chart_name: str = chart_name
        self._series_list: List[Dict] = series_list

        self._add_chart()
        self._scheduler()

    def _add_chart(self) -> None:
        chart = Chart(self._chart_name)
        self._algo.add_chart(chart)

        for series in self._series_list:
            chart.AddSeries(Series(series['name'], series['type'], series['unit']))

    def _scheduler(self) -> None:
        algo: QCAlgorithm = self._algo
        algo.schedule.on(
            algo.date_rules.every_day('SPY'), 
            algo.time_rules.before_market_close('SPY', 0), 
            self._update_chart
        )

    @abstractmethod
    def _update_chart(self):
        pass

class Benchmark(Charts):
    def __init__(
        self, 
        algo: QCAlgorithm, 
        benchmark_symbol: Symbol,
        init_cash: int
    ) -> None:

        series_list: List[Dict] = [
            {
            'name': f'Benchmark {benchmark_symbol.value}', 
            'type': SeriesType.LINE, 
            'unit': '$'
            }
        ]

        self._benchmark_values: List[float] = []
        self._benchmark_symbol: Symbol = benchmark_symbol
        self._init_cash: int = init_cash

        super().__init__(algo = algo, chart_name = 'Strategy Equity', series_list = series_list)
    
    def _update_chart(self):
        # wait for available data
        if not self._algo.portfolio.invested:
            return

        # print benchmark in main equity plot
        mkt_price_df: DataFrame = self._algo.history(self._benchmark_symbol, 2, Resolution.DAILY)
        if not mkt_price_df.empty:
            benchmark_price: float = mkt_price_df['close'].unstack(level= 0).iloc[-1]
            if len(self._benchmark_values) == 2:
                self._benchmark_values[-1] = benchmark_price
                benchmark_perf: float = self._init_cash * (self._benchmark_values[-1] / self._benchmark_values[0])
                self._algo.plot(self._chart_name, self._series_list[0]['name'], benchmark_perf)
            else:
                self._benchmark_values.append(benchmark_price)

# region imports
from AlgorithmImports import *
from enum import Enum
from dataclasses import dataclass
from charts import Benchmark
# endregion

class OptionTradeFunction(Enum):
    BUY = 1
    SELL = 2

@dataclass
class OptionLegSetup():
    trade_fn: OptionTradeFunction
    option_right: OptionRight
    dte: int
    otm: float

class OptionHedge(QCAlgorithm):

    _percentage_traded: float = 1.
    _DTE: int = 30
    _OTM: float = 0.01
    _benchmark_ticker: str = 'SPY'
    
    _hedge_ticker_setup: Dict[str, OptionLegSetup] = {
        'UDN' : OptionLegSetup(OptionTradeFunction.BUY, OptionRight.CALL, _DTE, _OTM), 
        'FXB' : OptionLegSetup(OptionTradeFunction.BUY, OptionRight.PUT, _DTE, _OTM) 
    }

    def initialize(self):
        self._init_cash: int = 100_000
        self.set_start_date(2010, 1, 1)
        self.set_cash(self._init_cash)

        self._market: Symbol = self.add_equity('SPY', Resolution.MINUTE).symbol

        # get last known price after subscribing option contract
        self.set_security_initializer(
            CustomSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices))
        )
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = False

        # option storage
        self._subscribed_contracts: Dict[Symbol, Optional[Symbol]] = {}

        self._benchmark_symbol: Symbol = self.add_equity(self._benchmark_ticker, Resolution.MINUTE).symbol

        for ticker in self._hedge_ticker_setup:
            eq: Equity = self.add_equity(ticker, Resolution.MINUTE)
            eq.set_data_normalization_mode(DataNormalizationMode.RAW)
            self._subscribed_contracts[eq.symbol] = None

        self._contracts_added: Set = set()
        self._recent_day: int = -1

        # charting
        Benchmark(self, self._benchmark_symbol, self._init_cash)

    def on_data(self, slice: Slice) -> None:
        if self.is_warming_up:
            return
        
        if not self.is_market_open(self._market):
            return
        
        # once a day execution
        if self._recent_day != self.time.day:
            self._recent_day = self.time.day

            # trade options
            self._option_strategy(self._percentage_traded)

    def _option_strategy(self, percentage_traded: float) -> None:
        # contract is assigned and it is about to expire -> remove contract
        for symbol, c in self._subscribed_contracts.items():
            if c is not None and self.time + timedelta(days=1) >= c.id.date:
                self.remove_option_contract(c)
                self._subscribed_contracts[symbol] = None

        for symbol, c in self._subscribed_contracts.items():
            if c is None:
                self._subscribed_contracts[symbol] = self._filter_options(
                    symbol, 
                    self._hedge_ticker_setup[symbol.value].option_right, 
                    self._hedge_ticker_setup[symbol.value].dte, 
                    self._hedge_ticker_setup[symbol.value].otm
                )

            if c is not None and self.current_slice.contains_key(c):
                self._trade_option(c, percentage_traded, self._hedge_ticker_setup[symbol.value].trade_fn, tag='')

                # market buy and hold
                if not self.portfolio[self._benchmark_symbol].invested:
                    self.set_holdings(self._benchmark_symbol, 1.)

    def _filter_options(
        self, 
        symbol: Symbol, 
        option_right: OptionRight, 
        dte: int, 
        moneyness: float
    ) -> Optional[Symbol]:
    
        contracts: Iterable[Symbol] = self.option_chain_provider.get_option_contract_list(symbol, self.time)
        
        underlying_price: float = self.securities[symbol].price
        
        if option_right == OptionRight.CALL:
            options: List[Symbol] = [
                i for i in contracts if i.id.option_right == OptionRight.CALL and   
                i.id.strike_price >= underlying_price * (1 + moneyness) and
                i.id.date >= self.time + timedelta(days=dte)
            ]
        else:
            options: List[Symbol] = [
                i for i in contracts if i.id.option_right == OptionRight.PUT and   
                i.id.strike_price <= underlying_price * (1 - moneyness) and
                i.id.date >= self.time + timedelta(days=dte)
            ]

        if len(options) > 0:
            expiry: datetime.datetime = min(list(map(lambda x: x.id.date, options)))

            index: int = 0 if option_right == OptionRight.PUT else -1
            contract: Symbol = sorted([o for o in options if o.id.date == expiry], key = lambda x: underlying_price - x.id.strike_price)[index]

            # prevent multiple subscriptions
            if contract not in self._contracts_added:
                self._contracts_added.add(contract)
                option = self.add_option_contract(contract, Resolution.MINUTE)
                option.is_tradable = True

            return contract
        else:
            self.log(f'Connot find filtered options for {symbol}')
            return None

    def _trade_option(self, contract: Symbol, percentage_traded: float, option_trade_function: OptionTradeFunction, tag: str) -> None:
        if contract and not self.portfolio[contract].invested:
            trade_direction: int = 1 if option_trade_function == OptionTradeFunction.BUY else -1
            q: float = self._get_amount_as_fraction_of_portfolio(contract, percentage_traded)
            # q: float = self._get_amount_as_fraction_of_cash(contract, percentage_traded)
            # q: float = self.calculate_order_quantity(contract, percentage_traded)

            self.market_order(
                contract, 
                trade_direction * q, 
                tag=tag
            )

    def _get_amount_as_fraction_of_portfolio(self, option_symbol: Symbol, fraction: float) -> float:
        multiplier: int = self.securities[option_symbol].contract_multiplier
        target_notional: float = self.portfolio.total_portfolio_value * fraction
        notional_of_contract: float = multiplier * self.securities[option_symbol.underlying].price
        amount: float = target_notional / notional_of_contract

        return amount

    def _get_amount_as_fraction_of_cash(self, option_symbol: Symbol, fraction: float) -> float:
        holding_value: float = self.portfolio.cash
        multiplier: int = self.securities[option_symbol].contract_multiplier
        target_notional: float = holding_value * fraction
        notional_of_contract: float = multiplier * self.securities[option_symbol.underlying].price
        amount: float = target_notional / notional_of_contract

        return amount

class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder) -> None:
        super().__init__(brokerage_model, security_seeder)    

    def initialize(self, security: Security) -> None:
        super().initialize(security)

        # overwrite the price model
        if security.type == SecurityType.OPTION: # option type
            security.price_model = OptionPriceModels.black_scholes()
            security.set_option_assignment_model(NullOptionAssignmentModel())