# region imports
from AlgorithmImports import *
from pandas.tseries.offsets import BDay
from dataclasses import dataclass
# endregion

@dataclass
class HoldingItem():
    quantity: int
    holding_period: int = 0
    FED_DAYS_flag: bool = False

# Your New Python File
class FedDays(PythonData):
    algo = None
    
    @staticmethod
    def set_algo(algo) -> None:
        FedDays.algo = algo

    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        if isLiveMode:
            # FedDays.algo.Log(f"Edited GetSource date {FedDays.algo.Time}")
            return SubscriptionDataSource("https://data.quantpedia.com/backtesting_data/economic/fed_days.json", SubscriptionTransportMedium.RemoteFile, FileFormat.UnfoldingCollection)

        return SubscriptionDataSource("https://data.quantpedia.com/backtesting_data/economic/fed_days.csv", SubscriptionTransportMedium.RemoteFile)
        
    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, isLiveMode: bool) -> BaseData:
        if isLiveMode:
            try:
                # FedDays.algo.Log(f"Reader")

                objects = []
                data = json.loads(line)
                end_time = None

                for index, sample in enumerate(data):
                    custom_data = FedDays()
                    custom_data.Symbol = config.Symbol

                    custom_data.Time = (datetime.strptime(str(sample["fed_date"]), "%Y-%m-%d") - BDay(1)).replace(hour=9, minute=31)
                    # FedDays.algo.Log(f"{custom_data.Time}")

                    end_time = custom_data.Time  
                    objects.append(custom_data)

                return BaseDataCollection(end_time, config.Symbol, objects) 

            except ValueError:
                # FedDays.algo.Log(f"Reader Error")
                return None
        else:
            if not (line.strip() and line[0].isdigit()):
                return None
            
            custom = FedDays()
            custom.Symbol = config.Symbol

            custom.Time = (datetime.strptime(line, "%Y-%m-%d") - BDay(1)).replace(hour=9, minute=31)
            custom.Value = 0.
            custom["fed_date_str"] = line

            return custom

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
from data_tools import FedDays, HoldingItem

class MetatronFEDDayPlusTOMDay(QCAlgorithm):
    
    _notional_value: float = 600_000
    _trade_exec_minute_offset: int = 15
    _observed_period: int = 20
    _traded_weight: float = .5
    _holding_manager: List[HoldingItem] = []

    # True -> Use indicator logic
    _indicator_flag: bool = True
    _traded_asset: str = 'SPY'
    
    # FED days Strategy
    _close_hour: int = 12

    # ToM Strategy
    _holding_period: int = 2

    def initialize(self) -> None:
        self.set_start_date(2000, 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

        security: Security = self.add_equity(self._traded_asset, Resolution.MINUTE)
        security._mom: Momentum = self.mom(self._traded_asset, self._observed_period, Resolution.DAILY)

        self._traded_asset: Symbol = security.symbol
        self._fed_days_symbol: Symbol = self.add_data(FedDays, 'fed_days', Resolution.DAILY, TimeZones.NEW_YORK).symbol

        FedDays.set_algo(self)
        self.set_warm_up(1, Resolution.DAILY)

        # 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 slice.contains_key(self._traded_asset):
            return

        asset_security: Security = self.securities[self._traded_asset]
        items_to_delete: List[HoldingItem] = []

        # Liquidate.
        if self.portfolio[self._traded_asset].invested:
            for holding_item in self._holding_manager:
                if holding_item.FED_DAYS_flag:
                    if self.time.hour == self._close_hour:
                        self.log(f"FOMC meeting day; submitting an market order to close opened position...")

                        self.market_order(self._traded_asset, -holding_item.quantity)
                        items_to_delete.append(holding_item)
                else:
                    if self._day_close_flag:
                        holding_item.holding_period -= 1
                        if holding_item.holding_period == 0:
                            self.log(f'Liquidating opened position after holding {self._holding_period} days.')

                            self.market_order(self._traded_asset, -holding_item.quantity)
                            items_to_delete.append(holding_item)
                    
        for item in items_to_delete:
            self._holding_manager.remove(item)

        if self.is_warming_up: return
        if self._day_close_flag:
            self._day_close_flag = False
            if self.securities.contains_key(self._fed_days_symbol) and self.securities[self._fed_days_symbol].get_last_data():
                if self.time.date() == self.securities[self._fed_days_symbol].get_last_data().time.date():
                    # New fed day data arrived.
                    if self.securities[self._traded_asset].get_last_data():
                        self.log(f"New FOMC meeting data arrived: {self.time}; submitting an market order...")

                        quantity: int = self._notional_value * self._traded_weight // slice[self._traded_asset].price
                        self.market_order(self._traded_asset, quantity)
                        self._holding_manager.append(HoldingItem(quantity, FED_DAYS_flag=True))

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

        trade_flag: bool = False

        if self._indicator_flag:
            if asset_security._mom.is_ready:
                trade_flag = True if asset_security._mom.current.value > 0 else False  
        else:
            trade_flag = True  

        if trade_flag:    
            # Turn of the month.
            self.log(f'Turn of the month; submiting market order...')   

            quantity: int = self._notional_value * self._traded_weight // slice[self._traded_asset].price
            self.market_order(self._traded_asset, quantity)
            self._holding_manager.append(HoldingItem(quantity, holding_period=self._holding_period))

    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):
    FED_DAYS = 1
    TOM = 2