from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel
import random


class HolidaySpecial(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 9, 25)  # Set Start Date
        self.SetEndDate(2020, 12, 25)  # Set Start Date
        starting_mochi_cash = 1000000
        self.SetCash(starting_mochi_cash)  # Set Strategy Cash

        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddAlpha(MochiAlphaModel(self, starting_mochi_cash))
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(lambda time: None))


class MochiAlphaModel(AlphaModel):

    def __init__(self, algorithm, starting_mochi_cash):
        self.Name = 'MochiPAlphaModel'
        self.mochi_wealth = [starting_mochi_cash]
        self.CreateMochi()

    def Update(self, algorithm, data):
        # Clear counters and insights, return if data slice is empty.
        mochi_insights = []
        if not data.HasData:
            return []

        # Is Mochi feeling like working today?:
        if random.random() > self.mochi_diligence:
            return[]

        things_mochi_likes_now = []
        # What companies does Mochi like?
        for symbol in data.keys():
            if random.random() > self.mochi_grumpyness:
                things_mochi_likes_now.append(symbol)

        maximum_things_mochi_likes = len(things_mochi_likes_now)
        picks_a_few = min(self.mochi_hoardiness, maximum_things_mochi_likes)
        if picks_a_few == 0:
            mochi_is_tired = True
            return []

        for symbol in random.choices(things_mochi_likes_now, k=random.randrange(picks_a_few)):
            is_this_going_up = random.getrandbits(1)
            how_much = random.random()
            mochi_direction = self.mochi_directions[is_this_going_up]
            and_when = timedelta(days=self.mochi_patience)
            how_much_does_mochi_invest = random.random()

            one_mochi_insight = Insight(symbol, and_when,
                                        InsightType.Price,
                                        mochi_direction, how_much, 1,
                                        self.Name, how_much_does_mochi_invest)
            mochi_insights.append(one_mochi_insight)

        self.mochi_wealth.append(algorithm.Portfolio.TotalPortfolioValue)
        is_mochi_winning = self.mochi_wealth[-1] > self.mochi_wealth[0]
        is_mochi_doing_good = self.mochi_wealth[-1] > self.mochi_wealth[-2] and is_mochi_winning
        if not is_mochi_doing_good:
            self.CreateMochi()

        return mochi_insights

    def CreateMochi(self):
        self.mochi_diligence = random.random()
        self.mochi_patience = random.randrange(253)
        self.mochi_grumpyness = random.random()
        self.mochi_directions = [InsightDirection.Down, InsightDirection.Up]
        self.mochi_hoardiness = random.randrange(500)