from AlgorithmImports import *
from EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel, PortfolioBias

class InsightWeightingPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that generates percent targets based on the
    Insight.Weight. The target percent holdings of each Symbol is given by the Insight.Weight from the last
    active Insight for that symbol.
    For insights of direction InsightDirection.Up, long targets are returned and for insights of direction
    InsightDirection.Down, short targets are returned.
    If the sum of all the last active Insight per symbol is bigger than 1, it will factor down each target
    percent holdings proportionally so the sum is 1.
    It will ignore Insight that have no Insight.Weight value.'''

    def __init__(self, rebalance = Resolution.Daily, portfolioBias = PortfolioBias.LongShort):
        '''Initialize a new instance of InsightWeightingPortfolioConstructionModel
        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__(rebalance, portfolioBias)

    def ShouldCreateTargetForInsight(self, insight):
        '''Method that will determine if the portfolio construction model should create a
        target for this insight
        Args:
            insight: The insight to create a target for'''
        # Ignore insights that don't have Weight value
        return insight.Weight is not None

    def DetermineTargetPercent(self, activeInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        result = {}

        # We will adjust weights proportionally in case the sum is > 1 so it sums to 1.
        weightSums = sum(self.GetValue(insight) for insight in activeInsights if self.RespectPortfolioBias(insight))
        weightFactor = 1.0
        if weightSums > 1:
            weightFactor = 1 / weightSums
        for insight in activeInsights:
            result[insight] = (insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat) * self.GetValue(insight) * weightFactor
        return result

    def GetValue(self, insight):
        '''Method that will determine which member will be used to compute the weights and gets its value
        Args:
            insight: The insight to create a target for
        Returns:
            The value of the selected insight member'''
        return abs(insight.Weight)

'''
Bearish Stocks per Option Alpha list - Framework algo
Short stocks with high OI [TSLA, NVDA, AAPL, MSFT, AMD]
Trend Filter: 5 SMA?
Short entry: P X down BB1(40, 1)  Pick one w/lowest SD
Exit: Portfolio DD
Rebalance frequency: Daily

10/10 working but insufficient margin errors
10/11 implement EP Chan BB(20) MR strategy using Zscore. 
    Working but flat insights not liquidating until next month
10/12 MR when BB within 1 SD
'''

from QuantConnect import *
from AlgorithmImports import *

class MyAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2022, 9, 1)  # Set start date for backtest
        self.SetEndDate(2023, 10, 1)
        self.SetCash(100000)  # Set initial cash balance
        #self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.01
    
        # Define a manual universe of high OI stocks per OA
        tickers = ['TSLA',  
                   'NVDA',  
                   'AAPL',  
                   'MSFT', 
                   'AMD'
                   ]
                                      
        symbols = [ Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in tickers ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.UniverseSettings.Resolution = Resolution.Daily  # Set data resolution for the universe
        self.Settings.RebalancePortfolioOnSecurityChanges = False #Stop PCM from rebalancing when universe changes
        self.Settings.RebalancePortfolioOnInsightChanges = True

        # Use the custom alpha model for selecting stocks and generating insights
        self.SetAlpha(BollingerAlphaModel())

        # Set the portfolio construction and execution models
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None))
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: Expiry.EndOfMonth(time), portfolioBias = PortfolioBias.Long))
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(lambda time: Expiry.EndOfMonth(time), portfolioBias = PortfolioBias.Short))
        #self.SetPortfolioConstruction( InsightWeightingPortfolioConstructionModel(self.RebalanceFunction) )
        #self.SetPortfolioConstruction(RiskParityPortfolioConstructionModel(lambda time: Expiry.EndOfMonth(time), portfolioBias = PortfolioBias.Long))

        self.SetExecution(ImmediateExecutionModel())
        #self.SetRiskManagement(NullRiskManagementModel())
        self.SetRiskManagement(MaximumDrawdownPercentPortfolio(0.01))

        self.SetWarmup(41)
        self.month = -1
        
    def RebalanceFunction(self, time):
        if self.month == -1:
            self.month = time.month
            return time
        if self.month != time.month:
            self.month = time.month
            return time
        return None


class BollingerAlphaModel(AlphaModel):
    def __init__(self):

        self.bb1 = {}

        self.month = None
     

    def Update(self, algorithm, data):
        insights = []
        scores = {}
        entryZscore = 1
        exitZscore = 0
        selected_symbols = []  # Initialize selected_symbols to an empty list 
        
        if algorithm.IsWarmingUp:
            return []

        ## If it has already run Update this month, then return nothing
        if algorithm.Time.month == self.month:
            return []
        algorithm.Log('Update() called: ' + str(algorithm.Time))
        ## Update self.month so that it won't do anything in Update until a month has gone by
        self.month = algorithm.Time.month        

        for symbol in self.bb1.keys():
            if data.ContainsKey(symbol) and self.bb1[symbol].StandardDeviation.Current.Value != 0:
                #score = self.bb1[symbol].PercentB.Current.Value
                score = (self.bb1[symbol].Price.Current.Value - self.bb1[symbol].MiddleBand.Current.Value)/self.bb1[symbol].StandardDeviation.Current.Value
                scores[symbol] = score
                #algorithm.Debug(f"%B score for {symbol}: {score}")
                algorithm.Debug(f"Z score for {symbol}: {score}")
        
        # Select symbols with negative scores
        #negative_scores = {k: v for k, v in scores.items() if v < 0}
        
        # Select symbols for short entries (zscore > entryZscore)
        shorts_scores = {k: v for k, v in scores.items() if v > entryZscore}

        # Select top symbol based on score
        selected = max(shorts_scores.items(), key=lambda x: x[1], default=None)
        if selected is not None:
            selected_symbols = [selected[0]]
            algorithm.Debug(f"Selected symbol: {selected_symbols[0]}")
        else:
            algorithm.Debug("No symbols with zscores > entryZscore")

        for symbol in selected_symbols:
            insights.append(Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Down, None, None, None, 1.0))
               
        '''#Emit Flat insights to liquidate shorts
        exit_scores = {k: v for k, v in scores.items() if v <= exitZscore}
        if exit_scores is not None:
            exit_symbols = list(exit_scores.keys())        
    
        # Find symbols that are both in exit_symbols and in algorithm.Insights
        existing_symbols = [insight.Symbol for insight in algorithm.Insights if insight.Direction == InsightDirection.Down]
        common_symbols = set(exit_symbols).intersection(existing_symbols)
        
        for symbol in common_symbols:
            insights.append(Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Flat))'''

        return insights

        

    def OnSecuritiesChanged(self, algorithm, changes):
        
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            self.bb1[symbol] = algorithm.BB(symbol, 20, 1.0, MovingAverageType.Simple, Resolution.Daily)
            
        for removed in changes.RemovedSecurities:
            symbol = removed.Symbol
            if symbol in self.mom1:
                self.bb1.pop(symbol)