#region imports
from AlgorithmImports import *
#endregion

class FundamentalFactorAlphaModel(AlphaModel):
   
    def __init__(self, str_period=7, str_multiplier=3):
        self.rebalanceTime = datetime.min
        # Dictionary containing set of securities in each sector
        # e.g. {technology: set(AAPL, TSLA, ...), healthcare: set(XYZ, ABC, ...), ... }
        self.sectors = {}
        self.symbolData = {}
        self.exclusions = ['SPY']
        self.str_period = str_period
        self.str_multiplier = str_multiplier

    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'''
        
        algorithm.Plot("myPlot", "SPY", self.spy.Close)

        if algorithm.Time <= self.rebalanceTime:
            return []
        
        # Set the rebalance time to match the insight expiry
        #self.rebalanceTime = Expiry.EndOfQuarter(algorithm.Time)
        self.rebalanceTime = Expiry.EndOfMonth(algorithm.Time)

        insights = []
        
        for sector in self.sectors:
            securities = self.sectors[sector]
            sortedByROE_pos = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=True)
            sortedByPM_pos = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=True)
            sortedByPE_pos = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=True)
            sortedByROE_neg = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.ROE.Value, reverse=False)
            sortedByPM_neg = sorted(securities, key=lambda x: x.Fundamentals.OperationRatios.NetMargin.Value, reverse=False)
            sortedByPE_neg = sorted(securities, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=False)

            # Dictionary holding a dictionary of scores for each security in the sector
            scores_pos = {}
            for security in securities:
                score = sum([sortedByROE_pos.index(security), sortedByPM_pos.index(security), sortedByPE_pos.index(security)])
                scores_pos[security] = score

            scores_neg = {}
            for security in securities:
                score = sum([sortedByROE_neg.index(security), sortedByPM_neg.index(security), sortedByPE_neg.index(security)])
                scores_neg[security] = score
            
            for key, sd in self.symbolData.items():
                # Add best 20% of each sector to long set (minimum 1)
                if self.spy.Close > sd.str.Current.Value:
                    #length = max(int(len(scores_pos)/5), 1)
                    length = 1
                    for security in sorted(scores_pos.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))
                        insights.append(Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Up))


                # Add worst 20% of each sector to short set (minimum 1)
                elif self.spy.Close < sd.str.Current.Value:
                    #length = max(int(len(scores_neg)/5), 1)
                    length = 1
                    for security in sorted(scores_neg.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.Down))
                        insights.append(Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Down))

        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:
            if security.Symbol.Value not in self.exclusions:                
                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:
            if security.Symbol.Value in self.exclusions: # Add SPY to symbolData
                self.symbolData[security.Symbol] = SymbolData(algorithm, security, self.str_period, self.str_multiplier)
                #self.indicators[security.Symbol] = SymbolData(algorithm, security, self.str_period, self.str_multiplier)
                self.spy = security
            if security.Symbol.Value not in self.exclusions:                
                sector = security.Fundamentals.AssetClassification.MorningstarSectorCode
                if sector not in self.sectors:
                    self.sectors[sector] = set()
                self.sectors[sector].add(security)

class SymbolData:
    def __init__(self, algorithm, security, str_period, str_multiplier):
        self.algorithm = algorithm
        self.Security = security
        self.str = SuperTrend(str_period, str_multiplier, MovingAverageType.Wilders)
        #self.ema = ExponentialMovingAverage(20,0.5)

        #self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, Resolution.Daily)
        #algorithm.RegisterIndicator(security.Symbol, self.str, self.Consolidator)

        algorithm.RegisterIndicator(self.Security.Symbol, self.str, Resolution.Daily)
        algorithm.WarmUpIndicator(security.Symbol, self.str, Resolution.Daily)

        #algorithm.RegisterIndicator(self.Security.Symbol, self.ema, Resolution.Daily)
        #algorithm.WarmUpIndicator(security.Symbol, self.ema, Resolution.Daily)


        self.str.Updated += self.OnSTRUpdated
        #self.ema.Updated += self.OnEMAUpdated


    def OnSTRUpdated(self, sender, updated):
        self.algorithm.Plot("myPlot", "STR", updated)

    def OnEMAUpdated(self, sender, updated):
        self.algorithm.Plot("myPlot", "EMA", updated)

# region imports
from AlgorithmImports import *
import numpy as np
# endregion

from AlphaModel import *
class VerticalTachyonRegulators(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 5, 30)
        self.SetEndDate(2023, 6, 1)
        self.SetCash(3000)
        self.EnableAutomaticIndicatorWarmUp = True
        # Universe selection
        self.month = 0
        self.num_coarse = 500
        self.UniverseSettings.Resolution = Resolution.Daily
        #
        # Add SPY to the universe 
        #
        self.spy = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.AddUniverseSelection(ManualUniverseSelectionModel(self.spy))
        #
        # Add rest of stocks to universe
        #
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
       
        # Alpha Model
        self.AddAlpha(FundamentalFactorAlphaModel())
        # Portfolio construction model
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(self.IsRebalanceDue))
       
        # Risk model
        #self.SetRiskManagement(NullRiskManagementModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.1))
        #self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.10))
        
        # Execution model
        self.SetExecution(ImmediateExecutionModel())

        chart = Chart("myPlot")
        self.AddChart(chart)        

    # Share the same rebalance function for Universe and PCM for clarity
    def IsRebalanceDue(self, time):
        # Rebalance on the first day of the Quarter or Month
        #if time.month == self.month or time.month not in [1, 4, 7, 10]:
        #if time.month == self.month or time.month not in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]:
        if time.month == self.month or time.month not in np.arange(1,13):
            return None
           
        self.month = time.month
        return time

    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]

        #sectors = [
        #    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
                                    and x.OperationRatios.NetMargin.Value != 0
                                    and x.ValuationRatios.PERatio != 0]

        return filtered_fine