from System.Collections.Generic import List
from QuantConnect.Data.UniverseSelection import *


class BasicTemplateAlgorithm(QCAlgorithm):
	
	def __init__(self):
		self.reb = 1
		self.num_coarse = 100
		self.num_fine = 20
		self.symbols = None
		self.first_month = 0


	def Initialize(self):
		self.SetCash(100000)
		self.SetStartDate(2007,1,4)
		self.SetEndDate(2017,1,1)
		
		
		self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
		
		self.UniverseSettings.Resolution = Resolution.Daily
		self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
		
		
		self.Schedule.On(self.DateRules.MonthStart(self.spy), 
		self.TimeRules.AfterMarketOpen(self.spy,5), Action(self.rebalance))
		
	
	def CoarseSelectionFunction(self, coarse):
		if self.reb != 1:
			return (List[Symbol]())
		
		selected = [x for x in coarse if (x.HasFundamentalData) 
		            and (float(x.Price) > 5)]
		
		sortedByDollarVolume = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
		top = sortedByDollarVolume[:self.num_coarse]
		list = List[Symbol]()
		for x in top:
			list.Add(x.Symbol)
		return list


	def FineSelectionFunction(self, fine):
		if self.reb != 1:
			return (List[Symbol]())
			
		self.reb = 0
			
		filtered_fine = [x for x in fine if x.OperationRatios.OperationMargin.Value
										and x.ValuationRatios.PriceChange1M 
										and x.ValuationRatios.BookValuePerShare]
										
		self.Log('remained to select %d'%(len(filtered_fine)))
		
		sortedByfactor1 = sorted(filtered_fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=True)
		sortedByfactor2 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PriceChange1M, reverse=True)
		sortedByfactor3 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.BookValuePerShare, reverse=True)
		
		


		stock_dict = {}
		
		for i,ele in enumerate(sortedByfactor1):
			rank1 = i
			rank2 = sortedByfactor2.index(ele)
			rank3 = sortedByfactor3.index(ele)
			score = sum([rank1*0.2,rank2*0.4,rank3*0.4])
			stock_dict[ele] = score
		
		self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
		sorted_symbol = [self.sorted_stock[i][0] for i in xrange(len(self.sorted_stock))]
		topFine = sorted_symbol[:self.num_fine]
		
		self.long = [x.Symbol for x in sorted_symbol[:20]]
		self.short = [x.Symbol for x in sorted_symbol[-20:]]

		list = List[Symbol]()
		for x in topFine:
			list.Add(x.Symbol)
		
		return list

	def OnData(self, data):
		pass
	
	def rebalance(self):
		if self.first_month == 0:
			self.first_month += 1
			return
		self.Liquidate()

		for i in self.long:
			self.SetHoldings(i,1.0/20)
		
		for i in self.short:
			self.SetHoldings(i,-1.0/20)

		self.reb = 1