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

class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, security):
        self.Security = security
        self.Symbol = security.Symbol
        self.Fast = None
        self.Slow = None

        # True if the fast is above the slow, otherwise false.
        # This is used to prevent emitting the same signal repeatedly
        self.FastIsOverSlow = False

    @property
    def SlowIsOverFast(self):
        return not self.FastIsOverSlow


class PairsTradingAlphaModel(AlphaModel):

    def __init__(self,universeObject,leverage=1):
        self.pair = [ ]
        self.spreadMean = SimpleMovingAverage(25)
        self.spreadStd = StandardDeviation(25)
        self.period = timedelta(hours=2)
        self.leverage = leverage
        
        
        self.lastDatetime = datetime(1800,1,1)
        self.rebalanceDelta = relativedelta(days=1)
        
        self.universeObject = universeObject
        
        self.hasOpenPositions = False
        
        self.leverage = 1.0
        
    def Update(self, algorithm, data):
        if algorithm.Time + self.rebalanceDelta < self.lastDatetime:
            return []
        
        outputInsights = []
        
        # 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(minutes=20), InsightType.Price, InsightDirection.Down, 0.0025, 1.00, None, prop*weight*self.leverage))
                elif weight > 0:
                    outputInsights.append(Insight(symb, timedelta(minutes=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 OnSecuritiesChanged(self, algorithm, changes):
        pass

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, 2)  # Set Start Date
        self.SetEndDate(2007, 3, 2)  # 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 = 100 #1000
        self.numberOfSymbolsFine = 500
        self.dollarVolumeBySymbol = {}
        
        self.historyPointsLookback = 30
        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