import xarray as xr
import numpy as np
import pandas as pd

import qnt.ta as qnta
import qnt.backtester as qnbt
import qnt.data as qndata


def load_data(period):
    
    futures = qndata.futures_load_data(assets=['F_CL'], tail=period, dims=('time','field','asset'))
    
    ap = qndata.blsgov.load_series_data('APU000072511', tail=period)
    
    # convert to pandas.DataFrame
    ap = pd.DataFrame(ap) 
    ap = ap.set_index('pub_date') 

    # remove yearly average data, see period dictionary
    ap = ap[ap['period'] != 'M13']
    
    # convert to xarray
    ap = ap['value'].to_xarray().rename(pub_date='time').assign_coords(time=pd.to_datetime(ap.index.values))
    
    # return both time series
    return dict(ap=ap, futures=futures), futures.time.values


def window(data, max_date: np.datetime64, lookback_period: int):
    # the window function isolates data which are needed for one iteration
    # of the backtester call
    
    min_date = max_date - np.timedelta64(lookback_period, 'D')
    
    return dict(
        futures = data['futures'].sel(time=slice(min_date, max_date)),
        ap = data['ap'].sel(time=slice(min_date, max_date))
    )


def strategy(data, state):
    
    close = data['futures'].sel(field='close')
    ap = data['ap']
    
    # the strategy complements indicators based on the Futures price with macro data
    # and goes long/short or takes no exposure:
    
    if ap.isel(time=-1) > ap.isel(time=-2) \
            and close.isel(time=-1) > close.isel(time=-20):
        return xr.ones_like(close.isel(time=-1)), 1
    
    elif ap.isel(time=-1) < ap.isel(time=-2) \
            and ap.isel(time=-2) < ap.isel(time=-3) \
            and ap.isel(time=-3) < ap.isel(time=-4) \
            and close.isel(time=-1) < close.isel(time=-40):
        return -xr.ones_like(close.isel(time=-1)), 1 
    
    # When the state is None, we are in the beginning and no weights were generated.
    # We use buy'n'hold to fill these first days.
    elif state is None: 
        return xr.ones_like(close.isel(time=-1)), None
    
    else:
        return xr.zeros_like(close.isel(time=-1)), 1


weights, state = qnbt.backtest(
    competition_type='futures',
    load_data=load_data,
    window=window,
    lookback_period=365,
    start_date="2006-01-01",
    strategy=strategy,
    analyze=True,
    build_plots=True
)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-11-6c5d77a629e2> in <module>
     75     strategy=strategy,
     76     analyze=True,
---> 77     build_plots=True
     78 )

~/book/qnt/backtester.py in backtest(competition_type, strategy, load_data, lookback_period, test_period, start_date, end_date, window, step, analyze, build_plots, collect_all_states)
    325         data, time_series = extract_time_series(data)
    326         print("Run strategy...")
--> 327         result = strategy_wrap(data, None)
    328         result, state = unpack_result(result)
    329         log_info("---")

<ipython-input-11-6c5d77a629e2> in strategy(data, state)
     61     # We use buy'n'hold to fill these first days.
     62     elif state is None:
---> 63         return xr.ones_like(close.isel(time=-1)), None
     64 
     65     else:

/usr/local/lib/python3.7/site-packages/xarray/core/dataarray.py in isel(self, indexers, drop, missing_dims, **indexers_kwargs)
   1128         # lists, or zero or one-dimensional np.ndarray's
   1129 
-> 1130         variable = self._variable.isel(indexers, missing_dims=missing_dims)
   1131 
   1132         coords = {}

/usr/local/lib/python3.7/site-packages/xarray/core/variable.py in isel(self, indexers, missing_dims, **indexers_kwargs)
   1164 
   1165         key = tuple(indexers.get(dim, slice(None)) for dim in self.dims)
-> 1166         return self[key]
   1167 
   1168     def squeeze(self, dim=None):

/usr/local/lib/python3.7/site-packages/xarray/core/variable.py in __getitem__(self, key)
    809         """
    810         dims, indexer, new_order = self._broadcast_indexes(key)
--> 811         data = as_indexable(self._data)[indexer]
    812         if new_order:
    813             data = duck_array_ops.moveaxis(data, range(len(new_order)), new_order)

/usr/local/lib/python3.7/site-packages/xarray/core/indexing.py in __getitem__(self, key)
   1300     def __getitem__(self, key):
   1301         array, key = self._indexing_array_and_key(key)
-> 1302         return array[key]
   1303 
   1304     def __setitem__(self, key, value):

IndexError: index -1 is out of bounds for axis 0 with size 0