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

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

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:
        
        # csv parsing
        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

##################
#  ECB Strategy  #
##################

@dataclass
class TradePosition:
    buy_date: date
    sell_date: date
    ticket: OrderTicket | None = None

        
class ECBMeetingsData(PythonData):
    _ecb_holding_period: int

    @staticmethod
    def set_holding_period(period: int) -> None:
        ECBMeetingsData._ecb_holding_period = period

    def get_source(self, config: SubscriptionDataConfig, date: datetime | date, isLiveMode: bool) -> SubscriptionDataSource:
        url: str = 'https://data.quantpedia.com/backtesting_data/calendar/ECB_MEETINGS.csv'
        return SubscriptionDataSource(url, SubscriptionTransportMedium.REMOTE_FILE)

    def reader(self, config: SubscriptionDataConfig, line: str, date: datetime | date, isLiveMode: bool) -> BaseData:
        data = ECBMeetingsData()
        data.symbol = config.symbol

        if not line[0].isdigit(): return None
        split = line.split(';')
        
        data.time = datetime.strptime(split[0], "%Y-%m-%d") - BDay(self._ecb_holding_period + 1)
        
        return data


class CustomFeeModel(FeeModel):
    def get_order_fee(self, parameters: OrderFeeParameters) -> OrderFee:
        fee = parameters.security.price * parameters.order.absolute_quantity * 0.00001
        return OrderFee(CashAmount(fee, "USD"))

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, TradePosition, CustomFeeModel, ECBMeetingsData

class MetatronFEDDayPlusTOM(QCAlgorithm):
    
    _notional_value: float = 1_000_000
    _trade_exec_minute_offset: int = 15
    _observed_period: int = 20
    _traded_weight: float = .33
    _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

    # ECB Strategy
    _ecb_holding_period: int = 1 # Day

    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

        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(self._observed_period, Resolution.DAILY)

        # ECB init
        self._ecb_positions: list[TradePosition] = []        
        self._ewg_security: Security = self.add_equity("EWG", Resolution.MINUTE)
        self._ewg_security.set_fee_model(CustomFeeModel())
        self._ewg: Symbol = self._ewg_security.symbol
        self._ecb_meetings: Symbol = self.add_data(ECBMeetingsData, 'ecb_meetings', Resolution.DAILY).symbol
        ECBMeetingsData.set_holding_period(self._ecb_holding_period)

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

        self.schedule.on(self.date_rules.every_day(self._ewg),
                    self.time_rules.before_market_close(self._ewg, self._trade_exec_minute_offset),
                    self._trade_ecb)

    def on_data(self, slice: Slice) -> None:
        if self.is_warming_up: return
        if slice.contains_key(self._ecb_meetings) and slice[self._ecb_meetings]:
            ecb_date: datetime = slice[self._ecb_meetings].time
            self._ecb_positions.append(TradePosition(ecb_date.date(), (ecb_date + BDay(self._ecb_holding_period)).date()))
            self.log(f"New ECB buying date arrived: {ecb_date}. Active ECB dates: {len(self._ecb_positions)} (should == 1)")

        if not slice.contains_key(self._traded_asset):
            return

        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
        asset_security: Security = self.securities[self._traded_asset]

        if self._indicator_flag:
            if asset_security.mom.is_ready:
                trade_flag = True if asset_security.mom.current.value > 0 else False  
                self.log(f'{asset_security.symbol} momentum value: {asset_security.mom.current.value}; signal: {asset_security.mom.current.value > 0}')
            else:
                self.log(f'Momentum indicator is not ready for {asset_security.symbol}')
        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

    def _trade_ecb(self) -> None:
        # Open position
        for position in self._ecb_positions:
            if position.buy_date == self.time.date():
                quantity_ewg: float = self._notional_value * self._traded_weight // self.securities[self._ewg].price
                buy_ticket = self.market_order(self._ewg, quantity_ewg)
                position.ticket = buy_ticket
                self.log(f"Open ECB order | Quantity filled: {buy_ticket.quantity_filled}")

        # Close position
        positions_to_remove = []
        for position in self._ecb_positions:
            if self.time.date() >= position.sell_date:
                if position.ticket and self.portfolio[self._ewg].invested:
                    sell_ticket = self.market_order(self._ewg, -self.portfolio[self._ewg].quantity)
                    self.log(f"Close ECB order | Quantity filled: {sell_ticket.quantity_filled}")
                positions_to_remove.append(position)
        
        for position in positions_to_remove:
            self._ecb_positions.remove(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):
    FED_DAYS = 1
    TOM = 2