from QuantConnect.Data.UniverseSelection import *
import numpy as np
import pandas as pd

### <summary>

### </summary>

class FiftyTwoWeeksHighEffectInStocks(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015,10,15)  #Set Start Date
        self.SetEndDate(2018,10,31)
        self.SetCash(100000)            #Set Strategy Cash
    
        self.symbols = None
        self.rebalance = 1
        self.windows = {}
        self.ratios = {}
        self.mktCap = {}
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
    
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol 
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.At(9,31), self.month_rebalance)
        
        
        #We define self.symbols in FineSelectionFunction, if we use warmup then Ondata and OnSecuritiesChanged will have no idea what self.symbols is
        # self.SetWarmup(252)  


        
    def CoarseSelectionFunction(self, coarse):
        if self.rebalance or self.symbols is None:
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 3)]
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
            return [x.Symbol for x in filtered[:500]]    #Pick 1000 most liquid stocks
        else:
            return self.symbols
    
    def FineSelectionFunction(self, fine):
        if self.rebalance or self.symbols is None:
            #Group the stocks by industry sector
            self.BasicMaterialsIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.BasicMaterials]
            self.ConsumerCyclicalIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.ConsumerCyclical]
            self.FinancialServicesIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.FinancialServices]
            self.RealEstateIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.RealEstate]
            self.ConsumerDefensiveIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.ConsumerDefensive]
            self.HealthcareIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Healthcare]
            self.UtilitiesIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Utilities]
            self.CommunicationServicesIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.CommunicationServices]
            self.EnergyIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Energy]
            self.IndustrialsIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Industrials]
            self.TechnologyIndustry = [x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology]
     
            self.pools = self.BasicMaterialsIndustry + self.ConsumerCyclicalIndustry + self.FinancialServicesIndustry + self.RealEstateIndustry + self.ConsumerDefensiveIndustry + self.HealthcareIndustry + self.UtilitiesIndustry + self.CommunicationServicesIndustry + self.EnergyIndustry + self.IndustrialsIndustry + self.TechnologyIndustry
            
            hist = self.History([i.Symbol for i in self.pools], 252, Resolution.Daily)   #The hist.index is in this format: "X R735QTJ8XC9
            for x in self.pools:
                self.mktCap[x.Symbol] = float(x.EarningReports.BasicAverageShares.ThreeMonths) * hist.loc[str(x.Symbol)]['close'][-1]  #str(x.Symbol) is "X R735QTJ8XC9X" ; #x.Symbol.Value is "X" so we should use str(x.Symbol) here

            
            self.symbols = [x.Symbol for x in self.pools]
            for symbol in self.symbols:
                self.windows[symbol] = RollingWindow[Decimal](252)
                for historical_close in hist.loc[str(symbol)]['close']:
                    self.windows[symbol].Add(historical_close)
                    
            self.rebalance = 0
            self.Log(f'{self.Time}:FineSelection is running! The size of self.symbols is {len(self.symbols)}')
            return self.symbols
        else:
            return self.symbols
            
    def OnData(self, data):  #Update the rolling windows
        if self.symbols is not None:
            for symbol in self.symbols:
                self.Log(symbol)

                if data.ContainsKey(symbol) and data[symbol] is not None :
                    self.Log(data[symbol].Close)
                    self.windows[symbol].Add(data[symbol].Close)
                    
                    self.Log("02")
                else:
                    self.windows[symbol].Add(0)
                    self.Log("03")

      

    def month_rebalance(self):
        
        self.Log(f'{self.Time} month_rebalance is running')
        self.rebalance = 1
        
        #Calculate the weighted close/52-week-high ratios for each sector
        
        industry_list = [self.BasicMaterialsIndustry,
                        self.ConsumerCyclicalIndustry, 
                        self.FinancialServicesIndustry,
                        self.RealEstateIndustry,
                        self.ConsumerDefensiveIndustry,
                        self.HealthcareIndustry,
                        self.UtilitiesIndustry,
                        self.CommunicationServicesIndustry,
                        self.EnergyIndustry,
                        self.IndustrialsIndustry,
                        self.TechnologyIndustry]

        weighted_scores = []
        for i in range(0,len(industry_list)):
            mktCap_sum = sum([self.mktCap[x.Symbol] for x in industry_list[i]])     #The sum of mkt capitalization of i-th sector
            self.Log(f"{self.Time} {i}-th maktCap is good")
            
            curPrice = np.array([self.windows[x.Symbol][0] for x in (industry_list[i])])   #The current prices of each stock in i-th sector
            
            # self.Log(curPrice)
            # self.Log(f'curPrice of {i}-th industry:{curPrice}')

            self.Log(f"{self.Time} {i}-th curPrice is good")
            
            
            fiftyTwoHigh = np.array( [ max(self.windows[x.Symbol]) for x in industry_list[i] ] ) #The 52-week high of each stock in i-th sector
            x = industry_list[i][0]


            # self.Log(f'fiftyTwoHigh of {i}-th industry:{fiftyTwoHigh}')
            self.Log(f'{self.Time} {i}-th fiftyTwoHigh is good')
            
           
            
            ratios = curPrice/fiftyTwoHigh                                       # Ratio (current prices/52-week high) of each stock in i-th sector
            # self.Log(f'ratios of {i}-th industry:{ratios}}')
            self.Log(f'{self.Time} {i}-th ratios is good')
            
            weights = np.array([self.mktCap[x.Symbol]/mktCap_sum for x in industry_list[i]]) #Weighted ratio for i-th sector
            self.Log(f'{self.Time} {i}-th weights is good')
            
            weighted_scores.append(sum(ratios*weights))   #Add weighted ratio to weighted_scores -> list
        
        self.Log("OK1")
        # self.Log(f'weighted_scores: {weighted_scores}')
      
        
        

        long_industry = industry_list[np.argmax(weighted_scores)]
        short_industry = industry_list[np.argmin(weighted_scores)]

        self.Log('long_industry and short_industry is good')
        for x in long_industry:
            self.SetHoldings(x.Symbol, 0.5/len(long_industry))
        for x in short_industry:
            self.SetHoldings(x.Symbol, -0.5/len(short_industry))