from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import StandardScaler
import numpy as np
import pandas as pd
import talib as ta

class MyAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # Find more symbols here: http://quantconnect.com/data
        self.AddEquity("BTC-USD", Resolution.Daily)
        self.symbols = [self.Securities["BTC-USD"].Symbol]
        
        # Initialize models and scalers
        self.models = [GradientBoostingClassifier(), RandomForestClassifier(), SVC(), KNeighborsClassifier()]
        self.scalers = [StandardScaler() for _ in range(4)]
        
        # Initialize other variables
        self.lookback = 14
        self.slice = None
        self.is_trained = False

    def OnData(self, data):
        if not self.is_trained:
            self.TrainModel(data)
            self.is_trained = True
        else:
            self.Predict(data)

    def TrainModel(self, data):
        df = pd.DataFrame()
        for symbol in self.symbols:
            history = self.History(symbol, self.lookback, Resolution.Daily)
            if 'close' not in history.columns:
                return
            history = history[['close']]
            history['pct_change'] = history['close'].pct_change()
            history['RSI'] = ta.RSI(history['close'].values, timeperiod=14)
            history['MACD'], _, _ = ta.MACD(history['close'].values, fastperiod=12, slowperiod=26, signalperiod=9)
            history['upper'], history['middle'], history['lower'] = ta.BBANDS(history['close'].values, timeperiod=5)
            history.dropna(inplace=True)
            df = pd.concat([df, history])
        df.dropna(inplace=True)

        # Define features and target
        features = df[['RSI', 'MACD', 'upper', 'middle', 'lower']]
        target = np.where(df['pct_change'] > 0, 1, 0)

        # Train models and scalers
        for i in range(4):
            self.scalers[i].fit(features)
            features_scaled = self.scalers[i].transform(features)
            self.models[i].fit(features_scaled, target)

    def Predict(self, data):
        df = pd.DataFrame()
        for symbol in self.symbols:
            history = self.History(symbol, self.lookback, Resolution.Daily)
            if 'close' not in history.columns:
                return
            history = history[['close']]
            history['pct_change'] = history['close'].pct_change()
            history['RSI'] = ta.RSI(history['close'].values, timeperiod=14)
            history['MACD'], _, _ = ta.MACD(history['close'].values, fastperiod=12, slowperiod=26, signalperiod=9)
            history['upper'], history['middle'], history['lower'] = ta.BBANDS(history['close'].values, timeperiod=5)
            history.dropna(inplace=True)
            df = pd.concat([df, history])
        df.dropna(inplace=True)

        # Define features
        features = df[['RSI', 'MACD', 'upper', 'middle', 'lower']]

        # Make predictions
        predictions = []
        for i in range(4):
            features_scaled = self.scalers[i].transform(features)
            predictions.append(self.models[i].predict(features_scaled))

        # Ensemble prediction
        prediction = 1 if sum(predictions) > len(predictions) / 2 else 0

        # Place orders based on prediction
        if prediction == 1:
            self.SetHoldings("BTC-USD", 1)
        else:
            self.SetHoldings("BTC-USD", -1)