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
        #creating the SPY symbol (for the index)
        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 or not algorithm.TradingCalendar.GetTradingDay().BusinessDay:  
            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 + [self.spy]

    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]] + [self.spy]

class BuyPerformance(AlphaModel):
    
    def __init__(self):
        self.lastMonth = -1
        self.newAdds = []
        self.spy = Symbol.Create('SPY', SecurityType.Equity, Market.USA)

    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))

        # Generate Pos (+) insight if the symbol has been added to 'newAdds', has data to trade and has not been invested before
        for added in self.newAdds:
            if not algorithm.Securities[added].Invested and algorithm.Securities[added].HasData and algorithm.Securities[added].IsTradable:
                algorithm.Debug('Positive Insight : ' + str(added))
                if added not in data.Bars:
                    continue 
                insights.append(Insight(added, timedelta(30), InsightType.Price, InsightDirection.Up))

        # Generate Flat insight if the symbol is already in the Portfolio, has been invested and it's not in newadds
        for x in algorithm.Portfolio:
            holding = x.Value
            symbol = holding.Symbol
            if holding.Invested and (symbol not in self.newAdds):
                if symbol not in data.Bars:
                    continue
                insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Flat))

        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) and (symbol != self.spy):
                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 and (symbol != self.spy):
                algorithm.Debug('removed symbol symbol : ' + str(symbol))
                self.newAdds.remove(symbol)

        #Until now we keep excluding the index SPY. We will execute the trade in the portfolio builder

class PortfolioBuilder(PortfolioConstructionModel):
    def __init__(self):
        
        self.lastMonth = -1
        self.spy = Symbol.Create('SPY', SecurityType.Equity, Market.USA)

    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 (just for symbols, excluding SPY)

        target_array = []
        target = None

        new = []
        #I am creating a vector of new symbols added to the algorithm
        # i will use this vector 'new' to find the lentgh and allocate the right % of the portfolio
        for insight in insights:
            if not algorithm.Securities[insight.Symbol].Invested:
                new.append(insight.Symbol)

        for insight in insights:
           
            if insight.Direction == InsightDirection.Flat:
                target = PortfolioTarget(insight.Symbol, 0) #Generate the sell order on the flat insight

            if insight.Direction == InsightDirection.Up:
                #Calculate shares using the half Portfolio value divided in equal amount with the number of symbols recently added.
                target = PortfolioTarget.Percent(algorithm, insight.Symbol, 1/(2*len(new)))
            if target is not None:
                target_array.append(target)

        #Now i will create the target for the index 'SPY' who's shorting the other half of the portfolio value
        for x in algorithm.ActiveSecurities:
            holding = x.Value
            symbol = holding.Symbol

            if not holding.Invested and symbol == self.spy:
                target = PortfolioTarget.Percent(algorithm, symbol, -0.5)
                target_array.append(target)
        
        return target_array

'''
class Trailing_SL_TP(RiskManagementModel):

    def __init__(self):
        pass

    def ManageRisk(self, algorithm, targets):
        pass
'''