from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
import scipy as sp
from collections import deque

class MomentumReversalCombinedWithVolatility(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2020, 1, 1)  # Set Start Date
        #self.SetEndDate(2018, 8, 1)    # Set Start Date       
        self.SetCash(100000)           # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.dataDict = {}
        # 1/6 of the portfolio is rebalanced every month
        self.portfolios = deque(maxlen=6)
        self.AddEquity("SPY", Resolution.Daily)
        self.Schedule.On(self.DateRules.MonthStart("SPY"),self.TimeRules.AfterMarketOpen("SPY"), self.Rebalance)
        # the lookback period for volatility and return is six months
        self.lookback = 20*6
        self.filteredFine = None
        self.monthly_rebalance = False

    def CoarseSelectionFunction(self, coarse):
        # update the price of stocks in universe everyday
        for i in coarse:
            if i.Symbol not in self.dataDict:
                self.dataDict[i.Symbol] = SymbolData(i.Symbol, self.lookback)
            self.dataDict[i.Symbol].Update(self.Time, i.AdjustedPrice)
        
        if self.monthly_rebalance:
            # drop stocks which have no fundamental data or have too low prices
            filteredCoarse = [x.Symbol for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            return filteredCoarse
        else:
            return Universe.Unchanged #[]

    def FineSelectionFunction(self, fine):
  
        if self.monthly_rebalance:  
            sortedFine = sorted(fine, key = lambda x: x.EarningReports.BasicAverageShares.Value * self.dataDict[x.Symbol].Price, reverse=True)
            # select stocks with large size 
            topFine = sortedFine[:int(0.5*len(sortedFine))]
            self.filteredFine = [x.Symbol for x in topFine]
            return self.filteredFine
        else:
            return Universe.Unchanged #[]
            
        
    def Rebalance(self):
        self.monthly_rebalance = True

           
    def OnData(self, data):
        
        for symbol, symbolData in self.dataDict.items():
            if symbolData.IsReady():
                self.Quit(f"Custom Sharpe metric for {symbol}: {symbolData.Sharpe()}")
        
        
        if self.monthly_rebalance and self.filteredFine:

            filtered_data = {symbol: symbolData for (symbol, symbolData) in self.dataDict.items() if symbol in self.filteredFine and symbolData.IsReady()}
            
            self.filteredFine = None
            self.monthly_rebalance = False
            
            # if the dictionary is empty, then return    
            if len(filtered_data) < 100: return
            # sort the universe by volatility and select stocks in the top high volatility quintile
            sortedBySharpe = sorted(filtered_data.items(), key=lambda x: x[1].Sharpe(), reverse = True)
            sortedBySharpe = dict(sortedBySharpe)
            
            for i in sortedBySharpe:
                self.Debug( "Sharpe ratio of "+str(i)+"is "+str(sortedBySharpe[i].Sharpe()))
                

class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    
    def __init__(self, symbol, lookback):
        # Std over the last 6 months
        self.symbolSTD = StandardDeviation(lookback)
        
        # Returns over the last 6 months (exluding most recent 10 days)
        self.roc = RateOfChange(lookback-10)
        self.delayed_roc = IndicatorExtensions.Of(Delay(10), self.roc)
        
        self.riskFreeRate = 0.1
        
        # Add warm-up...
        
        
    def Update(self, time, adj_close):
        self.Price = adj_close
        self.symbolSTD.Update(time, adj_close)
        self.roc.Update(time, adj_close)
        

    def IsReady(self):
        return self.symbolSTD.IsReady

        
    def Sharpe(self):
        return (self.delayed_roc.Current.Value-self.riskFreeRate)/self.symbolSTD.Current.Value