#region imports
from AlgorithmImports import *
#endregion
class PortfolioRebalanceOnCustomFuncRegressionAlgorithm(QCAlgorithm):
    def Initialize(self):
        ''' Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        self.UniverseSettings.Resolution = Resolution.Daily



        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2018, 1, 1)

        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.num_coarse = 500
        self.month= 0

        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        
        self.AddAlpha(FundamentalFactorAlphaModel())
        self.SetPortfolioConstruction(MyEqualWeightingPortfolioConstructionModel(self.IsRebalanceDue, ag=self))
        self.SetExecution(ImmediateExecutionModel())
   

    def CoarseSelectionFunction(self, coarse):
        # If not time to rebalance, keep the same universe
        if not self.IsRebalanceDue(self.Time): 
            return Universe.Unchanged

        # Select only those with fundamental data and a sufficiently large price
        # Sort by top dollar volume: most liquid to least liquid
        selected = sorted([x for x in 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]]


    def FineSelectionFunction(self, fine):
        # Filter the fine data for equities that IPO'd more than 5 years ago in selected sectors
        
        sectors = [
            MorningstarSectorCode.FinancialServices,
            MorningstarSectorCode.RealEstate,
            MorningstarSectorCode.Healthcare,
            MorningstarSectorCode.Utilities,
            MorningstarSectorCode.Technology]
        
        filtered_fine = [x.Symbol for x in fine if x.SecurityReference.IPODate + timedelta(365*5) < self.Time
                                    and x.AssetClassification.MorningstarSectorCode in sectors
                                    and x.OperationRatios.ROE.Value > 0.15
                                    and x.OperationRatios.NetMargin.Value > 0
                                    and 30 > x.ValuationRatios.PERatio > 0]
        # self.Debug(f'{self.Time} len(filtered_fine) {len(filtered_fine)}')
                
        return filtered_fine

    def IsRebalanceDue(self, time):
        # Rebalance on the first day of the Quarter
        if time.month == self.month or time.month not in [1, 4, 7, 10]:
            return None
            
        self.month = time.month
        self.Debug(f'{self.Time} utc {self.UtcTime}, IsRebalanceDue time {time}')
        return time



class MyEqualWeightingPortfolioConstructionModel(PortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that gives equal weighting to all securities.
    The target percent holdings of each security is 1/N where N is the number of securities.
    For insights of direction InsightDirection.Up, long targets are returned and
    for insights of direction InsightDirection.Down, short targets are returned.'''

    def __init__(self, rebalance = Resolution.Daily, portfolioBias = PortfolioBias.LongShort, ag=None):
        '''Initialize a new instance of EqualWeightingPortfolioConstructionModel
        Args:
            rebalance: Rebalancing parameter. If it is a timedelta, date rules or Resolution, it will be converted into a function.
                              If None will be ignored.
                              The function returns the next expected rebalance time for a given algorithm UTC DateTime.
                              The function returns null if unknown, in which case the function will be called again in the
                              next loop. Returning current time will trigger rebalance.
            portfolioBias: Specifies the bias of the portfolio (Short, Long/Short, Long)'''
        super().__init__()
        self.portfolioBias = portfolioBias 
        self.rebalanceTime = datetime.min
        
     

        # If the argument is an instance of Resolution or Timedelta
        # Redefine rebalancingFunc
        rebalancingFunc = rebalance
        if isinstance(rebalance, int):
            rebalance = Extensions.ToTimeSpan(rebalance)
        if isinstance(rebalance, timedelta):
            rebalancingFunc = lambda dt: dt + rebalance
        if rebalancingFunc:
            self.SetRebalancingFunc(rebalancingFunc)
            ag.Debug(f'pcm init rebalancingFunc')
        self.ag = ag

    def DetermineTargetPercent(self, activeInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        
        # self.ag.Debug(f'{self.ag.Time} pcm-1 triggered at {self.ag.Time}')
        # if self.ag.Time < self.rebalanceTime:
        #     return {}
        self.ag.Debug(f'{self.ag.Time} pcm triggered at {self.ag.Time} utc{self.ag.UtcTime}')
        # Set the rebalance time to match the insight expiry
        self.rebalanceTime = Expiry.EndOfQuarter(self.ag.Time)

        result = {}

        # give equal weighting to each security
        count = sum(x.Direction != InsightDirection.Flat and self.RespectPortfolioBias(x) for x in activeInsights)
        percent = 0 if count == 0 else 1.0 / count
        for insight in activeInsights:
            result[insight] = (insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat) * percent
        return result

    def RespectPortfolioBias(self, insight):
        '''Method that will determine if a given insight respects the portfolio bias
        Args:
            insight: The insight to create a target for
        '''
        return self.portfolioBias == PortfolioBias.LongShort or insight.Direction == self.portfolioBias

class FundamentalFactorAlphaModel(AlphaModel):
    
    def __init__(self):
        self.rebalanceTime = datetime.min
        # Dictionary containing set of securities in each sector
        # e.g. {technology: set(AAPL, TSLA, ...), healthcare: set(XYZ, ABC, ...), ... }
        self.sectors = {}

    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            New insights'''

        if algorithm.Time < self.rebalanceTime:
            return []
        
        # Set the rebalance time to match the insight expiry
        self.rebalanceTime = Expiry.EndOfQuarter(algorithm.Time)
        algorithm.Debug(f'{algorithm.Time} alpha rebalanceTime {self.rebalanceTime}')
        
        insights = []
        
        for sector in self.sectors:
            securities = self.sectors[sector]
            sortedByROE = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=True)
            sortedByPM = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=True)
            sortedByPE = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=False)

            # Dictionary holding a dictionary of scores for each security in the sector
            scores = {}
            for security in securities:
                score = sum([sortedByROE.index(security), sortedByPM.index(security), sortedByPE.index(security)])
                scores[security] = score
                
            # Add best 20% of each sector to longs set (minimum 1)
            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
                # Use Expiry.EndOfQuarter in this case to match Universe, Alpha and PCM
                insights.append(Insight.Price(symbol, Expiry.EndOfQuarter, InsightDirection.Up))
        
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        
        # Remove security from sector set
        for security in changes.RemovedSecurities:
            for sector in self.sectors:
                if security in self.sectors[sector]:
                    self.sectors[sector].remove(security)
        
        # Add security to corresponding sector set
        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)