from AlgorithmImports import *
from parser import MarketRegimeSignal


class MarketRegimeAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2020, 3, 2)
        self.set_end_date(2020, 3, 13)
        self.set_cash(100_000)
        self._spy = self.add_equity("SPY", Resolution.DAILY)
        # Store the symbol to reference it later
        self._regime_symbol = self.add_data(MarketRegimeSignal, "REGIME", Resolution.DAILY).symbol
        self.log(f"Initialized Custom Data: {self._regime_symbol.value}")

    def on_data(self, slice: Slice) -> None:
        self.log(f"on_data: {self.time}")
        slice_keys = [str(key) for key in slice.keys()]
        self.log(f"slice has: {slice_keys}")
        # Direct string check fallback
        if "REGIME" in slice:
            self.log("Success")
            # Extract data cleanly
            data = slice[self._regime_symbol]
            self.debug(f"Check data: {data.time.date()} | Regime: {data.regime} | Score: {data.score}")
            # Trading logic
            if data.regime == "risk_on":
                self.set_holdings(self._spy, 1.0)
            elif data.regime == "risk_off":
                self.liquidate(self._spy)

from AlgorithmImports import *

class MarketRegimeSignal(PythonData):
    def get_source(self, config, date, is_live_mode):
        return SubscriptionDataSource(
            "market-regime-signals.zip#signals.csv", 
            SubscriptionTransportMedium.OBJECT_STORE,
            FileFormat.CSV
        )

    def reader(self, config, line, date, is_live_mode):
        # Inspect the incoming raw line strings from the object store
        print(f"READER RECEIVED LINE: {line}")
        if "date" in line or not line.strip():
            return None
        csv = line.split(',')
        if len(csv) < 6:
            return None
        data = MarketRegimeSignal()
        data.symbol = config.symbol
        data.time = datetime.strptime(csv[0].strip(), "%Y-%m-%d")
        data.end_time = data.time + timedelta(days=1)
        data.ticker = csv[1].strip()
        data.score = float(csv[2].strip())
        data.regime = csv[3].strip()
        data.is_confirmed = csv[4].strip().lower() == 'true'
        data.note = csv[5].strip()
        data.value = data.score 
        print(f"READER SUCCESSFULLY PARSED: Date={data.time}, Regime={data.regime}")
        return data