from QuantConnect.Data.UniverseSelection import *
import operator
from math import ceil,floor
from datetime import datetime
import pandas as pd
import numpy as np
from scipy.optimize import minimize


class CoarseFineFundamentalComboAlgorithm(QCAlgorithm):
    

    def Initialize(self):
        self.SetStartDate(2011,2,01)  #Set Start Date
        self.SetEndDate(datetime.now())    #Set End Date
        self.SetCash(100000)            #Set Strategy Cash
        self.rebalence_flag = 0
        self.first_month_trade_flag = 1
        self.trade_flag = 0    

        self.UniverseSettings.Resolution = Resolution.Daily        
        
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.AddEquity("SPY")
        self.number_coarse = 2000
        self.num_portfolios = 6
        self.num_long = 8
        self.num_short = 6
        self.long = None
        self.short = None
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), Action(self.Rebalancing))


    def CoarseSelectionFunction(self, coarse):

        if self.rebalence_flag or self.first_month_trade_flag:

            CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
            sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=True) 
            top = sortedByDollarVolume[:self.number_coarse]

            return [i.Symbol for i in top]
        else:
            return []


    def FineSelectionFunction(self, fine):

        if self.rebalence_flag or self.first_month_trade_flag:
      
            filtered_fine = [x for x in fine if x.ValuationRatios.PriceChange1M 
                                            and x.EarningReports.TotalDividendPerShare.ThreeMonths
                                            and x.ValuationRatios.PERatio
                                            and x.ValuationRatios.BookValueYield]
    
            sortedByfactor1 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PriceChange1M, reverse=False)
            sortedByfactor2 = sorted(filtered_fine, key=lambda x: x.EarningReports.TotalDividendPerShare.ThreeMonths, reverse=True)
            sortedByfactor3 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.PERatio, reverse=False)
            sortedByfactor4 = sorted(filtered_fine, key=lambda x: x.ValuationRatios.BookValueYield, reverse=True)

            
            num_stocks = floor(len(filtered_fine)/self.num_portfolios)
        
            stock_dict = {}
            
            for i,ele in enumerate(sortedByfactor1):
                rank1 = i
                rank2 = sortedByfactor2.index(ele)
                rank3 = sortedByfactor3.index(ele)
                rank4 = sortedByfactor4.index(ele)
  
                score = [ceil(rank1/num_stocks),
                         ceil(rank2/num_stocks),
                         ceil(rank3/num_stocks),
                         ceil(rank4/num_stocks)]
                
                score = sum(score)
                stock_dict[ele] = score

            self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=True)
            sorted_symbol = [self.sorted_stock[i][0] for i in xrange(len(self.sorted_stock))]
            self.long = sorted_symbol[:self.num_long]
            self.short = sorted_symbol[-self.num_short:]
            
            
            self.rebalence_flag = 0
            self.first_month_trade_flag = 0
            self.trade_flag = 1
            
            return [i.Symbol for i in (self.long+self.short)]
        else:
            return []


    def OnData(self, data):
        if self.long is None and self.short is None: return
        if self.trade_flag or self.first_month_trade_flag:        
            holdings = self.long + self.short
            for i in self.Portfolio.Values:
                if (i.Invested) and (i.Symbol not in holdings):
                    self.Liquidate(i.Symbol)

            symbols = [ x.Symbol.Value for x in self.long] 
            for i in symbols:
                self.AddEquity(i, Resolution.Daily)
            history = self.History(symbols, 100, Resolution.Daily)
            if history is None: return
       
            data = {}
            for i in symbols:
                if i in history.index.levels[0]:
                    data[i] = history.loc[i]['close'] 
                
            
            df_price = pd.DataFrame(data,columns=data.keys()) 

            log_return = np.log(df_price / df_price.shift(1)).dropna()
            
            a = PortfolioOptimization(log_return, 0, len(data))
            opt_weight = a.opt_portfolio()            
                        
            for i in range(len(data)):
                self.SetHoldings(df_price.columns[i],opt_weight[i])

            # # Equally weighted 
            # for i in self.long:
            #     self.SetHoldings(i.Symbol,1.0/self.num_long)               

            if self.num_short != 0:        
                for i in self.short:
                    self.SetHoldings(i.Symbol,-0.2/self.num_short)

            self.trade_flag = 0

            
    def Rebalancing(self):
        self.rebalence_flag = 1
        
        
class PortfolioOptimization(object):

    import numpy as np
    import pandas as pd

    def __init__(self, df_return, risk_free_rate, num_assets):
        self.log_return = df_return
        self.risk_free_rate = risk_free_rate
        self.n = num_assets # numbers of risk assets in portfolio

    def annual_port_return(self, weights):
        # calculate the annual return of portfolio
        return np.sum(self.log_return.mean() * weights) * 252

    def annual_port_vol(self, weights):
        # calculate the annual volatility of portfolio
        return np.sqrt(np.dot(weights.T, np.dot(self.log_return.cov() * 252, weights)))

    def mc_mean_var(self):
        # apply monte carlo method to search for the feasible region of return
        returns = []
        vols = []
        for i in range(200):
            weights = np.random.rand(n)
            weights /= np.sum(weights)
            returns.append(self.annual_port_return(weights))
            vols.append(self.annual_port_vol(weights))
        return returns, vols

    def min_func(self, weights):
        return - self.annual_port_return(weights) / self.annual_port_vol(weights)

    def opt_portfolio(self):
        # maximize the sharpe ratio to find the optimal weights
        cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
        bnds = tuple((0, 1) for x in range(self.n))
        opt = minimize(self.min_func,
                       np.array(self.n * [1. / self.n]),
                       method='SLSQP',
                       bounds=bnds,
                       constraints=cons)
                     
        opt_weights = opt['x']
 
        return opt_weights