#region imports

from AlgorithmImports import *
from universes import *
import json
#endregion


"""
Simple Bollinger bands pair trading also
open long or short pair trade at a 2 std deviation, close at middle band cross in 2 steps:
  - intraday as soon as middle band crossed
  - end of day on re-cross in the "wrong way" to allow participation in the trend, we may  loose a little versus closing all @ middle band 
 but should have opportunities for large gain in longer trends

research.ipynb allows us to see and manipulate the saved states of the algo (in live mode only, nothing is saved during backtest)

We save the state for each trade as things could crash in the middle and we do not want to reconcile manually...

TODO implement update_adj_price before going live!
TODO ADD (lots) of logging before go live
TODO RECONCILIATION: make sure both legs of the pair trade actually execute to avoid being unhedged
TODO RECONCILIATION: fill self.pair_qty in 
TODO Add earnings check: do not enter 1 or 2 weeks (analysis to be completed) before earnings
TODO Add earnings check: exit 2 days before earnings to avoid randomness.
"""

class PairTrading(QCAlgorithm):

    def Initialize(self):

        ###############################################################
        ############## START ALGO CONFIG

        #relies on pairs list defined in universes.py 
        pairs_2021 = rw_2021
        # pairs_2021 = list(set((x,y) for x, y in rw_2021) | set((x,y) for x, y in rw_fab_pairs_2021))

        BACKTEST_START_YEAR = 2015

        # Notional size for each leg of the trade
        INITIAL_TRADE_SIZE = 15000

        # PARAM FOR MAX NUMBER OF CONCURRENT PAIR TRADE 
        # set to len(pairs_2021) to test the overall stability, taking the first n slots introduces a bit of randomness 
        # TODO is it worth ranking pairs "somehow" (selection score or recent performance...) before allocating new slots?
        self.MAX_POS = 2 #len(pairs_2021)

        # MAX PERCENTAGE OF PORTFOLIO IN SINGLE SECURITY, no more than 3 trades if small number of pairs else 5%, that seems reasonable
        self.MAX_CONCENTRATION = max(0.05, 3.0 / self.MAX_POS)

        # TODO MAX SECTOR/INDUSTRY CONCENTRATION
        
        year_pairs = {
            # 2010 : pairs_2021,
            # 2011 : pairs_2021,
            # 2012 : pairs_2021,
            # 2013 : pairs_2021,
            # 2014 : pairs_2021,
            2015 : pairs_2021,
            2016 : pairs_2021,
            2017 : pairs_2021,
            2018 : pairs_2021,
            2019 : pairs_2021,
            2020 : pairs_2021,
            2021 : pairs_2021,
            2022 : pairs_2021,
            }

        self.SetStartDate(BACKTEST_START_YEAR, 1, 1)  # Set Start Date
        self.SetCash(self.MAX_POS * INITIAL_TRADE_SIZE)  # Set Strategy Cash

        ###############################################################
       
        self.year_pair_name = {}
        #self.ids = {}
        
        self.ratios = {}
        self.bbs = {}

        added_sym = set()
        added_pairs = set()
        
        self.trade_on = {}
        self.pstate = {}
        self.pair_qty = {}
        
        # we will adjust the prices ourselves to allow the algo to deal with it without a restart once live
        self.price_adj_factor = {} 

        # we load state data for live before we initialise defaults
        self.load_live_state()

        for year, pairs in year_pairs.items():
            for p in pairs:
                s1 = p[0]
                s2 = p[1]
                
                
                self.add_symbol(s1, added_sym)
                self.add_symbol(s2, added_sym)
                pair_name = self.get_pair_name(s1, s2)
                if year not in self.year_pair_name:
                    self.year_pair_name[year] = set()
                self.year_pair_name[year].add(pair_name)

                if pair_name not in added_pairs:
                    self.ratios[pair_name] = None
                    self.bbs[pair_name] = BollingerBands(20, 2) # try sma instead of default ema
                    added_pairs.add(pair_name)
                    self.pstate[pair_name] = self.pstate.get(pair_name, 0)
                    self.pair_qty[pair_name] = self.pair_qty.get(pair_name, [0 , 0])
                    self.trade_on[pair_name] = self.trade_on.get(pair_name, 0)                
        
        self.pday = 0
        self.rebalance = False

        ref_sym = "SPY"#pairs_2021[0][0]# "LNG"
        self.AddEquity("SPY")
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(ref_sym, 15), self.rebalance_flag)

        self.SetWarmup(TimeSpan.FromDays(30))

    def add_symbol(self, symbol, added_sym):
        if symbol in added_sym:
            return
        security = self.AddEquity(symbol, Resolution.Minute)

        # we set leverage to 2.5 to allow for long short full account and margin portfolio 
        #security.MarginModel = PatternDayTradingMarginModel()
        security.SetLeverage(3)
        if self.LiveMode:
            security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        else:
            security.SetDataNormalizationMode(DataNormalizationMode.TotalReturn)

        added_sym.add(symbol)
        self.price_adj_factor[symbol] = 1.0        
    
    def rebalance_flag(self):
        self.rebalance = True

    def OnData(self, data: Slice):
        # workaround Schedule.On not called before end of warmup
        if self.IsWarmingUp: 
            if self.Time.hour == 15 and self.Time.minute == 30:
                self.rebalance = True      

        #if self.pday != self.Time.day:
        if self.rebalance:
            # golong = set()
            # goshort = set()
            # udpate indicators first
            for pair in self.ratios.keys():
                stock1, stock2 = self.get_2symbols_from_pair_name(pair)            
                if self.Securities[stock2].Price == 0:
                    continue
                
                #first let's deal with dividends and split etc...
                # Slice: Splits Splits;
                # Slice: Dividends Dividends;
                # Slice: Delistings Delistings;
                self.update_adj_prices(data)

                self.ratios[pair] = self.get_ratio(stock1, stock2)
                self.bbs[pair].Update(IndicatorDataPoint(self.Time, self.ratios[pair]))
            
            if not self.IsWarmingUp:
                # First, close all possible positions to know how many knew we will be abe to open
                self.close_positions()

                # Open new positions
                # but first count open spread trade live so we don't open more than self.MAX_POS
                s2qty = self.open_positions()
                
                # TODO add a trending component when crossing middle band: close half then trail with running min of 1 std
                

            self.rebalance = False
            self.pday = self.Time.day
        elif not self.IsWarmingUp: #intraday take profit asap
            self.take_profit_intraday(data)
    
    def get_pair_name(self, s1, s2):
        return s1 + '#' + s2

    def get_2symbols_from_pair_name(self, pair_name):
        return pair_name.split('#')

    def get_state(self, ratio, pair):
        state = 0
        if ratio > self.bbs[pair].UpperBand.Current.Value:
            state = 2
        elif ratio > self.bbs[pair].MiddleBand.Current.Value:
            state = 1
        elif ratio < self.bbs[pair].LowerBand.Current.Value:
            state = -2
        elif ratio < self.bbs[pair].MiddleBand.Current.Value:
            state = -1
        return state

    def get_ratio(self, stock1, stock2):
        return self.Securities[stock1].Price * self.price_adj_factor[stock1] / (self.Securities[stock2].Price * self.price_adj_factor[stock1])

    def update_adj_prices(self, data: Slice):
        # TODO deal with adj prices for live trading
        if self.LiveMode:
            raise Exception("not implemented")        

    def close_positions(self):
        for pair in self.ratios.keys():
            stock1, stock2 = self.get_2symbols_from_pair_name(pair)
            if self.Securities[stock2].Price == 0:
                continue
            if self.bbs[pair].IsReady:
                state = self.get_state(self.ratios[pair], pair)
                if not self.IsWarmingUp and len(self.ratios) <= 10: # cannot plot more than 10 series
                    self.Plot('state', pair, state)
        
                tag = f"{pair} {state} {self.pstate[pair]}  {self.ratios[pair]}"
                if pair in self.year_pair_name.get(self.Time.year, []) and state == 2:
                    pass
                elif pair in self.year_pair_name.get(self.Time.year, []) and state == -2:
                    pass   
                elif self.pstate[pair] < 0 and state > 0 and self.trade_on[pair] > 0:
                    self.trade_on[pair] = 0
                    q1, q2 = self.pair_qty[pair]
        
                    # because we could be trading the stock on 2 sides resulting qty could be 0
                    if q1 != 0 : 
                        self.MarketOrder(stock1, -q1, tag=tag+",close")
                    if q2 != 0 : 
                        self.MarketOrder(stock2, -q2, tag=tag+",close")
                    self.pair_qty[pair] = [0, 0]
                    self.save_live_state()
                elif self.pstate[pair] > 0 and state < 0 and self.trade_on[pair] < 0:
                    self.trade_on[pair] = 0
                    q1, q2 = self.pair_qty[pair]
        
                    if q1 != 0 : 
                        self.MarketOrder(stock1, -q1, tag=tag+",close")
                    if q2 != 0 : 
                        self.MarketOrder(stock2, -q2, tag=tag+",close")
                    self.pair_qty[pair] = [0, 0]
                    self.save_live_state()
        

    def open_positions(self):
        current_pos_count = len([ v for v in self.trade_on.values() if v != 0])
        for pair in self.ratios.keys():
            stock1, stock2 = self.get_2symbols_from_pair_name(pair)
        
            if self.Securities[stock2].Price == 0:
                continue
            if self.bbs[pair].IsReady:
                state = self.get_state(self.ratios[pair], pair)
                if current_pos_count >= self.MAX_POS:
                    self.pstate[pair] = state
                    break
        
                tag = f"{pair} {state} {self.pstate[pair]} {self.ratios[pair]}"
                if pair in self.year_pair_name.get(self.Time.year, []) and state == 2:
                    #if self.Portfolio[stock1].Quantity >= 0:
                    if self.trade_on[pair] >= 0:
                        self.trade_on[pair] = -2
                        weight = 1.0 / self.MAX_POS # len(self.year_pair_name[self.Time.year])
                        price1 = self.Securities[stock1].Price
                        price2 = self.Securities[stock2].Price
                        pf_value = self.Portfolio.TotalPortfolioValue
        
                        if abs(-weight + self.Portfolio[stock1].Quantity * price1 / pf_value) > self.MAX_CONCENTRATION:
                            # would be above concentration dont trade
                            continue
                        if abs(weight + self.Portfolio[stock1].Quantity * price2 / pf_value) > self.MAX_CONCENTRATION:
                            # would be above concentration dont trade
                            continue
        
                        weightval = weight * pf_value
                        s1qty = round(weightval / price1 , 0)
                        s2qty = round(weightval / price2 , 0)
        
                        # !!! rounding could make it 0 for just 1 piece
                        # TODO check that the resulting ratio is still close to the initial one so that we have a change to converge profitably!
                        # TODO linked to the above => beware of adjusted prices in backtest which might make the ratio possible or impossible to trade in the past
                        if s1qty != 0 and s2qty != 0: 
                            self.MarketOrder(stock1, -s1qty, tag=tag)
                            self.MarketOrder(stock2, s2qty, tag=tag)
                            self.pair_qty[pair] = [-s1qty, s2qty]
                            self.save_live_state()
                            current_pos_count += 1
                #elif self.pstate == -2 and state > -2:
                elif pair in self.year_pair_name.get(self.Time.year, []) and state == -2:
                    #if self.Portfolio[stock1].Quantity <= 0:
                    if self.trade_on[pair] <= 0:
                        self.trade_on[pair] = 2
                        weight = 1.0 / self.MAX_POS # len(self.year_pair_name[self.Time.year])
                        price1 = self.Securities[stock1].Price
                        price2 = self.Securities[stock2].Price
                        pf_value = self.Portfolio.TotalPortfolioValue
        
                        if abs(weight + self.Portfolio[stock1].Quantity * price1 / pf_value) > self.MAX_CONCENTRATION:
                            # would be above concentration dont trade
                            continue
                        if abs(-weight + self.Portfolio[stock1].Quantity * price2 / pf_value) > self.MAX_CONCENTRATION:
                            # would be above concentration dont trade
                            continue
                
                        weightval = weight * pf_value
                        s1qty = round(weightval / price1 , 0)
                        s2qty = round(weightval / price2 , 0)
                        if s1qty != 0 and s2qty != 0:                             
                            self.MarketOrder(stock1, s1qty, tag=tag)
                            self.MarketOrder(stock2, -s2qty, tag=tag)
                            self.pair_qty[pair] = [s1qty, -s2qty]
                            self.save_live_state()
                            current_pos_count += 1
        
                self.pstate[pair] = state
        

    def take_profit_intraday(self, data):
        # TODO loop only on trade_on pairs => faster
        for pair in self.ratios.keys():
            if self.trade_on[pair] == 0:
                continue
            stock1, stock2 = self.get_2symbols_from_pair_name(pair)
            if self.ratios[pair] != None and data.Bars.ContainsKey(stock1) and data.Bars.ContainsKey(stock2):
                ratio = data.Bars[stock1].Close / data.Bars[stock2].Close
        
                if self.trade_on[pair] == 2 and ratio > self.bbs[pair].MiddleBand.Current.Value or \
                    self.trade_on[pair] == -2 and ratio < self.bbs[pair].MiddleBand.Current.Value:
                    self.trade_on[pair] = self.trade_on[pair] / 2 # ie -1 or 1                        
                    q1, q2 = self.pair_qty[pair]
                    rq1 = round(q1/2,0)
                    rq2 = round(q2/2,0)
                    if rq1 != 0 and rq2 != 0: # otherwise will get unbalanced
                        self.MarketOrder(stock1, -rq1, tag="half")
                        self.MarketOrder(stock2, -rq2, tag="half")
            
                        self.pair_qty[pair] = [q1-rq1, q2-rq2]  
                        self.save_live_state()

    def save_live_state(self):
       if self.LiveMode:
            self.ObjectStore.Save("pair_qty", json.dumps(self.pair_qty))
            self.ObjectStore.Save("trade_on", json.dumps(self.trade_on))
            self.ObjectStore.Save("pstate", json.dumps(self.pstate))

    def load_live_state(self):
       if self.LiveMode:
            if self.ObjectStore.ContainsKey("pair_qty"):
                self.pair_qty = json.loads(self.ObjectStore.Read("pair_qty"))
            if self.ObjectStore.ContainsKey("trade_on"):
                self.trade_on = json.loads(self.ObjectStore.Read("trade_on"))
            if self.ObjectStore.ContainsKey("pstate"):
                self.pstate = json.loads(self.ObjectStore.Read("pstate"))

#region imports
from AlgorithmImports import *
#endregion


# Your New Python File
rw_2021 = [['KO','PEP']#,['F','GM'],
]