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 GridSearchCV
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.pipeline import Pipeline
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.model = None
        self.scaler = None
        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.model = RandomForestClassifier()
        self.scaler = preprocessing.StandardScaler()
        
        # Hyperparameter tuning using grid search
        parameters = {'n_estimators': [10, 50, 100, 150, 200]}
        clf = GridSearchCV(self.model, parameters)
        
        # Train models and scalers
        features = self.scaler.fit_transform(df[['RSI', 'MACD', 'BB_Upper', 'BB_Middle', 'BB_Lower']])
        clf.fit(features, np.sign(df['Return']))
        
        # Update the model with the best parameters
        self.model = clf.best_estimator_
        
        # Debugging element
        self.Debug("Models trained")
        
    def OnData(self, data):
        self.slice = data
        if not self.model 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 the model to make a prediction
        features_scaled = self.scaler.transform(df)
        prediction = self.model.predict(features_scaled)

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

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