from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from pypfopt.efficient_frontier import EfficientFrontier
from pypfopt import risk_models
from pypfopt import expected_returns

class AlphaFiveUSTreasuries(QCAlgorithm):

    def Initialize(self):

        #1. Required: Five years of backtest history
        self.SetStartDate(2001, 1, 1)
        self.SetEndDate(2003, 1, 1)
    
        #2. Required: Alpha Streams Models:
        self.SetBrokerageModel(BrokerageName.AlphaStreams)
    
        #3. Required: Significant AUM Capacity
        self.SetCash(100000)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseUniverseSelection)
        
        self.AddEquity("SPY",Resolution.Daily)

        #4. Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
    
        self.assets = []

        self.symbols = {}
        
        self.portfolioValue = RollingWindow[Decimal](500)
        
        self.SetWarmup(500)
        
        # initializing security with last known price 
        # self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        
        # fine tune how frequent training is to obtain best results
        # self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("SPY", 30), self.EveryDayAfterMarketOpen)
        # self.Schedule.On(self.DateRules.MonthEnd(), self.TimeRules.AfterMarketOpen("SPY",0), self.EveryMonthAfterMarketClose)
        
    # Portfolio Optimization using Mean-Variance
    # def OptimalPortfolio(self, coarse):
    #     pass 
    

    def CoarseUniverseSelection(self, coarse):
        # if self.lastMonth == algorithm.Time.month:
        #     return Universe.Unchanged
        # self.lastMonth = algorithm.Time.month 
        
        # if self.IsWarmingUp:
        #     return 

        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData], key=lambda x: x.DollarVolume, reverse=True)

        self.Debug(f"test going in sort dollar volume")
 
        for i in range(10):
            if len(self.assets) <10 and sortedByDollarVolume[i] not in self.assets:
                self.assets.append(sortedByDollarVolume[i].Symbol)
    
        self.df = pd.DataFrame()

        # if self.df.iloc[0:,1].count() == 252:
        # self.assets.append()
        
        # for i in range(len(self.assets)):
        # self.symbols[self.assets[i]] = self.AddEquity( self.assets[i] ,Resolution.Hour) 

        for x in range(len(self.assets)):
            self.AddEquity( self.assets[x] ,Resolution.Daily) 

        # for asset in self.assets, count column less than 252 then drop
        self.df = self.History(self.assets, 252)
        
        # check if historical data 252 days available in asset
        self.df = self.df["close"].unstack(level=0)
        
        weights = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
    
        self.returns = self.df.pct_change()

        # show annualized covariance matrix 
        self.cov_matrix_annual = self.returns.cov() * 252 
        
        # self.Debug(self.cov_matrix_annual.iloc[:,1])
        
        # calculate portfolio variance
        self.port_variance = np.dot(weights.T, np.dot(self.cov_matrix_annual, weights))
        
        # calculate portfolio volatility
        self.port_volatility = np.sqrt(self.port_variance)
        
        # self.Debug(f"returns mean")
        # self.Debug(self.returns.mean())
        
        # Calculate annual portfolio return
        self.portfolioSimpleAnnualReturn = np.sum(self.returns.mean() * weights) * 252
        
        # show expected annual return, volatility (risk), variance
        self.percent_var = str(round(self.port_variance, 2) * 100) + '%'
        self.percent_vols = str(round(self.port_volatility, 2) * 100) + '%'
        self.percent_ret = str(round(self.portfolioSimpleAnnualReturn, 2) * 100) + '%'
        
        self.Debug(f'Expected annual return: ')
        self.Debug(self.percent_ret)
        
        self.Debug('Annual volatility / risk: ')
        self.Debug(self.percent_vols)
        
        self.Debug('Annual variance: ')
        self.Debug(self.percent_var)
    
        # Portfolio Optimization
        
        # Calulate the expected returns and annualized sample covariance matrix of asset returns
        self.mu = expected_returns.mean_historical_return(self.df)
        
        self.S = risk_models.sample_cov(self.df)
    
        self.ef = EfficientFrontier(self.mu, self.S)
        
        try:
            self.weights = (self.ef).max_sharpe()
        except:
            # alternatively:
            # return Universe.Unchanged
            return []
        
        self.cleaned_weights = self.ef.clean_weights() 
        
        self.cleaned_weights = [round(float(x), 2) for x in self.cleaned_weights.values()]
        
        # fix cleaned weights
        # self.Debug(f"show updated weights")
        self.Debug(self.cleaned_weights)
        
        # self.Debug(self.ef.portfolio_performance(verbose = True))
    
        return [Symbol.Create(x, SecurityType.Equity, Market.USA) for x in self.symbols.values()]