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 QuantConnect.Securities import Crypto
from QuantConnect.Brokerages import BrokerageName
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(137) # initial investment
        self.crypto = self.AddCrypto("BTCUSD", Resolution.Daily).Symbol
        self.lookback = 30
        self.minimumOrderSize = 0.000016

        self.Schedule.On(self.DateRules.EveryDay(self.crypto), self.TimeRules.At(0, 0), self.TrainModels)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.At(0, 0), self.AddWeeklyInvestment) # add weekly investment
        self.models = []
        self.scalers = []
        self.slice = None
        
        self.rsi = RelativeStrengthIndex(14)
        self.macd = MovingAverageConvergenceDivergence(12, 26, 9)
        self.bb = BollingerBands(20, 2)
        
        self.SetBrokerageModel(BrokerageName.GDAX, AccountType.Cash) # Set brokerage model
        
    def AddWeeklyInvestment(self):
        self.Portfolio.CashBook["USD"].AddAmount(68.5) # add $68.5 every week

    def TrainModels(self):
        history = self.History(self.crypto, self.lookback, Resolution.Daily)
        
        if len(history) < self.lookback:
            return
        
        df = history.loc[self.crypto].copy()
        df['Return'] = df['close'].pct_change()
        df.dropna(inplace=True)
        
        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)
        
        self.models = []
        self.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)
        
    def OnData(self, data):
        self.slice = data
        if not self.models or not self.slice.ContainsKey(self.crypto):
            return

        price = self.slice[self.crypto].Close
        self.rsi.Update(self.Time, price)
        self.macd.Update(self.Time, price)
        self.bb.Update(self.Time, price)

        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]
        })

        predictions = []
        for model, scaler in zip(self.models, self.scalers):
            features_scaled = scaler.transform(df)
            prediction = model.predict(features_scaled)
            predictions.append(prediction)

        if all(prediction == 1 for prediction in predictions):
            orderSize = self.CalculateOrderQuantity(self.crypto, 1.0)
            if abs(orderSize) >= self.minimumOrderSize:
                self.SetHoldings(self.crypto, 1.0)

        elif all(prediction == -1 for prediction in predictions):
            holdings = self.Portfolio[self.crypto].Quantity
            if abs(holdings) >= self.minimumOrderSize:
                self.Liquidate(self.crypto)
        
        holdings = self.Portfolio[self.crypto].Quantity
        if holdings > 0 and self.Portfolio[self.crypto].UnrealizedProfitPercent < -0.10:  # 10% stop-loss
            self.Liquidate(self.crypto)
        elif holdings < 0 and self.Portfolio[self.crypto].UnrealizedProfitPercent < -0.10:  # 10% stop-loss for short position
            self.Liquidate(self.crypto)