from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
import numpy as np
class AlphaFiveUSTreasuries(QCAlgorithm):

    def Initialize(self):

        #1. Required: Five years of backtest history
        self.SetStartDate(2014, 1, 1)
        self.SetEndDate(2014, 1, 10)
    
        #2. Required: Alpha Streams Models:
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
    
        #3. Required: Significant AUM Capacity
        self.SetCash(100000)

        self.AddUniverse(self.CoarseSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Hour
    
        self.spy = self.AddEquity("SPY", Resolution.Hour).Symbol
        
        #4. Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())
    
        # self.assets = ["IEF", "SHY", "TLT", "IEI", "SHV", "TLH", "EDV", "BIL",
        #               "SPTL", "TBT", "TMF", "TMV", "TBF", "VGSH", "VGIT",
        #               "VGLT", "SCHO", "SCHR", "SPTS", "GOVT"]

        self.assets = []

        self.symbols = {}
        
        self.portfolioValue = RollingWindow[Decimal](500)
        
        self.SetWarmup(500)
        
        # Add Equity ------------------------------------------------ 
        # add equities from coarse universe to self.symbols
        # for i in range(len(self.assets)):
        #     self.symbols[self.assets[i]] = self.AddEquity(self.assets[i],Resolution.Hour).Symbol 
        
        # test different timings after market opens affects the returns
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("SPY", 0), self.EveryDayAfterMarketOpen)
    
    def CoarseSelectionFunction(self, coarse):
        # add the top 20 assets
        # Sort symbols by Dollar Volume and if fundamental data exists, descending
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 10],
                key=lambda x: x.DollarVolume, reverse=True)

        # for y in range(20):
        #     self.Debug(sortedByDollarVolume[y])
        #     self.Debug(f"Test going into range")

        # if self.lastMonth == algorithm.Time.month:
        #     return Universe.Unchanged
        # self.lastMonth = algorithm.Time.month
        self.Debug(f"before dictionary completion")
        
        self.assets = sortedByDollarVolume[:20]
        
        for x in range(len(self.assets)):
            if self.assets[x].Value is None:
                self.assets.RemoveSecurity(self.assets[x])
        
        self.Debug(f"test went through assets list")
        
        # for asset in self.assets:
        for x in self.assets:
            self.Debug(x)
            # if not self.Securities.ContainsKey(asset.Symbol):
            # self.symbols[self.assets[x]] = self.AddEquity(self.assets[x],Resolution.Hour).Symbol 
                

        # add assets list 
        # for asset in self.assets:
        #     self.Debug(f"test going in")
        #     self.Debug(asset)
            # self.Debug("{}".format(type(self.assets[x])))
            # create symbols and tickers dictionary
            #self.symbols[self.assets[x]] = self.AddEquity(self.assets[x],Resolution.Hour).Symbol 

        
        # add tickers and symbols to dictionary
        # for i in range(len(self.assets)):
        #     self.symbols[self.assets[x]] = self.AddEquity(self.assets[x],Resolution.Hour).Symbol 
        
        self.Debug(f"test dictionary completion")
        # for ticker, symbol in self.symbols:
        #     self.Debug(symbol)
        #     self.Debug(ticker)
        
        # return self.symbols
        return [x.Symbol for x in sortedByDollarVolume[:10]]
        
    def EveryDayAfterMarketOpen(self):
        for symbol, ticker in self.symbols.items():
            self.Debug(symbol, ticker)
            
        # for ticker, symbol in self.symbols:
        #     self.Debug(f"test going in")
        #     self.Debug(ticker, symbol)
        
        # training for model 
        # if not self.Portfolio.Invested:
        #     insights = []
        #     for ticker, symbol in self.symbols.items():
        #         insights.append( Insight.Price(symbol, timedelta(days=5), InsightDirection.Up, 0.01, None, None, 1/len(self.symbols)) )
        #     self.EmitInsights(insights)
        # else:
        #     qb = self 
        #     #==============================
        #     # Initialize instance of Random Forest Regressor
        #     regressor = RandomForestRegressor(n_estimators=100, min_samples_split=5, random_state = 1990)
    
        #     # Fetch history on our universe
        #     df = qb.History(qb.Securities.Keys, 500, Resolution.Hour)
            
        #     # Get train/test data
        #     returns = df.unstack(level=1).close.transpose().pct_change().dropna()
        #     X = returns
        #     y = [x for x in qb.portfolioValue][-X.shape[0]:]
        #     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 1990)
            
        #     # Fit regressor
        #     regressor.fit(X_train, y_train)
            
        #     # Get long-only predictions
        #     weights = regressor.feature_importances_
        #     symbols = returns.columns[np.where(weights)]
        #     selected = zip(symbols, weights)
        #     # ==============================
            
        #     insights = []
        #     for symbol, weight in selected:
        #         insights.append( Insight.Price(symbol, timedelta(days=5), InsightDirection.Up, 0.01, None, None, weight) )
        #     self.EmitInsights(insights)
        
    def OnData(self, data):
        self.portfolioValue.Add(self.Portfolio.TotalPortfolioValue)