#region imports
from AlgorithmImports import *
import numpy as np
from collections import deque
import statsmodels.api as sm
# from scipy import stats
import statistics as stat

import pickle
#endregion

class Q2PlaygroundAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2019, 3, 1)   # Set Start Date
        self.SetEndDate(2024, 6, 1)     # Set End Date
        self.SetCash(1000000)             # Set Strategy Cash
        self.SetSecurityInitializer(BrokerageModelSecurityInitializer(
            self.BrokerageModel, FuncSecuritySeeder(self.GetLastKnownPrices)
        ))

        #################################################################
        self.universe_settings.resolution = Resolution.DAILY

        self._momp = {}          # Dict of Momentum indicator keyed by Symbol
        self._lookback = 252     # Momentum indicator lookback period
        self._num_coarse = 100   # Number of symbols selected at Coarse Selection
        self._num_fine = 50      # Number of symbols selected at Fine Selection
        self._num_long = 5       # Number of symbols with open positions

        self._month = -1
        self._rebalance = False

        self.add_universe(self._coarse_selection_function, self._fine_selection_function)

    def _coarse_selection_function(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
        if self._month == self.time.month:
            return Universe.UNCHANGED

        self._rebalance = True
        self._month = self.time.month

        selected = sorted([x for x in coarse if x.has_fundamental_data and x.price > 5],
            key=lambda x: x.dollar_volume, reverse=True)

        return [x.symbol for x in selected[:self._num_coarse]]

    def _fine_selection_function(self, fine):
        '''Select security with highest market cap'''
        selected = sorted(fine, key=lambda f: f.market_cap, reverse=True)
        return [x.symbol for x in selected[:self._num_fine]]

    def on_data(self, data):
        # Update the indicator
        for symbol, mom in self._momp.items():
            mom.update(self.time, self.securities[symbol].close)

        if not self._rebalance:
            return

        # Selects the securities with highest momentum
        sorted_mom = sorted([k for k,v in self._momp.items() if v.is_ready],
            key=lambda x: self._momp[x].current.value, reverse=True)
        selected = sorted_mom[:self._num_long]

        # Liquidate securities that are not in the list
        for symbol, mom in self._momp.items():
            if symbol not in selected:
                self.liquidate(symbol, 'Not selected')

        # Buy selected securities
        for symbol in selected:
            self.set_holdings(symbol, 1/self._num_long)

        self._rebalance = False

    def on_securities_changed(self, changes):
        # Clean up data for removed securities and Liquidate
        for security in changes.removed_securities:
            symbol = security.symbol
            if self._momp.pop(symbol, None) is not None:
                self.liquidate(symbol, 'Removed from universe')

        for security in changes.added_securities:
            if security.symbol not in self._momp:
                self._momp[security.symbol] = MomentumPercent(self._lookback)

        # Warm up the indicator with history price if it is not ready
        added_symbols = [k for k,v in self._momp.items() if not v.is_ready]

        history = self.history(added_symbols, 1 + self._lookback, Resolution.DAILY)
        history = history.close.unstack(level=0)

        for symbol in added_symbols:
            ticker = symbol.id.to_string()
            if ticker in history:
                for time, value in history[ticker].dropna().items():
                    item = IndicatorDataPoint(symbol, time.date(), value)
                    self._momp[symbol].update(item)