from AlgorithmImports import *
from datetime import *


class EquitiesFundamentalSelectionModel(FineFundamentalUniverseSelectionModel):

    def __init__(self, universe_settings: UniverseSettings = None) -> None:
        self.month = 0
        self.num_coarse = 500
        super().__init__(self.SelectCoarse, self.SelectFine, universe_settings)

    '''
    Required method for all fundamental selection model
    The coarse fundamental selection relies on the CoarseFundamental data
    '''

    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        if not self.IsRebalanceDue(self.Time):
            return Universe.Unchanged

        selected = sorted([x for x in cast(Iterable[CoarseFundamental], coarse)
                          if x.HasFundamentalData and 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, fine: List[FineFundamental]) -> List[Symbol]:

        sectors = [MorningstarSectorCode.FinancialServices, MorningstarSectorCode.RealEstate,
                   MorningstarSectorCode.Healthcare, MorningstarSectorCode.Utilities, MorningstarSectorCode.Technology]

        filtered_fine = [x.Symbol for x in cast(Iterable[FineFundamental], fine)
                         if x.SecurityReference.IPODate + timedelta(5*365) < datetime and
                         x.AssetClassification.MorningstarEconomySphereCode in 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 IsRebalanceDue(self, time):
        if time.month == self.month or time.month not in [1, 4, 7, 10]:
            return None
        
        self.month = time.month
        return date

from AlgorithmImports import *
from datetime import *


class FundamentalFactorAlphaModel(AlphaModel):

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

    def Update(self, algorithm: QCAlgorithm, data: Slice) -> Iterable[Insight]:
        if algorithm.Time < self.rebalanceTime:
            return[]
        self.rebalanceTime = Expiry.EndOfQuarter(algorithm.Time)

        insights = []

        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 cast(Iterable[Security], 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.UniverseSelection = EquitiesFundamentalSelectionModel(self)
        self.AddUniverseSelection(self.UniverseSelection)

        self.AddAlpha(FundamentalFactorAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(
            EquitiesFundamentalSelectionModel.IsRebalanceDue))
        self.SetRiskManagement(NullRiskManagementModel)
        self.SetExecution(ImmediateExecutionModel())