# region imports
from AlgorithmImports import *
from QuantConnect.DataSource import *
from trade_manager import TradeManager
from dataclasses import dataclass
# endregion

# VIXY Short Model:
#   1. Strategy 1:
#       Hold VIXY ETF for 60 days if VIX value is above 20 for 4 days straight.
#       There are 60 EW brackets that are filled with new positions and each position is held for 60 days.
#
#   2. Strategy 2:
#       Hold VIXY ETF for 60 days if VIX value is above 18 at least for 1 day.
#       There are 60 EW brackets that are filled with new positions and each position is held for 60 days.
#
# 50% of notional values is allocated for each strategy. Signal is checked daily. Most recent daily VIX is used as a signal value.
# Trades are executed 15 minutes before market close.

class MetatronVIXYShortModel(QCAlgorithm):

    _notional_value: float = 80_000
    _trade_exec_minute_offset: int = 15
    _holding_period: int = 60
    _strategy_weight: float = 0.5

    # Strategy 1
    _volatility_threshold_1: float = 20.
    _lookback_1: int = 4
    
    # Strategy 2
    _volatility_threshold_2: float = 18.
    _lookback_2: int = 1

    def initialize(self) -> None:
        self.set_start_date(2020, 1, 1)
        self.set_cash(self._notional_value)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = True

        self._traded_asset: Symbol = self.add_equity('VIXY', Resolution.MINUTE).symbol
        self._volatility: Symbol = self.add_data(CBOE, 'VIX', Resolution.Daily).symbol
        
        self._volatility_values: RollingWindow = RollingWindow[float](max(self._lookback_1, self._lookback_2))

        self._traded_strategies: List[TradedStrategy] = [
            TradedStrategy(
                TradeManager(self, 0, self._holding_period, self._holding_period, self._strategy_weight), 
                self._volatility_threshold_1, 
                self._lookback_1
            ),

            TradedStrategy(
                TradeManager(self, 0, self._holding_period, self._holding_period, self._strategy_weight), 
                self._volatility_threshold_2, 
                self._lookback_2
            )
        ]

        # schedule function
        self._day_close_flag: bool = False

        self.schedule.on(
            self.date_rules.every_day(self._traded_asset), 
            self.time_rules.before_market_close(self._traded_asset, self._trade_exec_minute_offset), 
            self._before_eod
        )

    def on_data(self, slice: Slice) -> None:
        if self._day_close_flag:
            self.log(f'New day logic initiated')
            self._day_close_flag = False

            if self.securities[self._volatility].get_last_data() and self.securities[self._volatility].price != 0:
                self.log(f'Symbol {self._volatility} data present in the algorithm')
                self._volatility_values.add(self.securities[self._volatility].price)

            # liquidate opened symbols after holding period
            for ts in self._traded_strategies:
                ts.trade_manager.try_liquidate()

            if not self._volatility_values.is_ready:
                self.log(f'{self._volatility} history is not ready yet')
                return

            if slice.contains_key(self._traded_asset) and slice[self._traded_asset]:
                self.log(f'Symbol {self._traded_asset} data present in the algorithm')
                self.log(f'{self._volatility} values: {list(self._volatility_values)}')

                for ts in self._traded_strategies:
                    related_vix_values: List[float] = list(self._volatility_values)[:ts.lookback]
                    if all([x > ts.volatility_threshold for x in related_vix_values]):
                        self.log(f'Strategy {ts}: all VIX values are above threshold: {related_vix_values}. Trade initiated')
                        ts.trade_manager.add(self._traded_asset, False, self.securities[self._traded_asset].ask_price)
                    
    def _before_eod(self) -> None:
        self._day_close_flag = True

@dataclass
class TradedStrategy:
    trade_manager: TradeManager
    volatility_threshold: float
    lookback: int

# region imports
from AlgorithmImports import *
from dataclasses import dataclass
from typing import Optional
# endregion

# NOTE: Manager for new trades. It's represented by certain count of equally weighted brackets for long and short positions.
# If there's a place for new trade, it will be managed for time of holding period.

class TradeManager():
    def __init__(
        self, 
        algorithm: QCAlgorithm, 
        long_size: int, 
        short_size: int, 
        holding_period: int,
        weight: float = 0.5
    ) -> None:

        self._algorithm: QCAlgorithm = algorithm
        
        self._long_size: int = long_size
        self._short_size: int = short_size
        
        self._long_len: int = 0
        self._short_len: int = 0
    
        # array of ManagedSymbols
        self._symbols: List[ManagedSymbol] = []
        
        self._holding_period: int = holding_period
        self._weight: float = weight
    
    def add(
        self, 
        symbol: Symbol, 
        long_flag: bool, 
        price: float, 
        MOC_flag: bool = False
    ) -> None:
    
        quantity: Optional[int] = None
        
        # open new long trade
        if long_flag:
            if self._long_len < self._long_size:
                # there's a place for new trade
                quantity: int = int(self._algorithm._notional_value * self._weight / self._long_size / price)
                if quantity != 0:
                    if MOC_flag:
                        self._algorithm.market_on_close_order(symbol, quantity)
                    else:
                        self._algorithm.market_order(symbol, quantity)
                    self._long_len += 1
                else:
                    self._algorithm.log(f"Final position quantity is 0 for symbol: {symbol}")
            else:
                self._algorithm.log(f"There's no place for additional long trade {symbol}")

        # open new short trade
        else:
            if self._short_len < self._short_size:
                # there's a place for new trade
                
                quantity: int = int((self._algorithm._notional_value * self._weight / self._short_size / price))
                if quantity != 0:
                    if MOC_flag:
                        self._algorithm.market_on_close_order(symbol, -quantity)
                    else:
                        self._algorithm.market_order(symbol, -quantity)
                    self._short_len += 1
                else:
                    self._algorithm.log(f"Final position quantity is 0 for symbol: {symbol}")
            else:
                self._algorithm.log(f"There's no place for additional short trade {symbol}")
   
        if quantity:
            self._symbols.append(
                ManagedSymbol(symbol, self._holding_period, long_flag, quantity)
            )

    def try_liquidate(self) -> None:
        # decrement holding period and liquidate symbols
        symbols_to_delete: List[ManagedSymbol] = []

        for managed_symbol in self._symbols:
            managed_symbol.days_to_liquidate -= 1
            
            # liquidate
            if managed_symbol.days_to_liquidate == 0:
                self._algorithm.log(f'Liquidating {managed_symbol.symbol} with quantity {managed_symbol.quantity}')
                symbols_to_delete.append(managed_symbol)
                
                if managed_symbol.long_flag: 
                    self._algorithm.market_order(managed_symbol.symbol, -managed_symbol.quantity)
                    self._long_len -= 1
                else: 
                    self._short_len -= 1
                    self._algorithm.market_order(managed_symbol.symbol, managed_symbol.quantity)

        # remove symbols from management
        for managed_symbol in symbols_to_delete:
            self._symbols.remove(managed_symbol)

@dataclass   
class ManagedSymbol():
    symbol: Symbol 
    days_to_liquidate: int 
    long_flag: bool 
    quantity: int