from QuantConnect.DataSource import SmartInsiderTransaction
import numpy as np
from scipy.optimize import curve_fit

class SmartInsiderCorporateBuybacksAlgorithm(QCAlgorithm):

    def Initialize(self):
        # parameter: minimal information coefficient accpetable
        self.minIC = 0.05
        # parameter: minimal expected return accpetable
        self.minExpectedReturn = 0.005
        
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100000) 
        
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(timedelta(days=252)))
        self.SetExecution(ImmediateExecutionModel())
        
        # include "MMM", "BA", "GS", "HON", "JNJ", "MSFT", "TRV" as a result of the research
        self.symbols = [self.AddEquity(symbol, Resolution.Minute).Symbol for symbol in ["MMM", "GS", "HON", "JNJ", "MSFT", "TRV", "V"]]
        self.buybackSymbols = {symbol: self.AddData(SmartInsiderTransaction, symbol, Resolution.Daily).Symbol for symbol in self.symbols}
        
        # dict contains rolling windows storing tradebar data
        self.history = {symbol: RollingWindow[TradeBar](252*5) for symbol in self.symbols}
        
        # warm up rolling windows
        data = self.History(self.symbols, 252*5, Resolution.Daily)
        for symbol in self.symbols:
            for time, bar in data.loc[symbol].iterrows():
                tradeBar = TradeBar(time, symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
                self.history[symbol].Add(tradeBar)
            # set up consolidator for future auto-update
            self.Consolidate(symbol, Resolution.Daily, self.DailyBarHandler)
            
        # schedule daily check for entering position
        self.Schedule.On(self.DateRules.EveryDay("MMM"), self.TimeRules.At(10, 0), self.Entry)
        
    def DailyBarHandler(self, bar):
        self.history[bar.Symbol].Add(bar)
        
    def Entry(self):
        ''' check entry signal '''
        for symbol, tSymbol in self.buybackSymbols.items():
            self.transaction = self.History(tSymbol, timedelta(days=1), Resolution.Daily)
            if self.transaction.empty: continue
            
            # buyback%
            self.transaction = self.transaction.usdvalue.unstack("symbol") / self.transaction.usdmarketcap.unstack("symbol")
            # sum up for daily
            self.transaction.index = self.transaction.index.date
            self.transaction = self.transaction.groupby(self.transaction.index).sum()
            # log
            self.transaction = np.log(self.transaction)
            
            self.GetData(symbol)
            
            # find IC
            ic = self.GetIC(symbol)
            # discontinue if IC too low
            if ic < self.minIC: continue
        
            # find expected return
            expectation = self.GetExpectation(symbol)
            # discontinue if expected return too low
            if expectation < self.minExpectedReturn: continue
            
            # emit insight for entry
            self.EmitInsights(Insight.Price(symbol, timedelta(days=252), InsightDirection.Up, None, None, None, expectation*ic))
    
    def GetData(self, symbol):
        ''' get the processed dataframe
        Args:
            symbol: the symbol that we wish to get the processed dataframe'''
        # get historical close price data
        data = pd.DataFrame(self.history[symbol])[::-1]
        history = data.applymap(lambda bar: bar.Close)
        history.index = data.applymap(lambda bar: bar.EndTime.date()).values.flatten().tolist()
        # 1y forward
        history = history.pct_change(252).shift(-252).dropna()
        
        # get buyback transaction data
        transactionHistory = self.History(self.buybackSymbols[symbol], 252*5, Resolution.Daily)
        if transactionHistory.empty: return -1
        # buyback%
        transactionHistory = transactionHistory.usdvalue.unstack("symbol") / transactionHistory.usdmarketcap.unstack("symbol")
        # sum up for daily
        transactionHistory.index = transactionHistory.index.date
        transactionHistory = transactionHistory.groupby(transactionHistory.index).sum()
        
        # concatenate the dataframes to left only the slices with data
        df = pd.concat([history, np.log(transactionHistory)], axis=1).replace([np.inf, -np.inf], np.nan).dropna()
        
        self.df = df
        
    def GetIC(self, symbol):
        ''' get the correlation coefficient as information coefficient
        Args:
            symbol: the symbol that we wish to get IC between 1y forward return and buyback transaction
        Return:
            (float) the correlation coefficient/IC'''
        return self.df.corr().values[0, 1]
        
    def GetExpectation(self, symbol):
        ''' get expected return in 1 year given the buyback%
        Args:
            symbol: the Symbol that we wish to get ins expected return
        Return:
            (float) expected return%'''
        
        def Function(x, a, b, c):
            ''' the function to be fitted '''
            return a * np.exp(-b * x) + c
        
        try:
            popt, pcov = curve_fit(Function, self.df.iloc[:, 1].values, self.df.iloc[:, 0].values)
        # in case no optimal coefficients can be fit
        except:
            return -1 
        
        return Function(self.transaction, *popt).values