from AlgorithmImports import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class EquitiesFundamentalSelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self, universe_settings: UniverseSettings = None) -> None:
        self.nextRefreshTimeUtc = datetime.min
        self.month = 0
        self.num_coarse = 500
        self.sectors = [MorningstarSectorCode.FinancialServices, MorningstarSectorCode.RealEstate,
                        MorningstarSectorCode.Healthcare, MorningstarSectorCode.Utilities, MorningstarSectorCode.Technology]
        super().__init__(True, universe_settings)

    '''
    Required method for all fundamental selection model
    The coarse fundamental selection relies on the CoarseFundamental data
    '''
    def SelectCoarse(self, algorithm: QCAlgorithm, coarse: List[CoarseFundamental]) -> List[Symbol]:
        self.nextRefreshTimeUtc = Expiry.EndOfQuarter(algorithm.UtcTime)
        selected = sorted([x for x in coarse if x.Price > 5], 
                          key=lambda x: x.DollarVolume, 
                          reverse=True)

        return [x.Symbol for x in selected[:self.num_coarse]]

    '''
    Optional method for a fundamental selection model
    The fine fundamental selection relies on the FineFundamental data
    '''
    def SelectFine(self, algorithm: QCAlgorithm, fine: List[FineFundamental]) -> List[Symbol]:
        filtered_fine = [x.Symbol for x in fine
                         if x.SecurityReference.IPODate + timedelta(5*365) < algorithm.Time and
                         x.AssetClassification.MorningstarSectorCode in self.sectors and
                         x.OperationRatios.ROE.Value > 0 and
                         x.OperationRatios.NetMargin.Value > 0 and
                         x.ValuationRatios.PERatio > 0]

        return filtered_fine

    '''
    Ensure that this model selection only runs once each quarter
    '''
    def GetNextRefreshTimeUtc(self):
        return self.nextRefreshTimeUtc

from AlgorithmImports import *

class FundamentalFactorAlphaModel(AlphaModel):

    def __init__(self) -> None:
        self.rebalanceTime = datetime.min
        self.sectors = {}

    def Update(self, algorithm: QCAlgorithm, data: Slice) -> Iterable[Insight]:
        insights = []

        if algorithm.Time < self.rebalanceTime:
            return insights
        self.rebalanceTime = Expiry.EndOfQuarter(algorithm.Time)

        for sector in self.sectors:
            securities = self.sectors[sector]
            sortedByROE = sorted(cast(
                Iterable[Security], securities), key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=True)
            sortedByPM = sorted(cast(
                Iterable[Security], securities), key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=True)
            sortedByPE = sorted(cast(
                Iterable[Security], securities), key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=True)

        scores = {}

        for security in cast(Iterable[Security], securities):
            score = sum([sortedByROE.index(security), sortedByPM.index(
                security), sortedByPE.index(security)])
            scores[security] = score

        length = max(int(len(scores)/5), 1)
        for security in sorted(scores.items(), key=lambda x: x[1], reverse=False)[:length]:
            symbol = security[0].Symbol
            insights.append(Insight.Price(
                symbol, Expiry.EndOfQuarter, InsightDirection.Up))

        return insights

    def OnSecuritiesChanged(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.RemovedSecurities:
            for sector in self.sectors:
                if security in self.sectors[sector]:
                    self.sectors[sector].remove(security)

        for security in changes.AddedSecurities:
            sector = security.Fundamentals.AssetClassification.MorningstarSectorCode
            if sector not in self.sectors:
                self.sectors[sector] = set()
            self.sectors[sector].add(security)

from AlgorithmImports import *
from EquitiesFundamentalSelectionModel import *
from FundamentalFactorAlphaModel import *

class EquitiesTemplate(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2013, 1, 1)  # Set Start Date
        self.SetEndDate(2014, 1, 1)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(EquitiesFundamentalSelectionModel(self.UniverseSettings))
        
        self.AddAlpha(FundamentalFactorAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda t: Expiry.EndOfQuarter(t)))