#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/Screener/Details/77


class BetaFactorInStocks(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2018, 1, 1)   
        self.set_end_date(2022, 1, 1)         
        self.set_cash(1000000)            

        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function)

        # Add Wilshire 5000 Total Market Index data from Dropbox 
        self._price5000 = self.add_data(Fred, Fred.Wilshire.PRICE_5000, Resolution.DAILY).symbol

        # Setup a RollingWindow to hold market return
        self._market_return = RollingWindow[float](252)
        # Use a ROC indicator to convert market price index into return, and save it to the RollingWindow
        self._roc = self.roc(self._price5000, 1)
        self._roc.updated += lambda _, updated: self._market_return.add(updated.value)
        # Warm up
        hist = self.history(self._price5000, 253, Resolution.DAILY)
        for t, value in hist.loc[self._price5000]['value'].items():
            self._roc.update(t, value)

        self._data = {}
        self._monthly_rebalance = False
        self._long = None
        self._short = None
        
        spy = Symbol.create("SPY", SecurityType.EQUITY, Market.USA)
        self.schedule.on(self.date_rules.month_start(spy), self.time_rules.after_market_open(spy), self._rebalance)
        self.set_warm_up(timedelta(365))

    def _coarse_selection_function(self, coarse):
        for c in coarse:
            if c.symbol not in self._data:
                self._data[c.symbol] = SymbolData(c.symbol)
            self._data[c.symbol].update(c.end_time, c.adjusted_price)

        if self._monthly_rebalance:
            filtered_data = {symbol: data for symbol, data in self._data.items() if data.last_price > 5 and data.is_ready()}

            if len(filtered_data) > 10:
                # sort the dictionary in ascending order by beta value
                sorted_beta = sorted(filtered_data, key=lambda x: filtered_data[x].beta(self._market_return))
                self._long = sorted_beta[:5]
                self._short = sorted_beta[-5:]
                return self._long + self._short
            else: 
                self._monthly_rebalance = False
        return []

    def _rebalance(self):
        self._monthly_rebalance = True

    def on_data(self, data):
        if not self._monthly_rebalance or self.is_warming_up: 
            return 
        
        # Liquidate symbols not in the universe anymore
        for symbol, security_holding in self.portfolio.items():
            if security_holding.invested and symbol not in self._long + self._short:
                self.liquidate(symbol)

        if self._long is None or self._short is None: 
            return
        
        longs = [symbol for symbol in self._long if self.securities[symbol].price]
        long_scale_factor = 0.4/sum(range(1,len(longs)+1))
        for rank, symbol in enumerate(longs):    
            self.set_holdings(symbol, (len(longs)-rank+1)*long_scale_factor)
        
        shorts = [symbol for symbol in self._short if self.securities[symbol].price]
        short_scale_factor = 0.4/sum(range(1,len(shorts)+1))
        for rank, symbol in enumerate(shorts):
            self.set_holdings(symbol, -(rank+1)*short_scale_factor)
            
        self._monthly_rebalance = False
        self._long = None
        self._short = None


class SymbolData:

    def __init__(self, symbol):
        self.symbol = symbol
        self.last_price = 0
        self._returns = RollingWindow[float](252)
        self._roc = RateOfChange(1)
        self._roc.updated += lambda _, updated: self._returns.add(updated.value)
        
    def update(self, time, price):
        if price != 0:
            self.last_price = price
            self._roc.update(time, price)
    
    def is_ready(self):
        return self._roc.is_ready and self._returns.is_ready
    
    def beta(self, market_ret):
        asset_return = np.array(list(self._returns), dtype=np.float32)
        market_return = np.array(list(market_ret), dtype=np.float32)
        return np.cov(asset_return, market_return)[0][1] / np.var(market_return)