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

class MetatronTurnaroundTuesday(QCAlgorithm):
    
    _notional_value: float = 400_000
    _trade_exec_minute_offset: int = 1
    _traded_weight: float = .5

    _history_period: int = 3
    _consecutive_days_period: int = 2
    _order_liquidation_date: Optional[datetime] = None

    _volatility_flag: bool = False
    _volatility_threshold: int = 16

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

        leverage: int = 4

        tickers: List[str] = ['HYG']#, 'VOO']

        self._traded_assets: Symbol = [
            self.add_equity(ticker, Resolution.MINUTE, leverage=leverage).symbol for ticker in tickers
        ]

        if self._volatility_flag:
            self._volatility: Symbol = self.add_data(CBOE, 'VIX', Resolution.DAILY).symbol

        self._trade_flag: bool = False
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.before_market_close(self._traded_assets[0], self._trade_exec_minute_offset), 
            self._selection
        )
        
        self.settings.daily_precise_end_time = False

    def on_data(self, slice: Slice) -> None:
        if self._order_liquidation_date:
            if self.portfolio[self._traded_assets[1]].invested and self.time >= self._order_liquidation_date:
                self.market_order(
                    self._traded_assets[1], 
                    -self.portfolio[self._traded_assets[1]].quantity
                )
                self._order_liquidation_date = None

        if not self._trade_flag:
            return
        self._trade_flag = False

        if all(slice.contains_key(symbol) and slice[symbol] for symbol in self._traded_assets):
            history: DataFrame = self.history(
                TradeBar, self._traded_assets, start=self.time - BDay(self._history_period), end=self.time
            )

            if not history.empty:
                closes: DataFrame = history.unstack(level=0).resample('B').last().close[-self._history_period:]
                if len(closes) < self._history_period:
                    return
                
                # Trade if asset had last down day.
                if closes[self._traded_assets[0]][-1] / closes[self._traded_assets[0]][-2] - 1 < 0:
                    self.market_order(
                        self._traded_assets[0], 
                        # self._notional_value * self._traded_weight // slice[self._traded_assets[0]].price,
                        self.calculate_order_quantity(self._traded_assets[0], 1),
                        tag='MarketOrderHYG'
                    )

                # Trade if there were last two consecutive down days.
                # if all(closes[self._traded_assets[1]][i] > closes[self._traded_assets[1]][i+1] for i in range(self._consecutive_days_period)):
                #     if self._volatility_flag:
                #         if self.securities[self._volatility].get_last_data().price > self._volatility_threshold:
                #             return

                    # self.market_order(
                    #     self._traded_assets[1], 
                    #     self._notional_value * self._traded_weight // slice[self._traded_assets[1]].price,
                    #     tag='MarketOrderVOO'
                    # )

    def on_order_event(self, orderEvent: OrderEvent) -> None:
        order_ticket: OrderTicker = self.transactions.get_order_ticket(orderEvent.order_id)
        symbol: Symbol = order_ticket.symbol

        if orderEvent.status == OrderStatus.FILLED:
            if 'MarketOrderHYG' in order_ticket.tag:
                self.market_on_open_order(symbol, -order_ticket.quantity) 
            if 'MarketOrderVOO' in order_ticket.tag:
                # Set date for liquidation.
                self._order_liquidation_date = self.securities[symbol].exchange.hours.get_next_trading_day(self.time)

    def _selection(self) -> None:
        if all(self.securities[symbol].get_last_data() for symbol in self._traded_assets):
            self._trade_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