from AlgorithmImports import *
from datetime import timedelta, datetime
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel


class Third_Attempt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetEndDate(2022, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.AddUniverseSelection(Highperformance())
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddAlpha(BuyPerformance())
    
        self.SetPortfolioConstruction(PortfolioBuilder())
        self.AddRiskManagement(Trailing_SL_TP())
        self.SetExecution(ImmediateExecutionModel())
        

class Highperformance (FundamentalUniverseSelectionModel):

    def __init__(self):
        super().__init__( True, None)
        self.lastMonth = -1
        #self.spy = Symbol.Create('SPY', SecurityType.Equity, Market.USA)

    def SelectCoarse(self, algorithm, coarse):
        #run the algorithm once a month, return Universe.Unchanged in case we are looking at exactly the same month
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged
        self.lastMonth = algorithm.Time.month

        sortedByVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        filteredByFundamentals = [x.Symbol for x in sortedByVolume if x.HasFundamentalData]

        return filteredByFundamentals

    def SelectFine(self, algorithm, fine):

        sorted_high = sorted([x for x in fine if x.MarketCap > 2e9
                            and 0.5 > x.OperationRatios.AVG5YrsROIC.FiveYears > 0.20
                            and 50 > x.ValuationRatios.PERatio > 20
                            and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices
                            and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.Healthcare], 
                            key = lambda x: x.ValuationRatios.PERatio, reverse=True)

        #fundamental_universe = [x.Symbol for x in sorted_high[:5]] + [self.spy]

        return [x.Symbol for x in sorted_high[:5]]

class BuyPerformance(AlphaModel):
    
    def __init__(self):
        self.lastMonth = -1
        self.newAdds = []

    def Update(self, algorithm, data):
        
        if algorithm.Time.month == self.lastMonth:
            return []
        self.lastMonth = algorithm.Time.month

        insights = []
        #printing the time (for troubleshooting purposes)
        algorithm.Debug(str(algorithm.Time))

        # For the new securities, if they are NOT invested yet and have Data (tradeable)
        #the algo will add them to the insights
        # So if the security is already invested, it will still be in the portfolio
        #no matter if it gets removed from the Universe.
        for added in self.newAdds:
            if not algorithm.Securities[added].Invested and algorithm.Securities[added].HasData:
                algorithm.Debug('Positive Insight : ' + str(added))
                insights.append(Insight(added, timedelta(30), InsightType.Price, InsightDirection.Up))

        
        
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        # When assets are added to the universe, they will trigger OnSecuritiesChanged() event.
        #From there, you can initialize any state or history required for the Alpha Model
        algorithm.Debug('\n -----ALPHA MODEL ----: ' + str(algorithm.Time))

        # Securities added into the universe will be added to self.newAdds
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.newAdds:
                algorithm.Debug('added symbol : ' + str(symbol))
                self.newAdds.append(symbol)
        
        # Securities removed from the universe will be removed from self.newAdds
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.newAdds:
                algorithm.Debug('removed symbol symbol : ' + str(symbol))
                self.newAdds.remove(symbol)


class PortfolioBuilder(PortfolioConstructionModel):
    def __init__(self):
        
        self.lastMonth = -1

    def CreateTargets (self, algorithm, insights):
        
        if not algorithm.Time.month == self.lastMonth:
            total_equity = algorithm.Portfolio.TotalPortfolioValue
        
        else:
            return[]


        self.lastMonth = algorithm.Time.month
        
        algorithm.Debug('\n -----PORTFOLIO CONSTRUCTION ----: ' + str(algorithm.Time))
        #Create a list of PortfolioTarget objects from Insights

        uniques = []

        for insight in insights:
            if not algorithm.Securities[insight.Symbol].Invested:
                uniques.append(insight.Symbol)
        
        # Now i m going to count how many securities are new, and how many securities are already invested coz I have to see
        #the length of my new portfolio after the new insights get executed
        #this is considered just because my algo does not sell securities until they hit stop loss
        invested = [x.Key for x in algorithm.Portfolio if x.Value.Invested]

        Stocks_In_Portfolio = len(uniques) + len(invested)

        if Stocks_In_Portfolio != 0:
            allocation = 0.95 * total_equity / Stocks_In_Portfolio
        algorithm.Debug('Total Porfolio Value : $' + str(total_equity))
        algorithm.Debug('Number of Stocks : ' + str(Stocks_In_Portfolio))
        algorithm.Debug('Number of Invested stocks : ' + str(len(invested)))
        algorithm.Debug('Number of uniques : ' + str(len(uniques)))
        algorithm.Debug('Individual Security Allocation : $' + str(allocation))
        # Now i need to calculate the new amount of shares per security that my portfolio needs to have. For this one i need to breakdown if the security has been invested
        # if the security is invested, check the existing security allocation and compare it with the allocation value above
        # if the security has recently been added, then we just need to buy the amoun of shares equivalent to the allocation

        # **************
        target_array = []

        

        for x in algorithm.ActiveSecurities:
            holding = x.Value
            symbol = holding.Symbol

            if holding.Invested:
                algorithm.Debug('Security (Invested) :' + str(symbol))
                algorithm.Debug('Old Amount of Shares :' + str(algorithm.Portfolio[symbol].Quantity))            
                shares = allocation / algorithm.Securities[symbol].Close
                algorithm.Debug('New Shares Target :' + str(shares))
                    
                target = PortfolioTarget(symbol, shares)
                

            elif not holding.Invested:
                if algorithm.Securities[symbol].Close != 0:
                    shares = allocation / algorithm.Securities[symbol].Close
                    target = PortfolioTarget(symbol, shares) 
                    algorithm.Debug('Security (NON Invested) :' + str(symbol))
                    algorithm.Debug('New Shares Target :' + str(shares))  

                else:
                    algorithm.Debug('DIVISION BY ZERO!!')             

            target_array.append(target)
            
        return target_array


class Trailing_SL_TP(RiskManagementModel):

    def __init__(self):

        #These two dictionaries will store the symbol and the respective TP and SL
        self.Take_Profit = dict()
        self.Stop_Loss = dict()

    def ManageRisk(self, algorithm, targets):
        
        #Risk_targets is the list to be returned at the end
        #Green is the % value of the take profit (e.g 1.30 is 30% on the upside)
        #Red is the % value of the Stop Loss (e.g 0.90 is 10% on the downside)
        risk_targets = list()
        Green = 1.15
        Red = 0.95

        for target in targets:
            symbol = target.Symbol

            #We just care if the security is invested coz it is the whole point of the risk management model 
            if algorithm.Securities[symbol].Invested:

                if symbol not in self.Take_Profit:

                    self.Take_Profit[symbol] = Green * algorithm.Portfolio[symbol].AveragePrice
                    self.Stop_Loss[symbol] = Red * algorithm.Portfolio[symbol].AveragePrice

                if (symbol in self.Take_Profit):

                    if algorithm.Securities[symbol].Price >= self.Take_Profit[symbol]:
                        algorithm.Debug('Security ' + str(symbol) + 'just hit the Take Profit at ' + str(algorithm.Portfolio[symbol].Price))
                        self.Take_Profit[symbol] = Green * algorithm.Portfolio[symbol].Price
                        self.Stop_Loss[symbol] = Red * algorithm.Portfolio[symbol].Price
                        algorithm.Debug('New TP : ' + str(self.Take_Profit[symbol]) + ' New SL :' + str(self.Stop_Loss[symbol]))

                    if algorithm.Securities[symbol].Price <= self.Stop_Loss[symbol]:
                        algorithm.Debug('Security Liquidated: ' + str(symbol))
                        risk_targets.append(PortfolioTarget(symbol, 0))

        return risk_targets