import pandas as pd
import numpy as np
from scipy.stats import multivariate_normal as mvn
from scipy.stats import norm

    
class HighDebtUniverse(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2004, 1, 1)
        self.SetEndDate(2004, 3, 1)
        # self.SetStartDate(2021, 1, 21)
        # self.SetEndDate(2021, 1, 18)
        self.SetCash(100000)
        self.SetBenchmark("SPY")
        self.SetWarmUp(1)
        self.UniverseSettings.Resolution = Resolution.Daily
        # self.UniverseSettings.Resolution = Resolution.Monthly
        self.AddUniverse(self.CoarseSelectionFunction, self.SelectFine)
        self.SetBrokerageModel(InteractiveBrokersBrokerageModel())
        self.SetExecution(ImmediateExecutionModel())
        self.EMPTY_SERIES = pd.Series()
        self.AddEquity("SPY",Resolution.Hour).Symbol
        self.stops = {}        # Keep track of stop loss orders so we can update them
        self.stoplevel=0.70
        # self.FirstFilter = 60
        self.FirstFilter = 100
        self.SecondFilter = 5
        self.WeekDay = 3
        self.numberOfSymbolsCoarse = 200
        self.numberOfSymbolsFine = 3000
        
        self.result = {}
        self.liquidate = set()
        
        self.mvn_lookback = 120
        self.mvn_window = 60
        # self.mvn_cdf_window = 20
        self.mvn_lookahead = 20
        self.mvn_threshold = 0.995
        self.mvn_threshold_min = 0.6
        self.mvn_rev_threshold = 0.005
        self.cond_lookback = 60
        self.prob = 0
        
        # self.lastMonth = -1
        # self.dollarVolumeBySymbol = {}
        
        self.rebalence_flag = 0
        # make sure to run the universe selection at the start of the algorithm even it's not the manth start
        self.first_month_trade_flag = 1
        self.downturn = 1
        self.cover = 1
        self.Schedule.On(self.DateRules.MonthEnd("SPY",2), self.TimeRules.At(1, 0), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.AfterMarketOpen("SPY", 0), Action(self.rebalance))
        # self.Schedule.On(self.DateRules.MonthEnd("SPY"), self.TimeRules.At(10, 0), Action(self.rebalance))
        
        monitorPlot = Chart('Monitor')
        monitorPlot.AddSeries(Series('Leverage', SeriesType.Line, 0))
        self.AddChart(monitorPlot)
        
        self.Schedule.On(self.DateRules.MonthEnd("SPY"), \
            self.TimeRules.BeforeMarketClose("SPY", 0), \
            self.Record)
    
    def Record(self):
        leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot('Monitor', 'Leverage', leverage)
    
    def monthly_rebalance(self):
        self.rebalence_flag = 1

    def CoarseSelectionFunction(self, coarse):  
        selected = []
        if not (self.rebalence_flag or self.first_month_trade_flag):
            return Universe.Unchanged
            
        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.DollarVolume > 3000000 and x.Price >0],
                                     key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]
        return [x.Symbol for x in sortedByDollarVolume]
        
    def factor_stock_selection(self, name, fine, fun_filter, filter_lmd, order):
        if filter_lmd:
            fact_lmd = lambda x: fun_filter(x) > 0 and filter_lmd(x)
        else:
            fact_lmd = lambda x: fun_filter(x) > 0
            
        fine = ([x for x in fine if fact_lmd(x) and x.CompanyReference.CountryId == "USA"
                                            and x.CompanyReference.PrimaryExchangeID in ["NYS","NAS"]
                                            and (self.Time - x.SecurityReference.IPODate).days > 180 and x.MarketCap > 2e9])
                                            
        if len(fine) == 0:
            return ([], 0)

        sortedByDEratio = sorted(fine, key=fun_filter, reverse=order)[:self.FirstFilter]
        symbols = [x.Symbol for x in sortedByDEratio]
        
        averages = dict()
        history = self.History(symbols, 150, Resolution.Daily).close.unstack(0)
        for symbol in symbols:
            if symbol in history:
                df = history[symbol].dropna()
            else:
                continue
            if df.empty or len(df)<130:
                continue
            mom=df[:-1].pct_change(126) .iloc[-1]
            if pd.notnull(mom):
                averages[symbol] = mom

        mom_list = averages.items()
        
        sortedbyMomentum = sorted(mom_list, key=lambda x: x[1], reverse=True)
        selectedSyms = [x[0] for x in sortedbyMomentum[:self.SecondFilter]]
        selectedPrices = history[selectedSyms]
        
        return (selectedSyms, (selectedPrices.iloc[-1] - selectedPrices.iloc[-10]).sum(), selectedPrices.iloc[-self.mvn_lookback:].sum(axis=1), name)
    
    def SelectFine(self, fine):
        if not (self.rebalence_flag or self.first_month_trade_flag):
            return Universe.Unchanged
            
        self.rebalence_flag = 0
        self.first_month_trade_flag = 0
        
        best_factor = self.factor_stock_selection('RevenueGrowth',fine, lambda x: x.OperationRatios.RevenueGrowth.ThreeMonths, lambda x: x.FinancialStatements.IncomeStatement.TotalRevenue.ThreeMonths > 0, True)
        self.selected_price = best_factor[2]
        self.selected_symbols = best_factor[0]
        self.Log('======SELECT FACTOR:' + str(best_factor[3]) + ', ret:' + str(best_factor[1]))
            
        return best_factor[0]
        
    def rebalance(self):
        self.result = {}
        self.liquidate = set()
        
        count = len(self.selected_symbols)
        percent = 0 if count == 0 else 1 / count

        for stock in self.selected_symbols:
            # self.Log('==== SELECTED: ' + "   " + stock.Value)
            self.result[stock] = percent
        
        invested = [ x.Symbol for x in self.Portfolio.Values if x.Invested]  
                
        for share in invested:
            if not share in self.result:
                self.liquidate.add(share)
                
        self.Log('====================== Day: ' + str(self.Time) + ' ================================')
        for security in self.result:
            self.SetHoldings(security, self.result[security])
            self.Log('==== BUYING: ' + security.Value + ': ' + str(self.result[security]))
        for security in self.liquidate:
            self.Log('==== LIQUIDATING: ' + security.Value)
            self.Liquidate(security)