from dateutil.relativedelta import relativedelta
from datetime import datetime
import pandas as pd


class PairsTradingAlphaModel(AlphaModel):

    def __init__(self,universeObject,leverage=1):
        self.pair = [ ]
        self.spreadMean = SimpleMovingAverage(25)
        self.spreadStd = StandardDeviation(25)
        self.period = timedelta(hours=1)
        self.leverage = leverage
        
        
        self.lastDatetime = datetime(1800,1,1)
        self.rebalanceDelta = relativedelta(days=1)
        
        self.universeObject = universeObject
        
        self.hasOpenPositions = False
        
        self.positions = {}
        #self.Portfolio = {}
        
        
        self.expiry = {}
        self.trail_cut_loss = {}
        
        self.leverage = 1.0
        
    def Update(self, algorithm, data):
        outputInsights = []
        
        if algorithm.Time + self.rebalanceDelta < self.lastDatetime:
        
            # Check for losses
            
            
            # Check whether want to liquidate position at the end of the day
            if algorithm.Time.hour == 15:
                for symb in algorithm.Portfolio.keys():
                    q = -1 * int(algorithm.Portfolio[symb].Quantity)
                    
                    if abs(q) > 0:
                        outputInsights.append(Insight(symb, timedelta(hours=5), InsightType.Price,InsightDirection.Flat, None, None, None, 0.00))
                    
                if outputInsights == []:
                    return []
                else:
                    return Insight.Group(outputInsights)
        
        
        
        
        # If no new positions to open
        if self.positions == {}:
            return outputInsights
        
        local_position = dict(self.positions)
        for c,symb in enumerate(self.positions.keys()):
            weight = self.positions[symb]["weight"]
            if weight < 0:
                outputInsights.append(Insight(symb, timedelta(hours=15), InsightType.Price,InsightDirection.Down, None, None, None, abs(weight)*self.leverage))
                pass
            elif weight > 0:
                outputInsights.append(Insight(symb, timedelta(hours=15), InsightType.Price, InsightDirection.Up, None, None, None, abs(weight)*self.leverage))
                pass
                
            try:
                del local_position[symb]
            except KeyError:
                pass
        
        self.positions = local_position
        # Open positions in the morning
        #if not self.hasOpenPositions and algorithm.Time.hour == 10:
        #    for pos,symb in enumerate(self.universeObject.unweightedComponentPortfolio["symbols"]):
        #        weight = self.universeObject.unweightedComponentPortfolio[0]["rescaledEtfWeights"][pos]
        #        prop = self.universeObject.unweightedComponentPortfolio[0]["explainedVarianceProportion"] 
        #        if weight < 0:
        #            outputInsights.append(Insight(symb, timedelta(hours=1), InsightType.Price, InsightDirection.Down, 0.0025, 1.00, None, prop*weight*self.leverage))
        #        elif weight > 0:
        #            outputInsights.append(Insight(symb, timedelta(hours=20), InsightType.Price, InsightDirection.Up, 0.0025, 1.00, None, prop*weight*self.leverage))
        #    self.hasOpenPositions = True
        
        # Close positions at the end of the day
        #if algorithm.Time.hour == 15 and self.hasOpenPositions:
        #    algorithm.Liquidate()
        #    self.hasOpenPositions = False
        
        #algorithm.Plot("Explained Variance", "Exp.Var", 100.0*self.universeObject.unweightedComponentPortfolio[0]["lastReturn"])
        #self.lastDatetime = algorithm.Time
        
        if outputInsights == []:
            return []
        else:
            return Insight.Group(outputInsights)
    
    def CancelAllOrders(self, algorithm, symbol):
        
        #if 'debug' in Roboto.MSGS: self.Debug("{} Cancelling all orders for {}".format(self.Time, symbol))
        openOrders = algorithm.Transactions.CancelOpenOrders(symbol)
        for oo in openOrders:
            if not (oo.Status == OrderStatus.CancelPending):
                r = oo.Cancel()
                #if not r.IsSuccess:
                #    if 'error' in Roboto.MSGS: self.Error("{} Failed to cancel open order {} of {} for reason: {}, {}".format(self.Time, oo.Quantity, oo.Symbol, r.ErrorMessage, r.ErrorCode))

    
    def OnSecuritiesChanged(self, algorithm, changes):
        
        #self.positions
        
        self.positions = {}
        
        for security in changes.AddedSecurities:
            #self.CancelAllOrders(algorithm,security.Symbol)
            if not security.Invested and (not security.Symbol in self.positions.keys()):
                pass
                self.positions[str(security.Symbol)] =  {"weight":
                    #weight = self.universeObject.beta_loading[0][c]
                    self.universeObject.beta_loading[0][list(self.universeObject.unweightedComponentPortfolio["symbols"]).index(security.Symbol)]
                    #self.universeObject.unweightedComponentPortfolio[0]["rescaledEtfWeights"][list(self.universeObject.unweightedComponentPortfolio["symbols"]).index(security.Symbol)]
                }
        
        pass
    
    def OnOrderEvent(self, algorithm, orderEvent):
        
        if orderEvent.Status == OrderStatus.Filled:
            
            order = self.Transactions.GetOrderById(orderEvent.OrderId)
            
            # if completely liquidating position, stop tracking position age
            if not algorithm.Portfolio[order.Symbol].Invested:
                try:
                    del self.expiry[order.Symbol]
                    del self.trail_cut_loss[order.Symbol]
                    #if 'debug' in Roboto.MSGS: self.Debug("{} No longer tracking {}".format(self.Time, order.Symbol))
                except Error:
                    pass
                #    if 'error' in Roboto.MSGS: self.Error("{} Key deletion failed for {}".format(self.Time, order.Symbol))
            
            # if position is completely new, start tracking position age
            else:
                if (order.Symbol not in self.expiry):
                    self.expiry[order.Symbol] = self.Time
                else:
                    pass
                    #if 'error' in Roboto.MSGS: self.Error("{} Key already existed for {}".format(self.Time, order.Symbol))
                
                if (order.Symbol not in self.trail_cut_loss):
                    self.trail_cut_loss[order.Symbol] = 0
                else:
                    pass
                    #if 'error' in Roboto.MSGS: self.Error("{} Key already existed for {}".format(self.Time, order.Symbol))

