#region imports
from AlgorithmImports import *

from decimal import Decimal
from math import floor
#endregion
# https://quantpedia.com/Screener/Details/53


class SentimentAndStyleRotationAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2010, 1, 1)
        self.set_end_date(2018, 7, 1)
        self.set_cash(10000000)

        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        
        self.add_data(CBOE, "VIX", Resolution.DAILY)
        self.add_data(PutCallRatio, "PutCallRatio", Resolution.DAILY)
        
        self._vix_sma_1 = SimpleMovingAverage(21)
        self._vix_sma_6 = SimpleMovingAverage(21*6)
        self._pc_ratio_sma_1 = SimpleMovingAverage(21)
        self._pc_ratio_sma_6 = SimpleMovingAverage(21*6)
        # initialize the indicator with the history request
        pc_ratio_history = self.history(["PutCallRatio"], 21*10, Resolution.DAILY)
        vix_history = self.history(["VIX"], 21*10, Resolution.DAILY)

        for t, value in vix_history.loc['VIX']['value'].items():
            self._vix_sma_6.update(t, value)
            self._vix_sma_1.update(t, value)
        for t, value in pc_ratio_history.loc['PutCallRatio']['value'].items():
            self._pc_ratio_sma_1.update(t, value)
            self._pc_ratio_sma_6.update(t, value)

        self.add_universe(self._coarse_selection_function, self._fine_selection_function)
        self.add_equity("SPY", Resolution.DAILY)
        
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.at(0, 0), self._rebalance)
        
        self._month_start = False
        self._selection = False
        self._months = -1        
    
    def _coarse_selection_function(self, coarse):
        if self._month_start:
            # drop stocks which have no fundamental data or have low price
            return [x.symbol for x in coarse if (x.has_fundamental_data)]
        else: 
            return Universe.UNCHANGED

    def _fine_selection_function(self, fine):
        if self._month_start:
            self._selection = True

            fine = [i for i in fine if i.earning_reports.basic_average_shares.three_months>0
                                    and i.earning_reports.basic_eps.twelve_months>0
                                    and i.valuation_ratios.pe_ratio>0
                                    and i.valuation_ratios.pb_ratio>0]
            
            sorted_market_cap = sorted(fine, key=lambda x: x.market_cap, reverse=True)
            decile_top1 = sorted_market_cap[:floor(len(sorted_market_cap)/10)]
            decile_top2 = sorted_market_cap[floor(len(sorted_market_cap)/10):floor(len(sorted_market_cap)*2/10)]
            decile_top3 = sorted_market_cap[floor(len(sorted_market_cap)*2/10):floor(len(sorted_market_cap)*3/10)]
            sorted_pb1 = sorted(decile_top1, key=lambda x: x.valuation_ratios.pb_ratio)
            sorted_pb2 = sorted(decile_top2, key=lambda x: x.valuation_ratios.pb_ratio)
            sorted_pb3 = sorted(decile_top3, key=lambda x: x.valuation_ratios.pb_ratio)
            # The value portfolio consists of all firms included in the quintile with the lowest P/B ratio
            pb_bottom1 = sorted_pb1[:floor(len(decile_top1)/5)]
            pb_bottom2 = sorted_pb2[:floor(len(decile_top2)/5)]
            pb_bottom3 = sorted_pb3[:floor(len(decile_top3)/5)]
            self._value_portfolio = [i.symbol for i in pb_bottom1 + pb_bottom2 + pb_bottom3]   

            # The growth portfolio consists of all firms included in the quintile with the highest P/B ratio
            pb_top1 = sorted_pb1[-floor(len(decile_top1)/5):]
            pb_top2 = sorted_pb2[-floor(len(decile_top2)/5):]
            pb_top3 = sorted_pb3[-floor(len(decile_top3)/5):]
            self._growth_portfolio = [i.symbol for i in pb_top1 + pb_top2 + pb_top3]   

            return self._value_portfolio + self._growth_portfolio
        else:
            return Universe.UNCHANGED

    def _rebalance(self):
        # rebalance every three months
        self._months += 1
        if self._months%3 == 0:
            self._month_start = True
            
    def _get_tradable_assets(self, symbols):
        tradable_assets = []
        for symbol in symbols:
            security = self.securities[symbol]
            if security.price and security.is_tradable:
                tradable_assets.append(symbol)
        return tradable_assets

    def on_data(self, data):
        if (not data.contains_key("VIX") or not data.contains_key("PutCallRatio") or 
            data["VIX"].value == 0 or data["PutCallRatio"].value == 0): 
            return

        self._vix_sma_1.update(self.time, data["VIX"].value)
        self._vix_sma_6.update(self.time, data["VIX"].value)
        self._pc_ratio_sma_1.update(self.time, data["PutCallRatio"].value)
        self._pc_ratio_sma_6.update(self.time, data["PutCallRatio"].value)

        if self._month_start and self._selection: 
            self._month_start = False
            self._selection = False

            self._value_portfolio = self._get_tradable_assets(self._value_portfolio)
            self._growth_portfolio = self._get_tradable_assets(self._growth_portfolio)

            stocks_invested = [x.key for x in self.portfolio if x.value.invested]
            for i in stocks_invested:
                if i not in self._value_portfolio+self._growth_portfolio:
                    self.liquidate(i)

            if self._vix_sma_1.current.value > self._vix_sma_6.current.value:
                if self._pc_ratio_sma_1.current.value < self._pc_ratio_sma_6.current.value:
                    long_weight = 1/len(self._value_portfolio)
                    for long_ in self._value_portfolio:
                        self.set_holdings(long_, long_weight)
                elif self._pc_ratio_sma_1.current.value > self._pc_ratio_sma_6.current.value:
                    short_weight = 1/len(self._value_portfolio)
                    for short in self._value_portfolio:
                        self.set_holdings(short, -short_weight)              

            else:
                long_weight = 1/len(self._value_portfolio+self._growth_portfolio)
                for long_ in self._value_portfolio+self._growth_portfolio:
                    self.set_holdings(long_, long_weight)

       
class PutCallRatio(PythonData):

    def get_source(self, config, date, is_live_mode):
        return SubscriptionDataSource(
            "https://cdn.cboe.com/resources/options/volume_and_call_put_ratios/totalpc.csv", 
            SubscriptionTransportMedium.REMOTE_FILE)

    def reader(self, config, line, date, is_live_mode):
        if not (line.strip() and line[0].isdigit()): 
            return None
        index = CBOE()
        index.symbol = config.symbol

        try:
            # Example File Format:
            # DATE       CALL      PUT       TOTAL      P/C Ratio
            # 11/1/06    976510    623929    1600439    0.64
            data = line.split(',')
            index.time = datetime.strptime(data[0], "%m/%d/%Y").strftime("%Y-%m-%d")
            index.value = Decimal(data[4])

        except ValueError:
            return None

        return index