from System import *
from System.Collections.Generic import List
from QuantConnect import *
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
import pandas as pd
import numpy as np

class CoarseFineFundamentalComboAlgorithm(QCAlgorithm):

    def Initialize(self):

		self.SetStartDate(2010,01,01)  #Set Start Date
		self.SetEndDate(2014,06,01)    #Set End Date
		self.SetCash(50000)            #Set Strategy Cash

		self.UniverseSettings.Resolution = Resolution.Daily
        
		self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
		self.AddEquity("SPY")
		self.numOfCourseSymbols = 110
		self.numOfPortfolio = 5
		self.numOfPortSymbols = self.numOfPortfolio * 15
		self.numOfMonths = 48
		self._changes = SecurityChanges.None
		
		self.flag1 = 1
		self.flag2 = 0
		self.pos = 0
		self.df_return = pd.DataFrame(index = range(self.numOfPortfolio+1))
		self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), Action(self.Rebalancing))


    # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def CoarseSelectionFunction(self, coarse):
		if self.flag1:
			CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
			sortedByVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True)
			top = sortedByVolume[:self.numOfCourseSymbols]

	    	# we need to return only the symbol objects
			list = List[Symbol]()
			for x in top:
				list.Add(x.Symbol)
			return list
		else:
			return(List[Symbol]())

    def FineSelectionFunction(self, fine):
		if self.flag1:
			self.flag1 = 0
			self.flag2 = 1
			# filter the fine by deleting equities wit zero factor value 
			filtered_fine = [x for x in fine if x.ValuationRatios.PERatio != 0 ]
			# sort the fine by reverse order of factor value
			sorted_fine = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PERatio, reverse=True)
			self.symbol = [str(x.Symbol.Value) for x in sorted_fine]
			# factor_value = [x.ValuationRatios.PERatio for x in sorted_fine]
			self.pos = self.pos + 1
 		 	return (List[Symbol]())
		else:
			return (List[Symbol]())
       	
    def Rebalancing(self):
		self.flag1 = 1

    def OnData(self, data):
		if self.pos > 0:
			if self.flag2 == 1:
				self.flag2 = 0
				sorted_symbol = self.symbol
				self.AddEquity("SPY") # add benchmark
				for x in sorted_symbol:
					self.AddEquity(x)
				history = self.History(21,Resolution.Daily)
				monthly_return =[]
				new_symbol_list =[]
				
				for j in range(len(sorted_symbol)):
					try:
						daily_price = []
						for slice in history:
							bar = slice[sorted_symbol[j]]
							daily_price.append(float(bar.Close))
						new_symbol_list.append(sorted_symbol[j])
				 		monthly_return.append(daily_price[-1] / daily_price[0] - 1)
					except:
						self.Log("No history data for " + str(sorted_symbol[j]))
						del daily_price

				monthly_return = monthly_return[:self.numOfPortSymbols]
				# devide the stocks into n portfolios
				reshape_return = np.reshape(monthly_return, (self.numOfPortfolio, len(monthly_return)/self.numOfPortfolio))
				# calculate the average return of different portfolios
				port_avg_return = np.mean(reshape_return,axis=1).tolist()
				# add return of "SPY" as the benchmark  to the end of the return list 
				benchmark_syl = self.AddEquity("SPY").Symbol 
				history_benchmark = self.History(21,Resolution.Daily)
				benchmark_daily_price = [float(slice[benchmark_syl].Close) for slice in history_benchmark]
				benchmark_monthly_return = (benchmark_daily_price[-1]/benchmark_daily_price[0]) - 1
				port_avg_return.append(benchmark_monthly_return)
				self.df_return[str(self.pos)] = port_avg_return	

				if self.pos == self.numOfMonths:
					self.Log(str(self.df_return))
					result = self.calculate_criteria(self.df_return)

					self.Log(str(result))

    def calculate_criteria(self,df_port_return):
    	
		total_return = (df_port_return + 1).T.cumprod().iloc[-1,:] - 1
		self.Log(str(total_return))
		annual_return = (total_return+1)**(1./6)-1

		excess_return = annual_return - np.array(annual_return)[-1]
		correlation = annual_return[0:5].corr(pd.Series([5,4,3,2,1],index = annual_return[0:5].index))
		# higher factor with higher return 
		if np.array(total_return)[0] > np.array(total_return)[-2]:
			loss_excess = df_port_return.iloc[-2,:] - df_port_return.iloc[-1,:]
			win_excess = df_port_return.iloc[0,:] - df_port_return.iloc[-1,:]
			loss_prob = loss_excess[loss_excess<0].count()/float(len(loss_excess))
			win_prob = win_excess[win_excess>0].count()/float(len(win_excess))
			win_port_excess_return = np.array(excess_return)[0]
			loss_port_excess_return = np.array(excess_return)[-2]
		# higher factor with lower return
		else:
			loss_excess = df_port_return.iloc[0,:] - df_port_return.iloc[-1,:]
			win_excess = df_port_return.iloc[-2,:] - df_port_return.iloc[-1,:]
			loss_prob = loss_excess[loss_excess<0].count()/float(len(loss_excess))
			win_prob = win_excess[win_excess>0].count()/float(len(win_excess))
			win_port_excess_return = np.array(excess_return)[-2]
			loss_port_excess_return = np.array(excess_return)[0]
	
		test_result = {"correelation":[],
						"win probality":[],
						"loss probality":[],
						"win portfolio excess return":[],
					   	"loss portfolio excess return":[]}
		test_result["correelation"].append(correlation)
		test_result["win probality"].append(win_prob)
		test_result["loss probality"].append(loss_prob)
		test_result["win portfolio excess return"].append(win_port_excess_return)
		test_result["loss portfolio excess return"].append(loss_port_excess_return)
		

		return pd.DataFrame(test_result)