#region imports
from AlgorithmImports import *
from QuantConnect.DataSource import *
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
import math
import numpy as np
import pandas as pd
import scipy as sp

class BasicTemplateAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 2, 1)    # Set End Date
        self.SetCash(100000)          # Set Strategy Cash
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.num_fine = 30
        self.period_sma= 21
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time:None))
        self.timedelta = timedelta(91)
        self.NextEntryTime = datetime.min
        self.sorted_stock = []


    def CoarseSelectionFunction(self, coarse):
        return [c.Symbol for c in coarse if c.HasFundamentalData]

    def FineSelectionFunction(self, fine):
        if self.Time < self.NextEntryTime:
            return Universe.Unchanged
        
        SoartedByMktCap = sorted(fine, key=lambda x: x.MarketCap, reverse=True)[:50] #should return the SP500
        fine = [x for x in SoartedByMktCap if x.EarningReports.TotalDividendPerShare.ThreeMonths
                            and x.ValuationRatios.PriceChange1M
                            and x.ValuationRatios.FCFYield ]

        sortedByfactor1 = sorted(fine, key=lambda x: x.EarningReports.TotalDividendPerShare.ThreeMonths, reverse=True)
        sortedByfactor2 = sorted(fine, key=lambda x: x.ValuationRatios.PriceChange1M, reverse=False)
        sortedByfactor3 = sorted(fine, key=lambda x: x.ValuationRatios.FCFYield, reverse=True)
        stock_dict = {}

        for rank1, ele in enumerate(sortedByfactor1):
            rank2 = sortedByfactor2.index(ele)
            rank3 = sortedByfactor3.index(ele)
            stock_dict[ele] = rank1/3 + rank2/3 + rank3/3

        #sorting for the best stocks out of all 3 conditions
        self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1], reverse=True)
        self.sorted_stock = [x[0].Symbol for x in self.sorted_stock[:self.num_fine]]
        return [symbol for symbol in self.sorted_stock]


    def OnData(self, slice):
        if self.Time < self.NextEntryTime: 
            return
        for kvp in self.Portfolio:
            symbol = kvp.Key
            security_holding = kvp.Value
            if (security_holding.Invested) and (symbol not in self.sorted_stock):
                self.Liquidate(symbol)                
    
        weight = 1/self.num_fine
        for symbol in self.sorted_stock:
            if symbol in slice.Bars:
                self.SetHoldings(symbol, weight)
            

        self.NextEntryTime = self.Time+self.timedelta #next check in a quarter