# Aggregate strategy:
#   1. TLT Calendar:
#       TLT is traded 2 days before the EOM. 
#       Signal is set on a third day before the EOM 15 minutes before market close and trade is held until EOM (last traded day of the month, 15 minutes before market close).
#
#   2. TLT Reversal Model:
#       TLT is bought after 3 down days (15 minutes before market close) and is held for 3 days.

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 MetatronTLTCalendarPlusReversalModel(QCAlgorithm):
    
    _notional_value: float = 400_000
    _trade_calendar: bool = True
    _trade_reversal_model: bool = True
    _trade_exec_minute_offset: int = 15
    
    # Calendar Strategy
    _eom_trigger_day_offset: int = 2

    # Reversal Strategy
    _reversal_days_lookback: int = 3
    _holding_period: int = 3

    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

        # initialize Object Store Helper
        self.object_store_helper: ObjectStoreHelper = ObjectStoreHelper(
            self,
            "Metatron/1_TLT_Calendar_plus_Reversal_Model/trade_signal.json"
        )

        # load state from Object Store if available
        signal_state: Dict = self.object_store_helper.load_state()

        self._trade_signal: Dict[TradedStrategy, bool] = signal_state['trade_signal']
        self._reversal_model_days_held = signal_state['reversal_model_days_held']

        self._traded_asset: Symbol = self.add_equity('TLT', Resolution.MINUTE).symbol

        # schedule functions
        self._month_close_flag: bool = False
        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
        )
        self.schedule.on(
            self.date_rules.month_end(self._traded_asset), 
            self.time_rules.before_market_close(self._traded_asset, self._trade_exec_minute_offset), 
            self._month_close
        )
    
    def on_data(self, slice: Slice) -> None:
        if not self._day_close_flag:
            return
        
        self.log(f'Day close flag set')
        self._day_close_flag = False

        # reset Calendar Strategy signal
        if self._trade_calendar:
            if self._month_close_flag:
                self.log(f'Month close signal check')
                self._month_close_flag = False

                if self._trade_signal[TradedStrategy.CALENDAR]:
                    self.log(f'Signal reset for {TradedStrategy.CALENDAR}')
                    self._trade_signal[TradedStrategy.CALENDAR] = False
                                
        # reset Reversal Model signal
        if self._trade_reversal_model:
            if self._trade_signal[TradedStrategy.REVERSAL_MODEL]:
                self._reversal_model_days_held += 1
                self.log(f'Days held counter incremented for {TradedStrategy.REVERSAL_MODEL}. Days held: {self._reversal_model_days_held}')

                if self._reversal_model_days_held == self._holding_period:
                    self.log(f'Signal reset for {TradedStrategy.REVERSAL_MODEL}. Days held count reached {self._reversal_model_days_held}/{self._holding_period} days')
                    self._reversal_model_days_held = 0
                    self._trade_signal[TradedStrategy.REVERSAL_MODEL] = False

        # close position
        if all(not signal for signal in self._trade_signal.values()):
            if self.portfolio[self._traded_asset].invested:
                self.log(f'No active signal, liquidating portfolio')
                self.liquidate(self._traded_asset)
                
        if slice.contains_key(self._traded_asset) and slice[self._traded_asset]:
            self.log(f'Symbol {self._traded_asset} data present in the algorithm')
            if self._trade_reversal_model:
                # Reversal Model Strategy
                history: DataFrame = self.history[TradeBar](
                    self._traded_asset, 
                    self._reversal_days_lookback, 
                    Resolution.DAILY
                )
                closes: List[float] = [bar.close for bar in history] + [slice[self._traded_asset].close]
                is_in_consecutive_downtrend: bool = all(closes[i] > closes[i + 1] for i in range(len(closes) - 1))
                if is_in_consecutive_downtrend:
                    if not self._trade_signal[TradedStrategy.REVERSAL_MODEL]:
                        self.log(f'Positive signal for {TradedStrategy.REVERSAL_MODEL}')
                        self._trade_signal[TradedStrategy.REVERSAL_MODEL] = True

            if self._trade_calendar:
                # Calendar Strategy
                offset = BMonthEnd()
                last_day_of_month: datetime = offset.rollforward(self.time)

                # find last trading day of the month
                while not self.securities[self._traded_asset].exchange.hours.is_date_open(last_day_of_month):
                    last_day_of_month = last_day_of_month - timedelta(days=1)

                # add one day to account for before the close execution
                trigger_day: datetime = last_day_of_month - BDay(self._eom_trigger_day_offset)
                if self.time == trigger_day:
                    if not self._trade_signal[TradedStrategy.CALENDAR]:
                        self.log(f'Positive signal for {TradedStrategy.CALENDAR}')
                        self._trade_signal[TradedStrategy.CALENDAR] = True
        else:
            self.log(f'Symbol {self._traded_asset} data NOT present in the algorithm')

        if any(signal for signal in self._trade_signal.values()):
            if not self.portfolio[self._traded_asset].invested:
                self.log(f'Trade executed for {self._traded_asset}')
                quantity: int = self._notional_value // self.securities[self._traded_asset].ask_price
                self.market_order(self._traded_asset, quantity)

        # NOTE do not save portfolio state for now, liquidate live positions every time live restart is needed
        # save state to Object Store
        # signal_state: Dict[str, Any] = {
        #     "trade_signal": {
        #         TradedStrategy.CALENDAR.name: self._trade_signal[TradedStrategy.CALENDAR],
        #         TradedStrategy.REVERSAL_MODEL.name: self._trade_signal[TradedStrategy.REVERSAL_MODEL]
        #     },
        #     "reversal_model_days_held": self._reversal_model_days_held
        # }
        # self.object_store_helper.save_state(signal_state)        

    def _before_eod(self) -> None:
        self._day_close_flag = True
    
    def _month_close(self) -> None:
        self._month_close_flag = True

# 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