# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from AlgorithmImports import *

### <summary>
### CapmAlphaRankingFrameworkAlgorithm: example of custom scheduled universe selection model
### Universe Selection inspired by https://www.quantconnect.com/tutorials/strategy-library/capm-alpha-ranking-strategy-on-dow-30-companies
### </summary>
class CapmAlphaRankingFrameworkAlgorithm(QCAlgorithm):
 '''CapmAlphaRankingFrameworkAlgorithm: example of custom scheduled universe selection model'''

 def initialize(self):
 ''' Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

 # Set requested data resolution
 self.universe_settings.resolution = Resolution.MINUTE

 self.set_start_date(2016, 1, 1) #Set Start Date
 self.set_end_date(2017, 1, 1) #Set End Date
 self.set_cash(100000) #Set Strategy Cash

 # set algorithm framework models
 self.set_universe_selection(CapmAlphaRankingUniverseSelectionModel())
 self.set_alpha(ConstantAlphaModel(InsightType.PRICE, InsightDirection.UP, timedelta(1), 0.025, None))
 self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
 self.set_execution(ImmediateExecutionModel())
 self.set_risk_management(MaximumDrawdownPercentPerSecurity(0.01))

class CapmAlphaRankingUniverseSelectionModel(UniverseSelectionModel):
 '''This universe selection model picks stocks with the highest alpha: interception of the linear regression against a benchmark.'''

 period = 21
 benchmark = "SPY"

 # Symbols of Dow 30 companies.
 _symbols = [Symbol.create(x, SecurityType.EQUITY, Market.USA)
 for x in ["AAPL", "AXP", "BA", "CAT", "CSCO", "CVX", "DD", "DIS", "GE", "GS",
 "HD", "IBM", "INTC", "JPM", "KO", "MCD", "MMM", "MRK", "MSFT",
 "NKE","PFE", "PG", "TRV", "UNH", "UTX", "V", "VZ", "WMT", "XOM"]]

 def create_universes(self, algorithm):

 # Adds the benchmark to the user defined universe
 benchmark = algorithm.add_equity(self.benchmark, Resolution.DAILY)

 # Defines a schedule universe that fires after market open when the month starts
 return [ ScheduledUniverse(
 benchmark.exchange.time_zone,
 algorithm.date_rules.month_start(self.benchmark),
 algorithm.time_rules.after_market_open(self.benchmark),
 lambda datetime: self.select_pair(algorithm, datetime),
 algorithm.universe_settings)]

 def select_pair(self, algorithm, date):
 '''Selects the pair (two stocks) with the highest alpha'''
 dictionary = dict()
 benchmark = self._get_returns(algorithm, self.benchmark)
 ones = np.ones(len(benchmark))

 for symbol in self._symbols:
 prices = self._get_returns(algorithm, symbol)
 if prices is None: continue
 A = np.vstack([prices, ones]).T

 # Calculate the Least-Square fitting to the returns of a given symbol and the benchmark
 ols = np.linalg.lstsq(A, benchmark)[0]
 dictionary[symbol] = ols[1]

 # Returns the top 2 highest alphas
 ordered_dictionary = sorted(dictionary.items(), key= lambda x: x[1], reverse=True)
 return [x[0] for x in ordered_dictionary[:2]]

 def _get_returns(self, algorithm, symbol):

 history = algorithm.history([symbol], self.period, Resolution.DAILY)
 if history.empty: return None

 window = RollingWindow(self.period)
 rate_of_change = RateOfChange(1)

 def roc_updated(s, item):
 window.add(item.value)

 rate_of_change.updated += roc_updated

 history = history.close.reset_index(level=0, drop=True).items()

 for time, value in history:
 rate_of_change.update(time, value)

 return [ x for x in window]