# https://quantpedia.com/strategies/investment-factor/
#
# The investment universe consists of all NYSE, Amex, and NASDAQ stocks. Firstly, stocks are allocated to five Size groups (Small to Big) at the end of each June
# using NYSE market cap breakpoints. Stocks are allocated independently to five Investment (Inv) groups (Low to High) still using NYSE breakpoints.
# The intersections of the two sorts produce 25 Size-Inv portfolios. For portfolios formed in June of year t, Inv is the growth of total assets for
# the fiscal year ending in t-1 divided by total assets at the end of t-1. Long portfolio with the highest Size and simultaneously with the lowest
# Investment. Short portfolio with the highest Size and simultaneously with the highest Investment. The portfolios are value-weighted.
#
# QC implementation changes:
#   - The investment universe consists of 3000 largest US stock traded on NYSE, AMEX and NASDAQ with price >= 1$.
#   - The portfolios are equally-weighted.

from AlgorithmImports import *

class InvestmentFactor(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetCash(100_000)

        self.symbol: Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        
        self.fundamental_count: int = 3_000
        self.fundamental_sorting_key = lambda x: x.MarketCap

        self.long: List[Symbol] = []
        self.short: List[Symbol] = []

        self.quantile: int = 5
        self.leverage: int = 3
        self.rebalance_month: int = 6
        self.min_share_price: float = 1.
        
        self.weight: Dict[Symbol, float] = {}
        
        self.selection_flag: bool = False
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.FundamentalSelectionFunction)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0
        self.Schedule.On(self.DateRules.MonthEnd(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol), self.Selection)

        self.settings.daily_precise_end_time = False

    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]:
        if not self.selection_flag:
            return Universe.Unchanged

        selected: List[Fundamental] = [
            x for x in fundamental if x.HasFundamentalData and x.Market == 'usa' and x.MarketCap != 0 and x.Price >= self.min_share_price and \
            ((x.SecurityReference.ExchangeId == "NYS") or (x.SecurityReference.ExchangeId == "NAS") or (x.SecurityReference.ExchangeId == "ASE")) and \
            not np.isnan(x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths) and not np.isnan(x.OperationRatios.TotalAssetsGrowth.OneYear) and \
            x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths != 0 and x.OperationRatios.TotalAssetsGrowth.OneYear != 0
        ]
        
        if len(selected) > self.fundamental_count:
            selected = [x for x in sorted(selected, key=self.fundamental_sorting_key, reverse=True)[:self.fundamental_count]]

        # Sorting by investment factor.
        sorted_by_inv_factor: List[Fundamental] = sorted(selected, key = lambda x: (x.OperationRatios.TotalAssetsGrowth.OneYear / x.FinancialStatements.BalanceSheet.TotalAssets.TwelveMonths), reverse=True)
        
        if len(sorted_by_inv_factor) >= self.quantile:
            quantile: int = int(len(sorted_by_inv_factor) / self.quantile)
            self.long = [x.Symbol for x in sorted_by_inv_factor[-quantile:]]
            self.short = [x.Symbol for x in sorted_by_inv_factor[:quantile]]
        
        return self.long + self.short
    
    def OnData(self, slice: Slice) -> None:
        if not self.selection_flag:
            return
        self.selection_flag = False

        # order execution
        targets: List[PortfolioTarget] = []
        for i, portfolio in enumerate([self.long, self.short]):
            for symbol in portfolio:
                if symbol in slice and slice[symbol]:
                    targets.append(PortfolioTarget(symbol, ((-1) ** i) / len(portfolio)))
        
        self.SetHoldings(targets, True)

        self.long.clear()
        self.short.clear()
    
    def Selection(self) -> None:
        if self.Time.month == self.rebalance_month:
            self.selection_flag = True

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