# region imports
from AlgorithmImports import *

from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
# endregion

class DeterminedTanKitten(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 1, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash

        self.AddUniverseSelection(FSTopMarketCapUniverseSelectionModel(sector = MorningstarSectorCode.FinancialServices,
                                                                    number = 3,
                                                                    universe_settings = self.UniverseSettings))

    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        for security in changes.AddedSecurities:
            self.Debug(f"Added {security.Symbol}")

        for security in changes.RemovedSecurities:
            self.Debug(f"Removed {security.Symbol}")

    # def OnData(self, data: Slice):
    #     if not self.Portfolio.Invested:
    #         self.SetHoldings("SPY", 0.33)
    #         self.SetHoldings("BND", 0.33)
    #         self.SetHoldings("AAPL", 0.33)


class FSTopMarketCapUniverseSelectionModel(FineFundamentalUniverseSelectionModel):
    
    def __init__(self, sector: MorningstarSectorCode, number: int,
                universe_settings: UniverseSettings = None) -> None:
        super().__init__(self.SelectCoarse, self.SelectFine, universe_settings)
        self.sector = sector
        self.number = number

    def SelectCoarse(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        #1. Filt to securities with fundamental data
        return [c.Symbol for c in coarse if c.HasFundamentalData]
    
    def SelectFine(self, fine: List[FineFundamental]) -> List[Symbol]:
        #2. Select financial sector
        filtered_fine = [x for x in fine if x.AssetClassification.MorningstarSectorCode == self.sector]
        #3. Order by market cap descending. 
        sorted_by_mkcap = sorted(filtered_fine, key=lambda x: x.MarketCap, reverse=True)
        #4. From different companies
        companyId = {}
        for c in sorted_by_mkcap:
            identifier = c.CompanyReference.CompanyId
            if not(companyId.get(identifier)):
                companyId[identifier] = c.Symbol 

            if len(companyId)>=self.number:
                return list(companyId.values())
        return list(companyId.values()) # Return Top "number" assets by highest Market Cap in fianancial sector