from calendar import monthrange

def last_date_of_month(current_date):
    year = current_date.year
    month = current_date.month
    _, lastday = monthrange(year, month)
    
    return datetime(year, month, lastday).date()

def last_year_date(current_date):
    year = current_date.year - 1
    return current_date.replace(year = year, day=1)

class SeasonalitySignals(AlphaModel):
    
    def __init__(self, current_date, num_longs=5, num_shorts=5):
        self.Name = "SeasonalitySignalsAlphaModel"
        self.symbolDataBySymbol = {}
        self.rebalanceDate = last_date_of_month(current_date)
        self.period = timedelta(days=30)
        self.num_longs = num_longs
        self.num_shorts = num_shorts
    
    def Update(self, algorithm, data):
        if algorithm.Time.date() < self.rebalanceDate:
            return []
        
        self.rebalanceDate = last_date_of_month(self.rebalanceDate + timedelta(1))
        algorithm.Log("{}: Updated rebalance date to {}...".format(algorithm.Time, self.rebalanceDate))
        
        algorithm.Log("{}: Creating insights ...".format(algorithm.Time))

        for symbol, tradebar in data.Bars.items():
            period = tradebar.EndTime
            price = tradebar.Close
            if symbol not in self.symbolDataBySymbol:
                algorithm.Log("{}: {} not in symbolDataBySymbol".format(algorithm.Time, str(symbol)))
            else:
                self.symbolDataBySymbol[symbol].ReturnPct.Update(period, price)
        
        annual_returns = [(symbol, symboldata.ReturnPct.Current.Value)
                            for symbol, symboldata in self.symbolDataBySymbol.items()]
        annual_returns = sorted(annual_returns, key=lambda x: x[1], reverse=True)
        
        longs = [Insight.Price(symbol, self.period, InsightDirection.Up)
                    for symbol, _ in annual_returns[:self.num_longs]]
        
        shorts = [Insight.Price(symbol, self.period, InsightDirection.Down)
                    for symbol, _ in annual_returns[-self.num_shorts:]]
        
      
        algorithm.Log("{}: Done creating insights ...".format(algorithm.Time))
        
        return Insight.Group(longs + shorts)

    def OnSecuritiesChanged(self, algorithm, changes):
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.symbolDataBySymbol[symbol] = SymbolData(symbol)
            
            end = last_year_date(algorithm.Time.date())
            start = end - timedelta(5)
            
            history = algorithm.History(symbol, start, end)
    
            self.symbolDataBySymbol[symbol].warmup(history)
        
        for security in changes.RemovedSecurities:
            self.symbolDataBySymbol.pop(security.Symbol, None)
        
        algorithm.Log("{}: Done checking for changes in universe".format(algorithm.Time))
            

class SymbolData:
    def __init__(self, symbol):
        self.Symbol = symbol
        self.ReturnPct = MomentumPercent(1)
    
    def warmup(self, history):
        latest_close = history.unstack(level=0).close
        period, price = latest_close.index[-1], latest_close.iloc[-1].values.item()
        self.ReturnPct.Update(period, price)

def CreateCoarseSelectionFunction(num_coarse):
    def CoarseSelectionFunction(coarse):
        selected = sorted([x for x in coarse if x.Price > 5], key=lambda x: x.DollarVolume, reverse=True)
        selected = selected[:num_coarse]
        return [s.Symbol for s in selected]
       
    return CoarseSelectionFunction

def CreateFineSelectionFunction():
    def FineSelectionFunction(fine):
        return [f.Symbol for f in fine]
    
    return FineSelectionFunction

from Alphas.HistoricalReturnsAlphaModel import HistoricalReturnsAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from QuantConnect.Algorithm.Framework.Selection import FineFundamentalUniverseSelectionModel
import universe
import alphamodel

class TransdimensionalMultidimensionalThrustAssembly(QCAlgorithm):

    def Initialize(self):
        # self.SetStartDate(2013, 1, 1)
        # self.SetEndDate(2013, 12, 31)
        
        self.SetStartDate(2010, 1, 1)
        self.SetEndDate(2010, 3, 31)
        
        # self.SetEndDate(2019, 8, 1)
        self.SetCash(100000)
        
        self.AddAlpha(alphamodel.SeasonalitySignals(self.Time.date()))
        
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        self.SetRiskManagement(TrailingStopRiskManagementModel(0.03))
        
      
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(FineFundamentalUniverseSelectionModel(universe.CreateCoarseSelectionFunction(20),
                                                                        universe.CreateFineSelectionFunction()))
        


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        # if not self.Portfolio.Invested:
        #    self.SetHoldings("SPY", 1)