import numpy as np
import statsmodels.api as sm
from arch.unitroot import PhillipsPerron
from statsmodels.tsa.stattools import adfuller
import numpy as np
import pandas as pd

def is_stationary(data, sig_level=.05):
    if adfuller(data)[1] > sig_level:
        return False
    elif PhillipsPerron(data).pvalue > sig_level:
        return False
    else:
        return True
        
def optimize(data_anchor:pd.Series, data:pd.DataFrame):
    data_train = data.copy()
    data_train['ones'] = 1 # adds a column of 1's for the alpha term of OLS
    res = sm.OLS(data_anchor, data).fit()
    
    if not is_stationary(res.resid):
        return None
        
    weights = res.params[:-1]  # remove the alpha term
    portfolio_values = data_anchor - (data * weights).sum(axis=1)
    return portfolio_values.mean(), portfolio_values.std(), weights

# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2020 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from enum import Enum
import portfolio_construction as po
import numpy as np
from System.Drawing import Color

class VerticalTransdimensionalAutosequencers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 5, 1)  # Set Start Date
        self.SetCash(1000000)  # Set Strategy Cash
        
         # GDAX (aka Coinbase) doesn't allow shorting
        self.SetBrokerageModel(BrokerageName.Bitfinex) 
        
        self.anchor = 'BTCUSD'
        tickers = [self.anchor, 'ETHUSD', 'BCHUSD', 'LTCUSD']
        
        self.symbols = []
        
        for ticker in tickers:
            symbol = self.AddCrypto(ticker, Resolution.Daily, Market.Bitfinex).Symbol
            self.symbols.append(symbol)
        
        self.data = RollingWindow[Slice](100)  # stores historical price data
        self.weights = None
        self.portfolio = None
        
        self.state = State.NEUTRAL
        
        self.SetWarmup(100)
        
        self.Train(self.DateRules.WeekStart('BTCUSD'), self.TimeRules.Midnight, self.Retrain)
        
        # below is for fancy charting
        stockPlot = Chart('Custom')
        stockPlot.AddSeries(Series('long', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Triangle))
        stockPlot.AddSeries(Series('short', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.TriangleDown))
        #stockPlot.AddSeries(Series('liquidate', SeriesType.Scatter, '$', Color.Blue, ScatterMarkerSymbol.Diamond))
        self.AddChart(stockPlot)
        
    def OnData(self, data):
        self.data.Add(data)
        
        if self.IsWarmingUp:
            return
        
        if self.portfolio is None:
            return
        
        if not all([data.Bars.ContainsKey(symbol) for symbol in self.symbols]):
            return
        
        anchor_price = data.Bars[self.anchor].Close
        others_prices = []
        for ticker in self.weights.index:
            close = data.Bars[ticker].Close
            others_prices.append(close)
       
        signal = self.portfolio.update(self, anchor_price, others_prices)
        
        scale_factor = 20
        
        # go long if not already
        if signal == 1 and self.state != State.LONG:
            self.state = State.LONG
            self.Liquidate()
            self.MarketOrder('BTCUSD', 1 * scale_factor)
            for ticker, weight in self.weights.iteritems():
                self.MarketOrder(ticker, -weight * scale_factor)
        # go short if not already
        elif signal == -1 and self.state != State.SHORT:
            self.state = State.SHORT
            self.Liquidate()
            self.MarketOrder('BTCUSD', -1 * scale_factor)
            for ticker, weight in self.weights.iteritems():
                self.MarketOrder(ticker, weight * scale_factor)
                
    def Retrain(self):
        if not self.data.IsReady:
            return
        
        try:
            # since RollingWindow is recent at top, we need to reverse it
            data = self.PandasConverter.GetDataFrame(self.data).iloc[::-1]
        except:
            self.reset()
            return
        
        # turn the closing prices for each equity into columns
        data = data['close'].unstack(level=0)
       
        data_wo_anchor = data.drop(self.anchor, axis=1)
        
        res = po.optimize(data[self.anchor], data_wo_anchor)
        
        if res is None:
            self.reset()
            return
        
        mean, std, self.weights = res

        self.portfolio = Portfolio(mean, std, self.weights)
        
    def reset(self):
        self.Liquidate()
        self.weights = None
        self.portfolio = None
        self.Log('Resetting')
        
class Portfolio:
    def __init__(self, mean, std, weights, std_factor=1.5):
        self.mean = mean
        self.std = std
        self.weights = np.array(weights)
        self.std_factor = std_factor
        self.tripped = Tripped.NEUTRAL
    
    def update(self, algorithm, anchor_price, others_prices):
        spread = anchor_price - np.sum(np.array(others_prices) * np.array(self.weights))
        
        upper = self.mean + self.std * self.std_factor
        lower = self.mean - self.std * self.std_factor
        
        algorithm.Plot('Custom', 'Spread Value', spread)
        algorithm.Plot('Custom', 'Short Value', upper)
        algorithm.Plot('Custom', 'Long Value', lower)
        
        if spread > upper:
            self.tripped = Tripped.UPPER
        # go short
        elif self.tripped == Tripped.UPPER and spread <= upper:
            self.tripped = Tripped.NEUTRAL
            algorithm.Plot('Custom', 'short', spread)
            return -1
        elif spread < lower:
            self.tripped = Tripped.LOWER
        # go long
        elif self.tripped == Tripped.LOWER and spread >= lower:
            self.tripped = Tripped.NEUTRAL
            algorithm.Plot('Custom', 'long', spread)
            return 1
        return 0
            
class State(Enum):
    LONG = 0
    NEUTRAL = 1
    SHORT = 2

class Tripped(Enum):
    LOWER = 0
    NEUTRAL = 1
    UPPER = 2