from AlgorithmImports import *
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import pandas as pd

class BTCUSDTLinearRegressionAlgo(QCAlgorithm):
    
    def Initialize(self):
        # 1. Algorithm Parameters
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash("USDT", 10000)

        # Optional: Set Binance brokerage model for a cash account
        self.SetBrokerageModel(BrokerageName.Binance, AccountType.Cash)

        # 2. Add BTC/USDT from Binance, Daily Resolution (produces QuoteBars)
        self.symbol = self.AddCrypto("BTCUSDT", Resolution.Daily, Market.Binance).Symbol
        
        # 3. RollingWindow to store the last 200 QuoteBars
        self.data = RollingWindow[QuoteBar](200)
        
        # 4. Warm-up to get 200 bars before starting trading
        self.SetWarmUp(200)

        # 5. Initialize ML Model
        self.model = LinearRegression()
        self.training_count = 0
        self.is_model_trained = False

        # 6. Configurable parameter: fraction of available cash to spend when buying
        #    Adjust this to 0.90 (90%), 0.99 (99%), etc., as you see fit
        self.allocationFraction = 0.50
        
        # 7. Configurable Training Frequency: "Daily" or "4H" or "6H"
        self.trainingFrequency = "4H"  # Train 4 times/day
        self.ConfigureTrainingSchedule()

    def ConfigureTrainingSchedule(self):
        """Schedules the TrainModel() calls according to the desired frequency."""
        if self.trainingFrequency == "Daily":
            # Train once per day at 00:00 UTC
            self.Schedule.On(
                self.DateRules.EveryDay(self.symbol),
                self.TimeRules.At(0, 0),
                self.TrainModel
            )
        elif self.trainingFrequency == "6H":
            # Train every 6 hours: 00:00, 06:00, 12:00, 18:00 (UTC)
            for hour in [0, 6, 12, 18]:
                self.Schedule.On(
                    self.DateRules.EveryDay(self.symbol),
                    self.TimeRules.At(hour, 0),
                    self.TrainModel
                )
        elif self.trainingFrequency == "4H":
            # Train every 4 hours: 00:00, 04:00, 08:00, 12:00, 16:00, 20:00 (UTC)
            for hour in [0, 4, 8, 12, 16, 20]:
                self.Schedule.On(
                    self.DateRules.EveryDay(self.symbol),
                    self.TimeRules.At(hour, 0),
                    self.TrainModel
                )
        else:
            self.Debug(f"Unknown training frequency: {self.trainingFrequency}. Using daily by default.")
            self.Schedule.On(
                self.DateRules.EveryDay(self.symbol),
                self.TimeRules.At(0, 0),
                self.TrainModel
            )

    def OnData(self, data):
        # 8. Ensure we have QuoteBars for the symbol
        if not data.QuoteBars.ContainsKey(self.symbol):
            return
        
        quote_bar = data.QuoteBars[self.symbol]
        if quote_bar is None:
            return
        
        # 9. Add the QuoteBar to our rolling window
        self.data.Add(quote_bar)

        # 10. Wait until rolling window is ready
        if not self.data.IsReady or self.data.Count < 200:
            return

        # 11. Skip predictions if the model is not trained
        if not self.is_model_trained:
            self.Debug("Model not trained yet. Skipping prediction.")
            return

        # 12. Retrieve the latest feature vector
        df = self.GetFeatureDataFrame()
        if df is None or len(df) < 1:
            return
        
        latest_features = df.iloc[-1, :-1].values.reshape(1, -1)
        
        # 13. Prediction: continuous → binary threshold at 0.5
        try:
            prediction_value = self.model.predict(latest_features)[0]
            prediction = 1 if prediction_value > 0.5 else 0
        except Exception as e:
            self.Debug(f"Model prediction failed: {e}")
            return

        # 14. Trading Logic: manually calculate quantity
        holdings = self.Portfolio[self.symbol].Quantity
        current_price = quote_bar.Close  # mid-price from the QuoteBar

        if prediction == 1 and not self.Portfolio.Invested:
            # Buy up to allocationFraction of current USDT
            available_usdt = self.Portfolio.CashBook["USDT"].Amount * self.allocationFraction
            if available_usdt > 0 and current_price > 0:
                quantity_to_buy = available_usdt / current_price
                self.MarketOrder(self.symbol, quantity_to_buy)

        elif prediction == 0 and self.Portfolio.Invested:
            # If we have BTC and the model says sell, liquidate
            self.Liquidate(self.symbol)

    def TrainModel(self):
        # 15. Prepare data for training
        df = self.GetFeatureDataFrame()
        if df is None or len(df) < 50:
            self.Debug("Insufficient data for training.")
            return

        # Separate features (X) and label (y)
        X = df.iloc[:, :-1]
        y = df.iloc[:, -1]
        
        # Chronological split: 80% train, 20% test
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, shuffle=False, random_state=42
        )

        # 16. Train the Linear Regression model
        self.model.fit(X_train, y_train)
        self.is_model_trained = True

        # 17. Evaluate performance (threshold regression to binary for accuracy)
        y_train_pred = self.model.predict(X_train)
        y_train_pred_binary = [1 if val > 0.5 else 0 for val in y_train_pred]
        train_accuracy = accuracy_score(y_train, y_train_pred_binary)

        y_test_pred = self.model.predict(X_test)
        y_test_pred_binary = [1 if val > 0.5 else 0 for val in y_test_pred]
        test_accuracy = accuracy_score(y_test, y_test_pred_binary)
        
        self.training_count += 1
        self.Debug(f"Training #{self.training_count}: "
                   f"Train Accuracy: {train_accuracy:.2%}, "
                   f"Test Accuracy: {test_accuracy:.2%}")

    def GetFeatureDataFrame(self):
        # Need 200 bars in rolling window
        if self.data.Count < 200:
            return None

        # Extract the midpoint close from QuoteBars
        close_prices = [qb.Close for qb in self.data]
        df = pd.DataFrame(close_prices, columns=["Close"])

        # --- Feature Engineering ---
        df["SMA_10"] = df["Close"].rolling(window=10).mean()
        df["SMA_50"] = df["Close"].rolling(window=50).mean()
        
        # RSI(14)
        delta = df["Close"].diff()
        gain = (delta.where(delta > 0, 0)).rolling(14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(14).mean()
        rs = gain / loss
        df["RSI"] = 100 - (100 / (1 + rs))
        
        # MACD & Signal
        df["MACD"] = df["Close"].ewm(span=12, adjust=False).mean() - df["Close"].ewm(span=26, adjust=False).mean()
        df["MACD_Signal"] = df["MACD"].ewm(span=9, adjust=False).mean()

        # Historical Volatility (HV_30)
        df["HV_30"] = df["Close"].pct_change().rolling(window=30).std() * np.sqrt(252)

        # Binary target: next day's Close > today's => 1, else 0
        df["Target"] = (df["Close"].shift(-1) > df["Close"]).astype(int)
        
        # Remove NaN rows introduced by rolling calculations
        df.dropna(inplace=True)

        return df