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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data.Market import TradeBar
from sklearn import preprocessing
from sklearn.model_selection import train_test_split 
from sklearn.neighbors import KNeighborsClassifier
from sklearn import svm
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn import metrics
import pandas as pd
import numpy as np
from datetime import datetime
from QuantConnect.Indicators import RelativeStrengthIndex, MovingAverageConvergenceDivergence, BollingerBands

class CryptoPredictionAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2016, 1, 1)
        self.SetCash(100000)
        self.crypto = self.AddCrypto("BTCUSD", Resolution.Daily).Symbol
        self.lookback = 30

        self.Schedule.On(self.DateRules.EveryDay(self.crypto), self.TimeRules.At(0, 0), self.TrainModels)
        self.models = []
        self.scalers = []
        self.slice = None
        
        # Debugging elements
        self.Debug("Algorithm initialized")
        self.Debug("Symbol: " + str(self.crypto))
        self.Debug("Lookback period: " + str(self.lookback) + " days")

        self.rsi = RelativeStrengthIndex(14)
        self.macd = MovingAverageConvergenceDivergence(12, 26, 9)
        self.bb = BollingerBands(20, 2)

    def TrainModels(self):
        # Get historical data
        history = self.History(self.crypto, self.lookback, Resolution.Daily)
        
        # Ensure we have enough data
        if len(history) < self.lookback:
            self.Debug("Not enough data to train models")
            return
        
        # Debugging element
        self.Debug("Training models")

        # Create dataframe and calculate features
        df = history.loc[self.crypto].copy()
        df['Return'] = df['close'].pct_change()
        df.dropna(inplace=True)
        
        # Insert your feature calculation here
        for i in range(len(df)):
            end_time = df.index[i]
            price = df.loc[end_time, "close"]
            self.rsi.Update(end_time, price)
            self.macd.Update(end_time, price)
            self.bb.Update(end_time, price)
            df.loc[end_time, 'RSI'] = self.rsi.Current.Value
            df.loc[end_time, 'MACD'] = self.macd.Current.Value
            df.loc[end_time, 'BB_Upper'] = self.bb.UpperBand.Current.Value
            df.loc[end_time, 'BB_Middle'] = self.bb.MiddleBand.Current.Value
            df.loc[end_time, 'BB_Lower'] = self.bb.LowerBand.Current.Value
        
        df.dropna(inplace=True)

        # Create models and scalers
        self.models = []
        self.scalers = []
        
        # Train models and scalers
        for model in [KNeighborsClassifier(n_neighbors=3), svm.SVC(kernel='poly'), RandomForestClassifier(n_estimators=150), GradientBoostingClassifier(n_estimators=150)]:
            scaler = preprocessing.StandardScaler()
            features = scaler.fit_transform(df[['RSI', 'MACD', 'BB_Upper', 'BB_Middle', 'BB_Lower']])
            model.fit(features, np.sign(df['Return']))
            self.models.append(model)
            self.scalers.append(scaler)
        
        # Debugging element
        self.Debug("Models trained")
        
    def OnData(self, data):
        self.slice = data
        if not self.models or not self.slice.ContainsKey(self.crypto):
            return

        # Update the indicators with the latest price
        price = self.slice[self.crypto].Close
        self.rsi.Update(self.Time, price)
        self.macd.Update(self.Time, price)
        self.bb.Update(self.Time, price)

        # Create a DataFrame for the latest data
        df = pd.DataFrame({
            'RSI': [self.rsi.Current.Value],
            'MACD': [self.macd.Current.Value],
            'BB_Upper': [self.bb.UpperBand.Current.Value],
            'BB_Middle': [self.bb.MiddleBand.Current.Value],
            'BB_Lower': [self.bb.LowerBand.Current.Value]
        })

        # Use each model to make a prediction
        predictions = []
        for model, scaler in zip(self.models, self.scalers):
            features_scaled = scaler.transform(df)
            prediction = model.predict(features_scaled)
            predictions.append(prediction)

        # Decide whether to buy or sell based on the predictions
        if all(prediction == 1 for prediction in predictions):
            self.SetHoldings(self.crypto, 1.0)
        elif all(prediction == -1 for prediction in predictions):
            self.Liquidate(self.crypto)

        # Debugging element
        self.Debug("Prediction: " + str(predictions))