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

# 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
# endregion

class MetatronGasolinePreHolidayEffect(QCAlgorithm):

    _notional_value: float = 200_000
    _trade_exec_minute_offset: int = 15
    _traded_asset: str = 'UGA'

    def initialize(self) -> None:
        self.set_start_date(2020, 1, 1)
        self.set_cash('USD', 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 = False
        
        self._offset_days: List[int] = [1, 5]

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

        self._holidays_dates: List[datetime.date] = list(
            map(lambda x: x.date(), list(self.securities[self._traded_asset].exchange.hours.holidays))
        )

        self._rebalance_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_close
        )

    def on_data(self, slice: Slice) -> None:
        if not self._rebalance_flag:
            return
        self._rebalance_flag = False

        if slice.contains_key(self._traded_asset) and slice[self._traded_asset]:
            if self.portfolio[self._traded_asset].invested:
                # close position 1 day before holidays
                if (self.time + BDay(self._offset_days[0])).date() in self._holidays_dates:
                    self.liquidate()

            else:
                # open position 5 days before holidays
                if (self.time + BDay(self._offset_days[1])).date() in self._holidays_dates:
                    quantity: int = self._notional_value // slice[self._traded_asset].price
                    self.market_order(self._traded_asset, quantity)

    def before_close(self) -> None:
        self._rebalance_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