# https://quantpedia.com/strategies/short-interest-effect-long-only-version/
#
# All stocks from NYSE, AMEX, and NASDAQ are part of the investment universe. The short-interest ratio is used as the predictor variable. 
# Stocks are sorted based on their short interest ratio, and the first percentile is held. The portfolio is equally weighted and rebalanced monthly.
# 
# QC Implementation changes:
#   - Universe consists of 500 most liquid stocks from NYSE, AMEX and NASDAQ.

from AlgorithmImports import *
from io import StringIO
from typing import List, Dict
from numpy import isnan

class ShortInterestEffect(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2017, 1, 1)
        self.SetCash(100_000)

        self.tickers_to_ignore: List[str] = ['GEN']

        self.quantile: int = 10
        self.leverage: int = 5

        self.weight: Dict[Symbol, float] = {}

        # source: https://www.finra.org/finra-data/browse-catalog/equity-short-interest/data
        text: str = self.Download('data.quantpedia.com/backtesting_data/economic/short_volume.csv')

        self.short_volume_df: DataFrame = pd.read_csv(StringIO(text), delimiter=';')
        self.short_volume_df['date'] = pd.to_datetime(self.short_volume_df['date']).dt.date
        self.short_volume_df.set_index('date', inplace=True)

        self.fundamental_count: int = 500
        self.fundamental_sorting_key = lambda x: x.DollarVolume

        self.selection_flag: bool = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelectionFunction)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.
        self.settings.daily_precise_end_time = False
        self.recent_month = -1

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            security.SetFeeModel(CustomFeeModel())
            security.SetLeverage(self.leverage)

    def FundamentalSelectionFunction(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # monthly rebalance
        if self.recent_month == self.Time.month:
            return Universe.Unchanged
        self.recent_month = self.Time.month
        self.selection_flag = True

        # check last date on custom data
        if self.Time.date() > self.short_volume_df.index[-1] or self.Time.date() < self.short_volume_df.index[0]:
            self.Liquidate()
            return Universe.Unchanged

        selected: List[Fundamental] = [
            x for x in fundamental 
            if x.HasFundamentalData 
            and x.Market == 'usa'
            and x.CompanyProfile.SharesOutstanding != 0
            and x.Symbol.Value in self.short_volume_df.columns
            and x.Symbol.Value not in self.tickers_to_ignore
        ]

        if len(selected) > self.fundamental_count:
            selected = [x for x in sorted(selected, key=self.fundamental_sorting_key, reverse=True)[:self.fundamental_count]]

        short_interest: Dict[Symbol, float] = {}

        # calculate short interest
        for stock in selected:
            symbol: Symbol = stock.Symbol
            ticker: str = symbol.Value

            if ticker in self.short_volume_df.columns:
                if isnan(self.short_volume_df[self.short_volume_df.index <= self.Time.date()][ticker][-1]):
                    continue
                short_interest[symbol] = self.short_volume_df[self.short_volume_df.index <= self.Time.date()][ticker][-1] / stock.CompanyProfile.SharesOutstanding

        if len(short_interest) >= self.quantile:
            # sorting by short interest ratio
            sorted_short_interest: List[Symbol] = sorted(short_interest, key = short_interest.get)
            quantile: int = int(len(sorted_short_interest) / self.quantile)
            long: List[Symbol] = sorted_short_interest[:quantile]

            # equally weighting
            for symbol in long:
                self.weight[symbol] = 1 / len(long)

        return list(self.weight.keys())

    def OnData(self, data: Slice) -> None:
        if not self.selection_flag:
            return
        self.selection_flag = False
        
        # trade execution
        portfolio:List[PortfolioTarget] = [PortfolioTarget(symbol, w) for symbol, w in self.weight.items() if symbol in data and data[symbol]]
        self.SetHoldings(portfolio, True)
        self.weight.clear()

# Custom fee model.
class CustomFeeModel(FeeModel):
    def GetOrderFee(self, parameters: OrderFeeParameters) -> OrderFee:
        fee: float = parameters.Security.Price * parameters.Order.AbsoluteQuantity * 0.00005
        return OrderFee(CashAmount(fee, "USD"))