from Alphas.PearsonCorrelationPairsTradingAlphaModel import PearsonCorrelationPairsTradingAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
#from Portfolio.MeanVarianceOptimizationPortfolioConstructionModel import MeanVarianceOptimizationPortfolioConstructionModel
from PortfolioConstruction import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from UniverseSelection import QC500UniverseSelectionModel
from PaperAlpha import PairsTradingAlphaModel

class FatOrangeMonkey(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2006, 5, 4)  # Set Start Date
        self.SetEndDate(2008, 6, 11)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        
        self.SetExecution(ImmediateExecutionModel())
        
        universeObject = QC500UniverseSelectionModel()
        
        self.SetUniverseSelection(universeObject)
        
        #self.AddAlpha(PearsonCorrelationPairsTradingAlphaModel(252, Resolution.Daily))
        
        self.AddAlpha(PairsTradingAlphaModel(universeObject))

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        #self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))

from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from itertools import groupby
from math import ceil
from dateutil.relativedelta import relativedelta
from datetime import datetime
import pandas as pd

from sklearn.decomposition import PCA

class QC500UniverseSelectionModel(FundamentalUniverseSelectionModel):
    '''Defines the QC500 universe as a universe selection model for framework algorithm
    For details: https://github.com/QuantConnect/Lean/pull/1663'''

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the QC500UniverseSelectionModel'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 500 #1000
        self.numberOfSymbolsFine = 500
        self.dollarVolumeBySymbol = {}
        
        self.historyPointsLookback = 60
        self.historyResolution = Resolution.Daily

        self.lastDatetime = datetime(1800,1,1)
        self.rebalanceDelta = relativedelta(days=1)
        
        self.targetExplainedVariance = 0.8  # Proportion of the variance explained by PCA components
        
        self.stock_std = None
        
        self.unweightedComponentPortfolio = {}
        self.beta_loading = {}
        
        # Toggle rescaling per component
        self.rescale_eigenportfolio = False
    def SelectCoarse(self, algorithm, coarse):
        '''Performs coarse selection for the QC500 constituents.
        The stocks must have fundamental data
        The stock must have positive previous-day close price
        The stock must have positive volume on the previous trading day'''
        if algorithm.Time + self.rebalanceDelta < self.lastDatetime:
            return Universe.Unchanged
            
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData  
                                                            and x.Volume > 0 
                                                            and x.Price > 0
                                                            and x.Market == "usa"],
                                     key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}

        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        # return the symbol objects our sorted collection
        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''Performs fine selection for the QC500 constituents
        The company's headquarter must in the U.S.
        The stock must be traded on either the NYSE or NASDAQ
        At least half a year since its initial public offering
        The stock's market cap must be greater than 500 million'''
        
        
        # Might want to add Volume, DollarVolume filters
        sortedBySector = sorted([x for x in fine if (algorithm.Time - x.SecurityReference.IPODate).days > 180
                                        and x.MarketCap > 5e8],
                               key = lambda x: x.CompanyReference.IndustryTemplateCode)

        count = len(sortedBySector)
        
        # If no security has met the QC500 criteria, the universe is unchanged.
        # A new selection will be attempted on the next trading day as self.lastMonth is not updated
        if count == 0:
            return Universe.Unchanged

        # Update self.lastMonth after all QC500 criteria checks passed
        self.lastDatetime = algorithm.Time

        #percent = self.numberOfSymbolsFine / count
        #sortedByDollarVolume = []
        

        # select stocks with top dollar volume in every single sector
        #for code, g in groupby(sortedBySector, lambda x: x.CompanyReference.IndustryTemplateCode):
        #    y = sorted(g, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse = True)
        #    c = ceil(len(y) * percent)
        #    sortedByDollarVolume.extend(y[:c])

        #sortedByDollarVolume = sorted(sortedByDollarVolume, key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)
        
        #selectedSymbols =  [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]
        
        
        # Calculate Returns matrix
        for c,s in enumerate(sortedBySector):
            history = algorithm.History(s.Symbol, self.historyPointsLookback+1, self.historyResolution)

            pct_col = history[["close"]].pct_change().reset_index(drop=True).rename(columns={"close":s.Symbol})
            
            if c == 0:
                comb_returns = pct_col
            else:
                comb_returns.insert(0,s.Symbol,pct_col[s.Symbol])
                
        # Drop top NA row, followed by any other columns with mising da
        comb_returns = comb_returns[1:].dropna(axis=1)
        symbols = comb_returns.keys()
        
        self.stock_std = comb_returns.std()
        standard_return = (comb_returns-comb_returns.mean())/self.stock_std
        
        # Calculate corralation matrix
        corr_mat = standard_return.corr()
        
        pca = PCA()
        pca_obj = pca.fit(corr_mat)
        
        
        # Find number of components required for the target explained variance
        sum_var = 0
        for num_vars,var in enumerate(pca_obj.explained_variance_ratio_):
            sum_var += var
            if sum_var > self.targetExplainedVariance:
                num_vars += 1
                break
        
        
        #componentReturns = pd.Dataframe()
        
        # Create a dictionary of weights per component
        self.unweightedComponentPortfolio["symbols"] =  symbols
        for i in range(num_vars):
            # "Normalise" components i.e., "buy" ETF at 1x leverage
            tmp_comp = pca_obj.components_[i]/self.stock_std
            
            
            if self.rescale_eigenportfolio:
                tmp_comp = tmp_comp/(abs(tmp_comp).sum())
            
            componentReturns_tmp = (comb_returns*tmp_comp).sum(axis=1)
            
            # Calculate historic returns per Component
            if i == 0:
                componentReturns = pd.DataFrame(componentReturns_tmp)
            else:
                componentReturns.insert(0,i,componentReturns_tmp)
                
            self.unweightedComponentPortfolio[i] = {
               "component":pca_obj.components_[i],
               "etfWeights":list(tmp_comp),
               "rescaledEtfWeights":list(tmp_comp/(abs(tmp_comp).sum())),
               "lastReturn":componentReturns_tmp.iloc[-1],
               "explainedVarianceProportion":pca_obj.explained_variance_ratio_[i]
           }
           
        
        # Calculate Beta loadings per component. Equation (11) page 16
        for i in componentReturns:
            tmp_ret = componentReturns[i]
            self.beta_loading[i] = comb_returns.apply(lambda x: tmp_ret.cov(x))/tmp_ret.var()
        
        return symbols

