#region imports
from AlgorithmImports import *

from collections import deque
import calendar
#endregion


class BetaAlphaModel(AlphaModel):

    def __init__(self):
        self.name = 'BetaAlphaModel'
        self._assets = {}
        self._month = datetime.min
        
    def update(self, algorithm, data):
        if data.contains_key("SPY"):
            self._market_price.append(float(algorithm.securities["SPY"].price))
        for key, value in self._assets.items():
            if data.contains_key(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].is_long and i not in long_:  
                        insights.append(Insight.price(i, timedelta(days=1), InsightDirection.FLAT))
                    if algorithm.portfolio[i].is_short 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 on_securities_changed(self, algorithm, changes):
        for added in changes.added_securities:
            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.removed_securities:
            if removed.symbol in self._assets:
                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)
        

#region imports
from AlgorithmImports import *

from alpha import BetaAlphaModel
from portfolio import MonthlyPortfolioConstructionModel
#endregion


class BetaAlgorithm(QCAlgorithmFramework):

    def initialize(self):

        self.set_start_date(2012, 1, 1)   #Set Start Date
        self.set_end_date(2018, 3, 1)    #Set End Date
        self.set_cash(100000)            #Set Strategy Cash

        self.universe_settings.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.set_universe_selection(ManualUniverseSelectionModel(symbols))
        self.set_alpha(BetaAlphaModel())
        self.set_portfolio_construction(MonthlyPortfolioConstructionModel())
        self.set_execution(ImmediateExecutionModel())
        self.set_risk_management(NullRiskManagementModel())
        self.add_equity("SPY", Resolution.DAILY)
        

#region imports
from AlgorithmImports import *
#endregion


class MonthlyPortfolioConstructionModel(PortfolioConstructionModel):

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

    def create_targets(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