from AlgorithmImports import *
from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import pandas as pd

# ------------------------------
# Custom Models: 0.1% Fee & 0 Slippage
# ------------------------------
class CustomFeeModel:
    """
    Applies a 0.1% transaction fee on each trade (open/close).
    """
    def GetOrderFee(self, parameters):
        orderValue = parameters.Security.Price * abs(parameters.Order.Quantity)
        fee = 0.001 * orderValue  # 0.1% of trade notional
        return OrderFee(CashAmount(fee, "USD"))

class CustomSlippageModel:
    """
    Sets slippage to 0.
    """
    def GetSlippageApproximation(self, asset, order):
        return 0

class MLTradingAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(10000)
        self.symbol = self.AddEquity("MARA", Resolution.Daily).Symbol
        self.Securities[self.symbol].SetFeeModel(CustomFeeModel())
        self.Securities[self.symbol].SetSlippageModel(CustomSlippageModel())
        self.data = RollingWindow[TradeBar](200)
        self.SetWarmUp(200)
        self.model = SVC(probability=True, random_state=42)
        self.training_count = 0
        self.is_model_trained = False
        self.allocation_fraction = 0.2
        self.spySymbol = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Securities[self.spySymbol].SetFeeModel(CustomFeeModel())
        self.Securities[self.spySymbol].SetSlippageModel(CustomSlippageModel())
        self.SetBenchmark(self.spySymbol)
        self.spyPriceStart = None
        self.initialCapital = None
        # Initialize counters for tracking performance against SPY
        self.beat_spy_count = 0
        self.total_comparisons = 0
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), 
                         self.TimeRules.At(10, 0), 
                         self.TrainModel)

    def OnData(self, data):
        if self.initialCapital is None:
            self.initialCapital = self.Portfolio.TotalPortfolioValue
        
        if self.spyPriceStart is None and data.ContainsKey(self.spySymbol):
            bar = data[self.spySymbol]
            if bar and bar.Close > 0:
                self.spyPriceStart = bar.Close

        if not data.ContainsKey(self.symbol):
            return
        
        trade_bar = data[self.symbol]
        if trade_bar is None:
            return
        
        self.data.Add(trade_bar)

        if not self.data.IsReady or self.data.Count < 200:
            return
        
        if not self.is_model_trained:
            self.Debug("Model is not trained yet. Skipping prediction.")
            return

        df = self.GetFeatureDataFrame()
        if df is None or len(df) < 1:
            return
        
        latest_features = df.iloc[-1, :-1].values.reshape(1, -1)
        
        try:
            prob_class = self.model.predict_proba(latest_features)[0][1]
            prediction = 1 if prob_class > 0.5 else 0
        except Exception as e:
            self.Debug(f"Error: Model prediction failed. {e}")
            return
        
        holdings = self.Portfolio[self.symbol].Quantity
        
        if prediction == 1 and holdings <= 0:
            self.SetHoldings(self.symbol, self.allocation_fraction) 
        elif prediction == 0 and holdings > 0:
            self.Liquidate(self.symbol)

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

        X = df.iloc[:, :-1]
        y = df.iloc[:, -1]
        
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, shuffle=False, random_state=42
        )

        self.model.fit(X_train, y_train)
        self.is_model_trained = True

        y_train_prob = self.model.predict_proba(X_train)[:, 1]
        y_train_pred_binary = [1 if val > 0.5 else 0 for val in y_train_prob]
        train_accuracy = accuracy_score(y_train, y_train_pred_binary)

        y_test_prob = self.model.predict_proba(X_test)[:, 1]
        y_test_pred_binary = [1 if val > 0.5 else 0 for val in y_test_prob]
        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):
        if self.data.Count < 200:
            return None
        
        close_prices = [bar.Close for bar in self.data]
        df = pd.DataFrame(close_prices, columns=["Close"])
        
        df["SMA_10"] = df["Close"].rolling(window=10).mean()
        df["SMA_50"] = df["Close"].rolling(window=50).mean()
        
        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))
        
        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()

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

        df["Target"] = (df["Close"].shift(-1) > df["Close"]).astype(int)
        
        df.dropna(inplace=True)

        return df

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            if self.spyPriceStart is None or self.spyPriceStart == 0:
                return

            strategyReturn = (self.Portfolio.TotalPortfolioValue / self.initialCapital - 1) * 100.0
            spyPriceNow = self.Securities[self.spySymbol].Price
            spyReturn = (spyPriceNow / self.spyPriceStart - 1) * 100.0

            # Increment total comparisons and beat count if strategy outperforms SPY
            self.total_comparisons += 1
            if strategyReturn > spyReturn:
                self.beat_spy_count += 1

            if strategyReturn > spyReturn:
                conclusion = "Strategy is beating SPY"
            elif strategyReturn < spyReturn:
                conclusion = "SPY is beating the Strategy"
            else:
                conclusion = "Strategy and SPY are at the same return"

            self.Debug(f"[Order Filled] Strategy Return: {strategyReturn:.2f}%, "
                       f"SPY B/H Return: {spyReturn:.2f}%. {conclusion}")

    def OnEndOfAlgorithm(self):
        # Print the final count of times the strategy beat SPY
        self.Debug(f"Number of times strategy beat SPY: {self.beat_spy_count} out of {self.total_comparisons}")