from AlgorithmImports import *
import numpy as np
from sklearn.linear_model import LogisticRegression
from datetime import date, datetime, timedelta
from sklearn.model_selection import train_test_split

class SleepyYellowGreenWhale(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2023, 2, 25)
        #2. Required: Brokerage Models:
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

    
        #3. Required: Significant AUM Capacity
        self.SetCash(10000)

        #4. Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        #spx = qb.AddCfd("SPX500USD").Symbol
        self.assets = ["SPY", "QQQ", "GLD"]


        self.symbols = {}
        
        self.portfolioValue = RollingWindow[Decimal](500)
        
        self.SetWarmup(500)
        self.take_profit = 0.70
        self.stop_loss = 0.17
        self.min_data = 100
        self.lookback = 100
        #self.hyperparameter_update_interval = 14
        #self.last_training = datetime.min
        self.units = 10
        #self.signal=0

    def EveryDayAfterMarketOpen(self):
        holdings = {}
        for symbol in self.assets:
            holdings[symbol] = self.Portfolio[symbol].Quantity

        for symbol in self.assets:
            holdings[symbol] = self.Portfolio[symbol].Quantity
            signal = self.GetSignal(symbol)
            if holdings[symbol] == 0:
                if signal == 1:
                    order = self.MarketOrder(symbol, self.units)
                    self.CreateStopLossAndTakeProfitOrders(order)
            elif holdings[symbol] > 0:
                if signal == -1:
                    self.Liquidate(symbol)
                elif signal == 1:
                    self.SetHoldings(symbol, 1/len(self.assets))
                else:
                    if signal == 1:
                        order = self.MarketOrder(symbol, self.units)
                        self.CreateStopLossAndTakeProfitOrders(order)


    def CreateStopLossAndTakeProfitOrders(self, order):
        self.stopMarketTicket = self.StopMarketOrder(order.Symbol, -self.units, self.stop_loss * order.AverageFillPrice, "Stop Loss")
        self.profitTargetTicket = self.LimitOrder(order.Symbol, self.units, self.take_profit * order.AverageFillPrice, "Take Profit")

    
    def GetSignal(self, symbol):
        df = self.History([symbol], self.lookback, Resolution.Daily)
        if df is None:
            return 0
    
        df['ma60'] = self.SMA(symbol, 60)
        df['ma150'] = self.SMA(symbol, 150)
        df['rsi30'] = self.RSI(symbol, 30)

        if len(df) >= self.min_data:
            X = df['open', 'high', 'low', 'close', 'volume', 'ma60', 'ma150', 'rsi30']
            X = (X - X.mean()) / X.std()
            X = X.values
            regressor = LogisticRegression(random_state=0, solver='lbfgs', max_iter=100)
            signal = np.where(df['close'].shift(-1) > df['close'], 1, -1)
            X_train, X_test, y_train, y_test = train_test_split(X, signal, test_size=0.2, random_state=1990)
            regressor.fit(X_train, y_train)
            return np.sign(regressor.predict(X)[-1])
        else:
            return 0
        
    def OnData(self, data):
        self.portfolioValue.Add(self.Portfolio.TotalPortfolioValue)