from AlgorithmImports import *
# from Portfolio.MeanReversionPortfolioConstructionModel import *
from scipy.stats import norm, zscore

"""
AE Changes
1. Trading securities with significant overlap/correlation. E.g. 
 QQQ, TQQQ, SQQQ
 Added HasFundamentalData filter in CoarseSelectionFunction to only trade stocks.
2. changed number_of_symbols to be an input parameter
3. added z_model_period as input parameter

Best performance via 30 symbol  s and z model period = 50
"""

class MeanReversionPortfolioAlgorithm(QCAlgorithm):
    # number_of_symbols = 10

    def Initialize(self):
        self.SetStartDate(2022, 12, 1)  # Set Start Date to Dec 2022
        self.SetEndDate(2024, 6, 18)  # Set End Date to June 18th 2024
        self.SetCash(100000)

        self.Settings.RebalancePortfolioOnInsightChanges = True
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        
        self.SetSecurityInitializer(
            lambda security: security.SetMarketPrice(
                self.GetLastKnownPrice(security)
            )
        )
        self.resolution = Resolution.Hour
        self.UniverseSettings.Resolution = self.resolution
        
        # Requesting data
        self.AddUniverse(self.CoarseSelectionFunction)
        # Feed in data for 100 trading hours before the start date
        # self.SetWarmUp(100, Resolution.Hour)

        # Get the input parameters
        self.number_of_symbols = 50 #int(self.GetParameter("number_of_symbols"))
        self.z_model_period = 30 #int(self.GetParameter("z_model_period"))

        self.AddAlpha(
            ZScoreAlphaModel(
                lookupPeriod = self.z_model_period, 
                resolution = Resolution.Daily
            )
        )
        '''
        self.SetPortfolioConstruction(
            InsightWeightingPortfolioConstructionModel(Resolution.Hour)
        )
        '''
# Test only rebalancing every day
        # self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(Resolution.Hour))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(Resolution.Daily))
        self.Settings.MinAbsolutePortfolioTargetPercentage = \
            0.00000000000000000000001
        self.AddRiskManagement(MaximumDrawdownPercentPerSecurity())
        #self.AddRiskManagement(TrailingStopRiskManagementModel())
        self.SetExecution(ImmediateExecutionModel())

    def CoarseSelectionFunction(self, coarse):
        hasdatasymbols = list(filter(lambda x: x.HasFundamentalData, coarse))
        sortedByDollarVolume = sorted(
            hasdatasymbols, key=lambda x: x.DollarVolume, reverse=True
        )
        highest_volume_symbols = \
            [x.Symbol for x in sortedByDollarVolume[:self.number_of_symbols]]
        return highest_volume_symbols

    def OnSecuritiesChanged(self, changes):
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Invested:
                self.Liquidate(security.Symbol)

class ZScoreAlphaModel(AlphaModel):
    '''Alpha model that uses an Z score to create insights'''
    def __init__(self,
                 lookupPeriod = 20,
                 resolution = Resolution.Daily, 
                 max_weight = 0.2):
        '''Initializes a new instance of the ZScoreAlphaModel class
        Args:
            lookupPeriod: Look up period of history
            resolution: Resoultion of the history'''
        self.lookupPeriod = lookupPeriod
        self.resolution = resolution
        self.predictionInterval = \
            Time.Multiply(Extensions.ToTimeSpan(resolution), lookupPeriod)
        self.symbolDataBySymbol = []

    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:
            The new insights generated'''
        insights = []
        df = algorithm.History(
            self.symbolDataBySymbol, self.lookupPeriod, self.resolution
        )
        if df.empty: 
            return insights

        # Make all of them into a single time index.
        df = df.close.unstack(level=0)
        # Mean of the stocks
        df_mean = df.mean()
        # standard deviation
        df_std = df.std()
        # get last prices
        df = df.iloc[-1]
        # calculate z_score
        z_score = (df.subtract(df_mean)).divide(df_std)
        magnitude = -z_score*df_std/df
        confidence = (-z_score).apply(norm.cdf)

        algorithm.Log(f'{algorithm.Time}: Average Z Score: {z_score.mean()}')
        algorithm.Log(f'{algorithm.Time}: Max Z Score: {z_score.max()}')
        algorithm.Log(f'{algorithm.Time}: Min Z Score: {z_score.min()}')

        # weight = confidence - 1 / (magnitude + 1)
        weights = {}
        for symbol in z_score.index:
            if z_score[symbol] > 3:
                insights.append(
                    Insight.Price(
                        symbol=symbol,
                        period=timedelta(hours=1), 
                        direction=InsightDirection.Down, 
                        magnitude=magnitude[symbol], 
                        confidence=confidence[symbol], 
                        sourceModel=None, 
                        weight=z_score[symbol] - 3
                    )
                )
                algorithm.Log(f'Down Insight for {symbol}, Zscore: {z_score[symbol]}, Magnitude: {magnitude[symbol]}, Confidence: {confidence[symbol]},Weight: {z_score[symbol] - 3}')

            elif z_score[symbol] < -3:
                insights.append(
                    Insight.Price(
                        symbol=symbol, 
                        period=timedelta(hours=1), 
                        direction=InsightDirection.Up, 
                        magnitude=magnitude[symbol], 
                        confidence=confidence[symbol], 
                        sourceModel=None,  
                        weight= abs(z_score[symbol]) - 3
                    )
                )
                algorithm.Log(f'Up Insight for {symbol}, Zscore: {z_score[symbol]}, Magnitude: {magnitude[symbol]}, Confidence: {confidence[symbol]},Weight: {z_score[symbol] - 3}')

        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'''
        for added in changes.AddedSecurities:
            if added.Symbol not in self.symbolDataBySymbol:
                #symbolData = SymbolData(added, self.fastPeriod, self.slowPeriod, algorithm, self.resolution)
                self.symbolDataBySymbol.append(added.Symbol)
        for removed in changes.RemovedSecurities:
            if removed.Symbol in self.symbolDataBySymbol:
                data = self.symbolDataBySymbol.remove(removed.Symbol)