from AlgorithmImports import *
from pandas.tseries.offsets import BDay
from pandas.tseries.offsets import BMonthEnd
from object_store_helper import ObjectStoreHelper
from typing import Any, List, Dict
from traded_strategy import TradedStrategy

class MetatronShortVolatilityStrategy(QCAlgorithm):
    
    _notional_value: float = 200_000
    _trade_exec_minute_offset: int = 15

    _sma_period: int = 20

    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

        leverage: int = 4

        traded_ticker: str = 'VXX' if self.live_mode else 'VIXY'
        security: Security = self.add_equity(traded_ticker, Resolution.MINUTE, leverage=leverage)
        self._traded_asset: Symbol = security.symbol
        
        self.log(f'Traded asset: {self._traded_asset.value}')

        # holding days setting
        self._before_expiration_delta_days: int = -4
        self._after_expiration_delta_days: int = 0
        assert self._before_expiration_delta_days < self._after_expiration_delta_days, 'delta days are not aligned properly'
        
        security.sma: SimpleMovingAverage = self.sma(self._traded_asset, self._sma_period, Resolution.DAILY)
        security.sma.updated += self.sma_updated
        self.last_market_price: float|None = None

        # load expiration days
        file: str = self.download('data.quantpedia.com/backtesting_data/calendar/vix_futures_expiration.csv')
        self.vix_expiration_dates: List[datetime.date] = list(
            map(lambda x: datetime.strptime(x, '%Y-%m-%d').date(), file.split('\r\n')[1:])
        )

        self.vix_symbols: List[Symbol] = [
            self.add_data(CBOE, 'VIX', Resolution.DAILY).symbol, 
            self.add_data(CBOE, 'VIX3M', Resolution.DAILY).symbol
        ]

        self.position_opened_this_month: bool = False

        self.recent_month: int = -1
        self.market_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.market_close
        )

        self.set_warmup(self._sma_period, Resolution.DAILY)
    
    def sma_updated(self, sender: SimpleMovingAverage, datapoint: IndicatorDataPoint) -> None:
        self.last_market_price = self.securities[self._traded_asset].price

    def close_volatility_position(self) -> None:
        if self.portfolio[self._traded_asset].invested:
            self.market_order(self._traded_asset, -self.portfolio[self._traded_asset].quantity)

    def market_close(self) -> None:
        self.market_close_flag = True

    def on_data(self, slice: Slice) -> None:
        if not (slice.ContainsKey(self._traded_asset) and slice[self._traded_asset] is not None) or \
            not all(self.securities.contains_key(symbol) and self.securities[symbol].get_last_data() for symbol in self.vix_symbols) or \
            not self.securities[self._traded_asset].sma.is_ready or self.last_market_price is None or \
            self.is_warming_up:
            return
        
        if not self.market_close_flag:
            return
        self.market_close_flag = False

        # get 1-day lagged VIX signal
        vix_in_contango: bool = self.securities[self.vix_symbols[1]].price / self.securities[self.vix_symbols[0]].price > 1.

        # check VIX ratio every day
        if self.portfolio.invested and not vix_in_contango:
            self.log(f'Liquidation signal triggered. VIX in contango: {vix_in_contango}. Terminating position')
            self.close_volatility_position()

        # new month -> reset flag
        if self.recent_month != self.time.month:
            self.position_opened_this_month: bool = False
            self.recent_month = self.time.month
        
        open_day_range: List[int] = list(range(self._before_expiration_delta_days, self._after_expiration_delta_days))

        # open position
        if any((self.time.date() + BDay((-1 * x) + 1)).date() in self.vix_expiration_dates for x in open_day_range): 
            self.log('Trying to open position')

            if vix_in_contango:
                self.log(f'VIX in contango: {vix_in_contango}')

                if not self.portfolio[self._traded_asset].invested:
                    if not self.position_opened_this_month:
                        above_sma: bool = self.last_market_price > self.securities[self._traded_asset].sma.current.value
                        if above_sma:
                            self.log(f'Last Market Price above SMA: {above_sma}')
                            self.log(f'Opening position')

                            quantity: int = self._notional_value // slice[self._traded_asset].price
                            self.market_order(self._traded_asset, -quantity)
                            
                            # NOTE handle sudden changes of term structure
                            # highly improbable, yet possible
                            self.position_opened_this_month = True

        # close day
        elif (self.time + BDay(-1 * self._after_expiration_delta_days)).date() in self.vix_expiration_dates:
            self.log('Upcoming expiration date. Terminating position')
            self.close_volatility_position()

# region imports
from AlgorithmImports import *
import json
from traded_strategy import TradedStrategy
# endregion

class ObjectStoreHelper:
    def __init__(
        self, 
        algorithm: QCAlgorithm, 
        path: str
    ) -> None:
        """
        Initializes ObjectStoreHelper with reference to the algorithm instance.
        """
        self._algorithm: QCAlgorithm = algorithm
        self._path: str = path

    def save_state(self, state: Dict) -> None:
        """
        Saves a dictionary `state` to the Object Store as JSON.
        """
        if not self._algorithm.live_mode:
            return

        json_data = json.dumps(state)
        self._algorithm.object_store.save(self._path, json_data)
        self._algorithm.log(f"Saved state to Object Store: {json_data}")

    def load_state(self) -> Dict:
        """
        Loads a JSON string from the Object Store and returns it as a dictionary.
        """
        if self._algorithm.object_store.contains_key(self._path) and self._algorithm.live_mode:
            json_data = self._algorithm.object_store.read(self._path)
            if json_data:
                self._algorithm.log(f"Loaded state from Object Store: {json_data}")
                result: Dict = json.loads(json_data)
                result['trade_signal'] = {TradedStrategy._member_map_[key]: value for key, value in result['trade_signal'].items() if key in TradedStrategy._member_map_}
                return result
        else:
            return {
                'trade_signal': {
                    TradedStrategy.CALENDAR: False,
                    TradedStrategy.REVERSAL_MODEL: False
                },
                'reversal_model_days_held': 0
            }

        return {}

# region imports
from AlgorithmImports import *
from enum import Enum
# endregion

class TradedStrategy(Enum):
    CALENDAR = 1
    REVERSAL_MODEL = 2