# region imports
from AlgorithmImports import *
from scipy import special
from scipy.stats import gamma, invweibull
# endregion

class MaxLossVaRShortPut(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2025, 1, 1)
        self.set_cash(100000)
        self.set_security_initializer(VolumeShareFillSecurityInitializer(self, 1))
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        
        # We want to trade the 95%VaR.
        self._alpha = 0.95
        self.lookback = 1000
        self.trade_period = 5
        self._orders = {}
        self.put_strikes = {}

        self.symbols = [self.add_equity(ticker).symbol for ticker in ["TQQQ", "SVXY", "TMF", "EDZ", "UGL", "UUP"]]
        # Rebalance weekly since we're trading the option expiring this week to avoid over-trading.
        self.schedule.on(
            self.date_rules.week_start(self.symbols[0]), 
            self.time_rules.after_market_open(self.symbols[0], 1), 
            self.rebalance
        )

    def rebalance(self):
        # Call the historical data to fit the Inverse Weibull distribution to model max loss.
        ret = self.history(self.symbols, 252+self.lookback+self.trade_period, Resolution.DAILY).close.unstack(0).pct_change().dropna()
        # Obtain the trade parameters, including strike price and expected loss.
        cvar = self.get_var(ret)

        # Short a put to earn credit in N% confidence that it will not be assigned.
        for symbol, strike in self.put_strikes.items():
            chain = self.option_chain(symbol)
            # Trade the week-expiring put to ensure short value and liquidity.
            filtered = [x for x in chain if x.right == OptionRight.PUT and x.expiry <= self.time + timedelta(self.trade_period + 1)]
            if filtered:
                expiry = max(x.expiry for x in filtered)
                put = sorted(
                    [x for x in filtered if x.expiry == expiry],
                    key=lambda x: abs(x.strike - strike)
                )
                if not put:
                    continue
                # Request the contract data for trading.
                put_symbol = self.add_option_contract(put[0]).symbol

                # Position sizing by the Kelly Criterion to maximize expected return.
                strike = put_symbol.id.strike_price
                weight = self.get_bet_weight(put_symbol, strike, cvar[symbol])
                # Obtain the actual number of contract to be ordered.
                quantity = weight * self.portfolio.total_portfolio_value / strike // self.securities[put_symbol].symbol_properties.contract_multiplier
                if quantity:
                    self._orders[put_symbol] = quantity

    def get_bet_weight(self, symbol, strike, cvar):
        price = self.securities[symbol].bid_price if self.securities[symbol].bid_price != 0 else self.securities[symbol].price
        # Use CVaR minus strike as the expected loss. 
        # Accounted for commission also (both buy and sell plus other costs), we estimate the cost to be $1.3 per contract.
        multiplier = self.securities[symbol].symbol_properties.contract_multiplier
        loss = (cvar - strike + 1.3 / multiplier) / price
        gain = (price - 1.3 / multiplier) / price
        if gain <= 0:
            return 0
        # Kelly criterion formula: win rate / expected loss - loss rate / expected win
        # Win and loss rate was set according to VaR.
        k = self._alpha / loss - (1 - self._alpha) / gain
        # Reducing 50% risk by 25% bet (Ed Thorp)
        # We limit the maximum position to be 10%.
        return min(0.1, k * 0.75)

    def get_var(self, ret):
        self.put_strikes.clear()
        cvar = {}
        
        # Obtain the rolling max loss to fit the Inverse Weibull distribution to model catastrophic loss.
        max_loss = ((1 + ret).rolling(self.trade_period).apply(np.prod, raw=True) - 1).rolling(self.lookback).min().iloc[self.lookback+self.trade_period:]

        for symbol in max_loss.columns:
            # Fit Inverse Weibull distribution.
            params = invweibull.fit(max_loss[symbol])
            shape, loc, scale = params

            # Get N% VaR of each symbol analytically as the N% confident level that the put will not be assigned.
            # We also need the CVaR to measure the loss of the bet.
            pi = scale * (-np.log(1 - self._alpha) ** (1 / shape))
            var_ = loc + pi
            cvar_ = var_ + scale * special.gamma(1 - 1 / shape) * gamma.cdf((scale / pi) ** shape, 1 - 1 / shape) / self._alpha

            cvar[symbol] = abs(cvar_) * self.securities[symbol].price
            self.put_strikes[symbol] = (1 - abs(var_)) * self.securities[symbol].price
        
        return cvar

    def on_data(self, slice):
        # Order when there is a quote to be more realistic and likely to be filled.
        for symbol, size in self._orders.copy().items():
            bar = slice.quote_bars.get(symbol)
            if bar:
                self.limit_order(symbol, -size, round(bar.high, 2))
                self._orders.pop(symbol)

    def on_assignment_order_event(self, assignment_event):
        # Liquidate the assigned underlyings to avoid volatility.
        self.market_order(
            assignment_event.symbol.underlying, 
            -assignment_event.fill_quantity * self.securities[assignment_event.symbol].symbol_properties.contract_multiplier, 
            tag="liquidate assigned"
        )

class VolumeShareFillModel(FillModel):
    def __init__(self, algorithm: QCAlgorithm, maximum_ratio: float = 1):
        self.algorithm = algorithm
        self.maximum_ratio = maximum_ratio
        self.absolute_remaining_by_order_id = {}

    def market_fill(self, asset, order):
        absolute_remaining = self.absolute_remaining_by_order_id.get(order.id, order.absolute_quantity)

        fill = super().market_fill(asset, order)
        # Set the fill amount to 100% of the previous bar.
        volume = asset.bid_size if order.quantity < 0 else asset.ask_size
        fill.fill_quantity = np.sign(order.quantity) * volume * self.maximum_ratio

        if (min(abs(fill.fill_quantity), absolute_remaining) == absolute_remaining):
            fill.fill_quantity = np.sign(order.quantity) * absolute_remaining
            fill.status = OrderStatus.FILLED
            self.absolute_remaining_by_order_id.pop(order.id, None)
        else:
            fill.status = OrderStatus.PARTIALLY_FILLED
            self.absolute_remaining_by_order_id[order.id] = absolute_remaining - abs(fill.fill_quantity)
            price = fill.fill_price

        return fill

class VolumeShareFillSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, algorithm: QCAlgorithm, fill_ratio: float = 1) -> None:
        super().__init__(algorithm.brokerage_model, FuncSecuritySeeder(algorithm.get_last_known_prices))
        self.fill_model = VolumeShareFillModel(algorithm, fill_ratio)
        
    def initialize(self, security: Security) -> None:
        super().initialize(security)
        security.set_fill_model(self.fill_model)
        security.set_slippage_model(VolumeShareSlippageModel(1, 0.5))