class EqualWeightingPortfolioConstructionModel(PortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that gives equal weighting to all securities.
    The target percent holdings of each security is 1/N where N is the number of securities.
    For insights of direction InsightDirection.Up, long targets are returned and
    for insights of direction InsightDirection.Down, short targets are returned.'''

    def __init__(self, rebalance = Resolution.Daily, portfolioBias = PortfolioBias.LongShort):
        '''Initialize a new instance of EqualWeightingPortfolioConstructionModel
        Args:
            rebalance: Rebalancing parameter. If it is a timedelta, date rules or Resolution, it will be converted into a function.
                              If None will be ignored.
                              The function returns the next expected rebalance time for a given algorithm UTC DateTime.
                              The function returns null if unknown, in which case the function will be called again in the
                              next loop. Returning current time will trigger rebalance.
            portfolioBias: Specifies the bias of the portfolio (Short, Long/Short, Long)'''
        self.portfolioBias = portfolioBias

        # If the argument is an instance of Resolution or Timedelta
        # Redefine rebalancingFunc
        rebalancingFunc = rebalance
        if isinstance(rebalance, int):
            rebalance = Extensions.ToTimeSpan(rebalance)
        if isinstance(rebalance, timedelta):
            rebalancingFunc = lambda dt: dt + rebalance
        if rebalancingFunc:
            self.SetRebalancingFunc(rebalancingFunc)

    def DetermineTargetPercent(self, activeInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        result = {}

        # give equal weighting to each security
        count = sum(x.Direction != InsightDirection.Flat and self.RespectPortfolioBias(x) for x in activeInsights)
        percent = 0 if count == 0 else 1.0 / count
        for insight in activeInsights:
            result[insight] = (insight.Direction if self.RespectPortfolioBias(insight) else InsightDirection.Flat) * insight.Weight
        return result

    def RespectPortfolioBias(self, insight):
        '''Method that will determine if a given insight respects the portfolio bias
        Args:
            insight: The insight to create a target for
        '''
        return self.portfolioBias == PortfolioBias.LongShort or insight.Direction == self.portfolioBias