from QuantConnect.Data.Custom.PsychSignal import *

class SP500SectorsETFUniverseExampleAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 4, 10)
        self.SetCash(100000)

        ## Set Universe Selection Model
        self.SetUniverseSelection(TechnologyETFUniverse())
        #self.SetUniverseSelection(SP500SectorsETFUniverse())

        ## Set Alpha Model
        self.SetAlpha(MyAlphaModel())

        ## Set Portfolio Construction Model
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        ## Set Execution Model
        self.SetExecution(ImmediateExecutionModel())

        ## Set Risk Management Model
        self.SetRiskManagement(NullRiskManagementModel())


class MyAlphaModel:

    def __init__(self):
        self.timeEntered = datetime(1998, 1, 1)
        self.period = timedelta(days=10)

    def Update(self, algorithm, data):

        insights = []
        
        # Scan our last time traded to prevent churn.
        if (algorithm.Time - self.timeEntered) <= self.period:
            return insights

        # Fetch the PsychSignal data for the active securities and trade on any
        for security in algorithm.ActiveSecurities.Values:
            if security.Type != SecurityType.Equity:
                continue
            tweets = security.Data.GetAll(PsychSignalSentiment)
            for sentiment in tweets:
                if sentiment.BullIntensity > 2.0 and sentiment.BullScoredMessages > 3:
                    insights.append(Insight.Price(security.Symbol, self.period, InsightDirection.Up, None, None))
                    self.timeEntered = algorithm.Time + self.period

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            if security.Type == SecurityType.Equity:
                algorithm.AddData(PsychSignalSentiment, security.Symbol)