from AlgorithmImports import *
import pandas as pd
from sklearn.tree import DecisionTreeClassifier

class WeeklyTreeBTC(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2025, 1, 1)
        self.SetCash(100000)

        # Subscribe to BTC and DOGE minute data
        self.btc = self.AddCrypto("BTCUSDT", Resolution.Minute, Market.Binance, True).Symbol
        self.doge = self.AddCrypto("DOGEUSDT", Resolution.Minute, Market.Binance, True).Symbol
        self.eth = self.AddCrypto("ETHUSDT", Resolution.Minute, Market.Binance, True).Symbol

        # Warm-up
        self.SetWarmUp(24*60, Resolution.Minute)

        # Top features
        self.top_features = [
            'BTCUSDT_hour_asia', 'BTCUSDT_hour_europe', 
            'ETHUSDT_rel_vol', 'BTCUSDT_rel_vol_ytd', 
            'ETHUSDT_rel_vol_ytd', 'ETHUSDT_vma5', 
            'DOGEUSDT_rel_vol_ytd', 'DOGEUSDT_vol_rank_15m'
        ]

        # Model placeholder
        self.model = None

        # Schedule weekly retraining
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), 
                         self.TimeRules.At(0,0), 
                         self.TrainWeeklyTree)

    def ComputeFeatures(self, hist, time):
        """Compute feature vector for top_features at given time"""
        # Reset index for safety
        hist_reset = hist.reset_index()
        features = {}

        # BTC hour features
        features['BTCUSDT_hour_asia'] = time.hour
        features['BTCUSDT_hour_europe'] = (time + pd.Timedelta(hours=6)).hour

        # ETH relative volume
        eth_hist = hist_reset[hist_reset['symbol'] == str(self.eth)]
        if not eth_hist.empty:
            eth_vol = eth_hist['volume'].mean()
            eth_bar = eth_hist[eth_hist['time'] == time]
            features['ETHUSDT_rel_vol'] = eth_bar['volume'].values[0] / eth_vol if eth_vol>0 else 0
            features['ETHUSDT_rel_vol_ytd'] = eth_bar['volume'].values[0] / eth_vol if eth_vol>0 else 0
            features['ETHUSDT_vma5'] = eth_hist.loc[eth_hist['time'] <= time].tail(5)['close'].mean()
        else:
            features['ETHUSDT_rel_vol'] = 0
            features['ETHUSDT_rel_vol_ytd'] = 0
            features['ETHUSDT_vma5'] = 0

        # BTC relative volume YTD
        btc_hist = hist_reset[hist_reset['symbol'] == str(self.btc)]
        if not btc_hist.empty:
            btc_vol = btc_hist['volume'].mean()
            btc_bar = btc_hist[btc_hist['time'] == time]
            features['BTCUSDT_rel_vol_ytd'] = btc_bar['volume'].values[0] / btc_vol if btc_vol>0 else 0

        # DOGE features
        doge_hist = hist_reset[hist_reset['symbol'] == str(self.doge)]
        if not doge_hist.empty:
            doge_vol = doge_hist['volume'].mean()
            doge_bar = doge_hist[doge_hist['time'] == time]
            features['DOGEUSDT_rel_vol_ytd'] = doge_bar['volume'].values[0] / doge_vol if doge_vol>0 else 0
            # vol_rank_15m = relative rank over last 15 minutes
            last_15 = doge_hist.loc[(doge_hist['time'] <= time)].tail(15)
            features['DOGEUSDT_vol_rank_15m'] = (doge_bar['volume'].values[0] / (last_15['volume'].max() if not last_15.empty else 1))

        return pd.DataFrame([features])

    def TrainWeeklyTree(self):
        # Pull 1 week of minute history
        hist = self.History([self.btc, self.doge, self.eth], 7*24*60, Resolution.Minute, dataNormalizationMode=DataNormalizationMode.Raw)
        if hist.empty:
            return

        hist = hist.reset_index()
        X_list = []
        y_list = []

        for time in hist['time'].unique()[:-1]:  # skip last bar (no next)
            features = self.ComputeFeatures(hist, time)
            # label = 1 if next BTC close > current BTC close
            btc_hist = hist[hist['symbol'] == str(self.btc)]
            btc_bar_idx = btc_hist[btc_hist['time'] == time].index
            if len(btc_bar_idx)==0 or btc_bar_idx[0] >= len(btc_hist)-1:
                continue
            current_close = btc_hist.loc[btc_bar_idx[0], 'close']
            next_close = btc_hist.loc[btc_bar_idx[0]+1, 'close']
            label = 1 if next_close > current_close else 0
            X_list.append(features.values[0])
            y_list.append(label)

        if len(X_list) == 0:
            return

        X = pd.DataFrame(X_list, columns=self.top_features)
        y = pd.Series(y_list)

        # Train decision tree
        self.model = DecisionTreeClassifier(max_depth=3)
        self.model.fit(X, y)
        self.Debug(f"Weekly tree trained at {self.Time}")

    def OnData(self, data):
        if self.model is None:
            return

        # Pull latest history for features
        hist = self.History([self.btc, self.doge, self.eth], 60, Resolution.Minute)
        if hist.empty:
            return

        features = self.ComputeFeatures(hist, hist.index.get_level_values('time')[-1])
        pred = self.model.predict(features)[0]

        # Execute BTC trade based on prediction
        if pred == 1:
            self.SetHoldings(self.btc, 1.0)  # full allocation
        else:
            self.Liquidate(self.btc)