from sklearn import linear_model
import numpy as np
import pandas as pd
from scipy import stats
from math import floor
from datetime import timedelta
from datetime import time

class PairsTradingAlgo(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2020,6,1)
        self.SetEndDate(self.StartDate + timedelta(10))
        self.SetCash(10000)
        # self.Settings.FreePortfolioValuePercentage = 0.05
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        tickers = ['XOM', 'CVX', 'SPY']
        
        for t in tickers:
            self.AddEquity(t, Resolution.Hour).Symbol
            
        self.SetBenchmark('SPY')
        
        self.SetAlpha(PairsAlpha())
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(TrailingStop())

class PairsAlpha(AlphaModel):
    
    def __init__(self):
        self.window_length = 13
        self.threshold = 1.5
        self.symbol_data = {}

    def Update(self, algo, data):
        insights = []
        
        for security in data.Keys:
            symbol = security.Value
            if symbol == 'SPY': continue
            if symbol not in self.symbol_data.keys():
                self.symbol_data[symbol] = RollingWindow[TradeBar](self.window_length)
            if data.Bars.ContainsKey(symbol):
                self.symbol_data[symbol].Add(data[security])
        
        symbols = [x for x in self.symbol_data.keys()]
        
        if not (self.symbol_data[symbols[0]].IsReady and self.symbol_data[symbols[1]].IsReady):
            return insights
        
        price_x = pd.Series([float(i.Close) for i in self.symbol_data[symbols[0]]], 
                             index = [i.Time for i in self.symbol_data[symbols[0]]])
                             
        price_y = pd.Series([float(i.Close) for i in self.symbol_data[symbols[1]]], 
                             index = [i.Time for i in self.symbol_data[symbols[1]]])

        spread_regr = self.regr(np.log(price_x), np.log(price_y))
        self.spread_mean = np.mean(spread_regr)
        self.spread_std = np.std(spread_regr)
        
        algo.Log('%s: %s' % (str(symbols[0]), str(algo.Portfolio[symbols[0]].Invested)))
        algo.Log('%s: %s' % (str(symbols[1]), str(algo.Portfolio[symbols[1]].Invested)))
        if not (algo.Portfolio[symbols[0]].Invested or algo.Portfolio[symbols[1]].Invested):
            if spread_regr[-1] > self.spread_mean + self.threshold * self.spread_std:
                insights.append(Insight.Price(symbols[0], timedelta(days=30), InsightDirection.Up))
                insights.append(Insight.Price(symbols[1], timedelta(days=30), InsightDirection.Down))
                algo.Log('insights emitted')
        
            elif spread_regr[-1] < self.spread_mean - self.threshold * self.spread_std:
                insights.append(Insight.Price(symbols[1], timedelta(days=30), InsightDirection.Up))
                insights.append(Insight.Price(symbols[0], timedelta(days=30), InsightDirection.Down))
                algo.Log('insights emitted')
                
        return insights
    
    def regr(self, x, y):
        regr = linear_model.LinearRegression()
        x_constant = np.column_stack([np.ones(len(x)), x])
        regr.fit(x_constant, y)
        beta = regr.coef_[0]
        alpha = regr.intercept_
        spread = y - x*beta - alpha
        return spread
        
class TrailingStop(RiskManagementModel):
    def __init__(self, drawdown_threshold = 0.01):
        self.drawdown_threshold = -abs(drawdown_threshold)
        self.max_profits = {}
        
    def ManageRisk(self, algo, targets):
        targets = []
        
        for kvp in algo.Securities:
            security = kvp.Value
            symbol = security.Symbol.Value
            current_profit = security.Holdings.UnrealizedProfitPercent
            
            if not security.Invested:
                continue
            
            if symbol not in self.max_profits.keys():
                self.max_profits[symbol] = current_profit
            self.max_profits[symbol] = np.maximum(self.max_profits[symbol], current_profit)
            
            drawdown = current_profit - self.max_profits[symbol]
            # algo.Log('%s - Current profit: %.3f, Max profit: %.3f, Current drawdown: %.3f' % (str(symbol), current_profit, self.max_profits[symbol], drawdown))
            if drawdown < self.drawdown_threshold:
                targets.append(PortfolioTarget(symbol, 0))
                algo.Log('%s position to be liquidated' % str(symbol))
                
                # Reset dictionary
                del self.max_profits[symbol]
        
        return targets