from Alphas.ConstantAlphaModel import ConstantAlphaModel
from AlphaModel import BetaAlphaModel
from Risk.NullRiskManagementModel import NullRiskManagementModel
from datetime import datetime
from PortfolioConstructionModel import MonthlyPortfolioConstructionModel

class BetaAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        self.SetStartDate(2012, 1, 1)   #Set Start Date
        self.SetEndDate(2018, 3, 1)    #Set End Date
        self.SetCash(100000)            #Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        tickers =  [
                    "EWJ",  # iShares MSCI Japan Index ETF
                    "EZU",  # iShares MSCI Eurozone ETF
                    "EFNL", # iShares MSCI Finland Capped Investable Market Index ETF
                    "EWW",  # iShares MSCI Mexico Inv. Mt. Idx
                    "ERUS", # iShares MSCI Russia ETF
                    "IVV",  # iShares S&P 500 Index
                    "ICOL", # Consumer Discretionary Select Sector SPDR Fund
                    "AAXJ", # iShares MSCI All Country Asia ex Japan Index ETF
                    "AUD",  # Australia Bond Index Fund
                    "EWQ",  # iShares MSCI France Index ETF
                    "BUND", # Pimco Germany Bond Index Fund
                    "EWH",  # iShares MSCI Hong Kong Index ETF
                    "EPI",  # WisdomTree India Earnings ETF
                    "EIDO"  # iShares MSCI Indonesia Investable Market Index ETF
                    "EWI",  # iShares MSCI Italy Index ETF
                    "GAF",  # SPDR S&P Emerging Middle East & Africa ETF
                    "ENZL", # iShares MSCI New Zealand Investable Market Index Fund
                    "NORW"  # Global X FTSE Norway 30 ETF
                    "EWY",  # iShares MSCI South Korea Index ETF
                    "EWP",  # iShares MSCI Spain Index ETF
                    "EWD",  # iShares MSCI Sweden Index ETF
                    "EWL",  # iShares MSCI Switzerland Index ETF
                    "GXC",  # SPDR S&P China ETF
                    "EWC",  # iShares MSCI Canada Index ETF
                    "EWZ",  # iShares MSCI Brazil Index ETF
                    "ARGT", # Global X FTSE Argentina 20 ETF
                    "AND",  # Global X FTSE Andean 40 ETF
                    "AIA",  # iShares S&P Asia 50 Index ETF
                    "EWO",  # iShares MSCI Austria Investable Mkt Index ETF
                    "EWK",  # iShares MSCI Belgium Investable Market Index ETF
                    "BRAQ", # Global X Brazil Consumer ETF
                    "ECH",  # iShares MSCI Chile Investable Market Index ETF
                    "CHIB", # Global X China Technology ETF
                    "EGPT", # Market Vectors Egypt Index ETF
                    "ADRU"] # BLDRS Europe 100 ADR Index ETF
        symbols = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in tickers]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.SetAlpha(BetaAlphaModel())
        self.SetPortfolioConstruction(MonthlyPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
        self.AddEquity("SPY", Resolution.Daily)

from datetime import timedelta, datetime
from QuantConnect.Algorithm.Framework.Alphas import AlphaModel, Insight, InsightType, InsightDirection
import numpy as np
from collections import deque
import calendar

class BetaAlphaModel(AlphaModel):

    def __init__(self):
        self.period = timedelta(hours=6)
        self.Name = 'BetaAlphaModel'
        self.assets = {}
        self.month = datetime.min
        
    def Update(self, algorithm, data):
        if data.ContainsKey("SPY"):
            self.market_price.append(float(algorithm.Securities["SPY"].Price))
        for key, value in self.assets.items():
            if data.ContainsKey(key):
                value.Price.append(float(algorithm.Securities[key].Price))

        insights = []

        if self.month != algorithm.Time.month:    
            self.month = algorithm.Time.month 
            beta_values = {} 
            market_return = np.diff(np.array(self.market_price))/np.array(self.market_price)[:-1]
            long = None
            for key, value in self.assets.items():
                if key != "SPY" and len(value.Price) == value.Price.maxlen:
                    asset_return = np.diff(np.array(value.Price))/np.array(value.Price)[:-1]
                    beta_values[key] = self.beta(asset_return, market_return)
            sorted_by_beta = sorted(beta_values, key = lambda x: beta_values[x])
            long = sorted_by_beta[:int(0.25*len(sorted_by_beta))]
            short = sorted_by_beta[-int(0.25*len(sorted_by_beta)):]
            
            # day: the weekday of first day of the month 
            # num_days: number of days in month  
            day, num_days = calendar.monthrange(algorithm.Time.year, algorithm.Time.month)
            insight_period = num_days - algorithm.Time.day - 1
            if long and short:
                invested = [x.Key for x in algorithm.Portfolio if x.Value.Invested]
                for i in invested:
                    if algorithm.Portfolio[i].IsLong and i not in long:  
                        insights.append(Insight.Price(i, timedelta(days=1), InsightDirection.Flat))
                    if algorithm.Portfolio[i].IsShort and i not in short:  
                        insights.append(Insight.Price(i, timedelta(days=1), InsightDirection.Flat))
                        
                for i in long:
                    insights.append(Insight.Price(i, timedelta(days=insight_period), InsightDirection.Up))
                for i in short:
                    insights.append(Insight.Price(i, timedelta(days=insight_period), InsightDirection.Down))
           
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            if added.Symbol.Value == "SPY":
                self.market_price = deque(maxlen=253)
                hist_SPY = algorithm.History(["SPY"], 500, Resolution.Daily)
                for i in hist_SPY.loc["SPY"].itertuples():
                    self.market_price.append(i.close)     

            if added not in self.assets and added.Symbol.Value != "SPY":
                hist = algorithm.History([added.Symbol.Value], 500, Resolution.Daily)
                if not hist.empty:
                    self.assets[added.Symbol] = SymbolData(added)
                    for i in hist.loc[added.Symbol.Value].itertuples():
                        self.assets[added.Symbol].Price.append(i.close)                

        for removed in changes.RemovedSecurities:
            self.assets.pop(removed.Symbol)

    def beta(self, asset_return, market_return):
        asset_return = np.array(asset_return, dtype=np.float32)
        market_return = np.array(market_return, dtype=np.float32)
        return np.cov(asset_return, market_return)[0][1]/np.var(market_return)
    
class SymbolData:
    def __init__(self, symbol):
        self.Symbol = symbol
        self.Price = deque(maxlen=253)

from datetime import datetime

class MonthlyPortfolioConstructionModel(PortfolioConstructionModel):

    def __init__(self):
        self.month = datetime.min

    def CreateTargets(self, algorithm, insights):
        if self.month == algorithm.Time.month:
            return []
        
        targets = []
        self.month = algorithm.Time.month
        if len(insights) != 0:
            percent = 1 / len(insights)
            for insight in insights:
                target = PortfolioTarget.Percent(algorithm, insight.Symbol, insight.Direction * percent)
                if not target is None:
                    targets.append(target)
        return targets