from io import StringIO
import pandas as pd
import datetime


class PensiveTanAnguilline(QCAlgorithm):

    ##################################################################
    ### Run the algo once to to store the objectStore              ###    
    ### Run the algo a second time to use it as a custom Benchmark ###
    ##################################################################

    ### START TO CUSTOM BENCHMARK ###
    
    portfolio_ObjectStore_Key_test = "test_v1.0.10"
    string_object = ""
    
    ### END - RELATED TO CUSTOM BENCHMARK ###


    START_DATE = datetime.datetime(2020,11,1)
    END_DATE = datetime.datetime(2020,12,15)

    def Initialize(self):
        self.SetStartDate(self.START_DATE.year, self.START_DATE.month, self.START_DATE.day)  
        self.SetEndDate(self.END_DATE.year, self.END_DATE.month, self.END_DATE.day)
        self.SetCash(2000000)
        
        self.symbols = ['A','AAL','AAP','AAPL','ABBV','ABC','ABMD','ABT','ACN','ADBE','ADI','ADM','ADP','ADSK','AEE','AEP','AES','AFL','AIG','AIZ','AJG','AKAM','ALB','ALGN','ALK','ALL','ALLE','ALXN','AMAT','AMCR','AMD','AME','AMGN','AMP','AMT','AMZN','ANET','ANSS','ANTM','AON','AOS','APA','APD','APH','APTV','ARE','ATO','ATVI','AVB','AVGO','AVY','AWK','AXP','AZO','BA','BAC','BAX','BBY','BDX','BEN','BFb','BIIB','BIO','BK','BKNG','BKR','BLK','BLL','BMY','BR','BRKb','BSX','BWA','BXP','C','CAG','CAH','CARR','CAT','CB','CBOE','CBRE','CCI','CCL','CDNS','CDW','CE','CERN','CF','CFG','CHD','CHRW','CHTR','CI','CINF','CL','CLX','CMA','CMCSA','CME','CMG','CMI','CMS','CNC','CNP','COF','COG','COO','COP','COST','CPB','CPRT','CRM','CSCO','CSX','CTAS','CTLT','CTSH','CTVA','CTXS','CVS','CVX','D','DAL','DD','DE','DFS','DG','DGX','DHI','DHR','DIS','DISCA','DISCK','DISH','DLR','DLTR','DOV','DOW','DPZ','DRE','DRI','DTE','DUK','DVA','DVN','DXC','DXCM','EA','EBAY','ECL','ED','EFX','EIX','EL','EMN','EMR','ENPH','EOG','EQIX','EQR','ES','ESS','ETN','ETR','ETSY','EVRG','EW','EXC','EXPD','EXPE','EXR','F','FANG','FAST','FB','FBHS','FCX','FDX','FE','FFIV','FIS','FISV','FITB','FLIR','FLS','FLT','FMC','FOXA','FRC','FRT','FTNT','FTV','GD','GE','GILD','GIS','GL','GLW','GM','GOOG','GOOGL','GPC','GPN','GPS','GRMN','GS','GWW','HAL','HAS','HBAN','HBI','HCA','HD','HES','HFC','HIG','HII','HLT','HOLX','HON','HPE','HPQ','HRL','HSIC','HST','HSY','HUM','HWM','IBM','ICE','IDXX','IEX','IFF','ILMN','INCY','INFO','INTC','INTU','IP','IPG','IPGP','IQV','IR','IRM','ISRG','IT','ITW','IVZ','JBHT','JCI','JEC','JKHY','JNJ','JNPR','JPM','K','KEY','KEYS','KHC','KIM','KLAC','KMB','KMI','KMX','KO','KR','KSU','L','LB','LDOS','LEG','LEN','LH','LHX','LIN','LKQ','LLY','LMT','LNC','LNT','LOW','LRCX','LUMN','LUV','LVS','LW','LYB','LYV','MA','MAA','MAR','MAS','MCD','MCHP','MCK','MCO','MDLZ','MDT','MET','MGM','MHK','MKC','MKTX','MLM','MMC','MMM','MNST','MO','MOS','MPC','MPWR','MRK','MRO','MS','MSCI','MSFT','MSI','MTB','MTD','MU','MXIM','NCLH','NDAQ','NEE','NEM','NFLX','NI','NKE','NLOK','NLSN','NOC','NOV','NOW','NRG','NSC','NTAP','NTRS','NUE','NVDA','NVR','NWL','NWSA','O','ODFL','OKE','OMC','ORCL','ORLY','OTIS','OXY','PAYC','PAYX','PBCT','PCAR','PEAK','PEG','PEP','PFE','PFG','PG','PGR','PH','PHM','PKG','PKI','PLD','PM','PNC','PNR','PNW','POOL','PPG','PPL','PRGO','PRU','PSA','PSX','PVH','PWR','PXD','PYPL','QCOM','QRVO','RCL','RE','REG','REGN','RF','RHI','RJF','RL','RMD','ROK','ROL','ROP','ROST','RSG','RTX','SBAC','SBUX','SCHW','SEE','SHW','SIVB','SJM','SLB','SLG','SNA','SNPS','SO','SPG','SPGI','SRE','STE','STT','STX','STZ','SWK','SWKS','SYF','SYK','SYY','T','TAP','TDG','TDY','TEL','TER','TFC','TFX','TGT','TJX','TMO','TMUS','TPR','TRMB','TROW','TRV','TSCO','TSLA','TSN','TT','TTWO','TWTR','TXN','TXT','TYL','UA','UAA','UAL','UDR','UHS','ULTA','UNH','UNM','UNP','UPS','URI','USB','V','VAR','VFC','VIAC','VLO','VMC','VNO','VNT','VRSK','VRSN','VRTX','VTR','VTRS','VZ','WAB','WAT','WBA','WDC','WEC','WELL','WFC','WHR','WLTW','WM','WMB','WMT','WRB','WRK','WST','WU','WY','WYNN','XEL','XLNX','XOM','XRAY','XRX','XYL','YUM','ZBH','ZBRA','ZION','ZTS']
        self.allSecurities = []
        for symbol in self.symbols:
            security = self.AddEquity(symbol, Resolution.Daily).Symbol 
            self.allSecurities.append(security)
        
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol 
        self.tradableSecurities = []
        
        self.SetWarmUp(timedelta(days = 30))
        
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday),
                self.TimeRules.BeforeMarketClose("SPY", 0),
                self.rebalance )
                
        
        self.investedCount = 0

        ### START - RELATED TO CUSTOM BENCHMARK ###

        # ObjectStore for custom Benchmark
        if self.ObjectStore.ContainsKey(self.portfolio_ObjectStore_Key_test):
            values = self.ObjectStore.Read(self.portfolio_ObjectStore_Key_test)
            self.Debug(f'{self.portfolio_ObjectStore_Key_test} key exists in object store.')
            self.Debug("str(values)" + "\n")
            self.Debug(str(values))
            history = pd.read_csv(StringIO(values), header=None, index_col=0, squeeze=True)
            history.index = pd.to_datetime(history.index)

            """
                Here comes the problems
                I tried two ways:
                    1) benchmark = [history.index.values,history.values]
                       self.SetBenchmark(benchmark)
                    
                    2) self.SetBenchmark(self.custom_function) # function definition below
                
                Both doesn't work
                    
            """
            
            # Doesn't work : 
            #benchmark = [history.index.values,history.values]
            #self.SetBenchmark(benchmark)

            # Doesn't work neither:
            #self.SetBenchmark(self.custom_function)
            

        ### END - RELATED TO CUSTOM BENCHMARK ###

    def OnData(self, data):
        
        ### START - RELATED TO CUSTOM BENCHMARK ###

        #self.Debug("self.Time : " + str(self.Time)  + " - FIN" + "\n" )
        # output : self.Time : 2021-02-01 00:00:00 - FIN

        #self.Debug("self.Portfolio.TotalPortfolioValue : " + str(self.Portfolio.TotalPortfolioValue) + " - FIN" + "\n")
        # output : self.Portfolio.TotalPortfolioValue : 2000000.0 - FIN
        
        ### Add only portfolio value since the start_date
        ### We don't store the portfolio value during the warmup period
        if(self.Time >= self.START_DATE):
            self.string_object += ('\n' + str(self.Time) + "," + str(self.Portfolio.TotalPortfolioValue) )
        

        #Save the object at the last day of the algorithm
        if( (self.END_DATE.year == self.Time.year)  and (self.END_DATE.month == self.Time.month)  and (self.END_DATE.day == self.Time.day) ):
            self.Debug(str(self.Time))
            self.ObjectStore.Save(self.portfolio_ObjectStore_Key_test, self.string_object)

        # self.Debug("string_object : " + self.string_object)
        

        ### END - RELATED TO CUSTOM BENCHMARK ###
        
        self.tradableSecurities = []
        for security in self.allSecurities:
            if data.ContainsKey(security.Value):
                self.tradableSecurities.append(security)
            else:
                continue

    def rebalance(self):
        
        '''Include sell statement for removed securities'''
        # Faire if len < 100; return
        if len(self.tradableSecurities) == 0:
            return

        if (len(self.tradableSecurities) < 100):
            self.Debug("Number of tradableSecurities : " + str(len(self.tradableSecurities)))
            self.Debug("We don''t rebalance")        
            return        
                
        for security in self.tradableSecurities:
            if not self.Portfolio[security].Invested:
                self.investedCount += 1
        
        INVESTED_FRACTION = 0.995 # Remove 0.5% for the fees ?
        weight_temp = INVESTED_FRACTION/len(self.tradableSecurities)
        WEIGHT_LIMIT = 0.03

        weight = self.weightLimitation(weight_temp, WEIGHT_LIMIT)
        
        for security in self.tradableSecurities:
            if self.Portfolio[security].Invested:
                self.SetHoldings([PortfolioTarget(security.Value, weight)])
                
        for security in self.tradableSecurities:
            if not self.Portfolio[security].Invested:
                self.SetHoldings([PortfolioTarget(security.Value, weight)])
                
    ### START - RELATED TO CUSTOM BENCHMARK ###
                
    def custom_function(self):
        
        time = self.Time
        date = datetime.datetime(time.year, time.month, time.day)
            
        t = pd.Timestamp(date)
        value = history.loc[t]
            
        return t, value

    ### END - RELATED TO CUSTOM BENCHMARK ###


    def weightLimitation(self, weight_temp,WEIGHT_LIMIT):
        if(weight_temp > WEIGHT_LIMIT):
            return WEIGHT_LIMIT
        else:
            return weight_temp