"""
This module contains class for ETF trick generation and futures roll function, described in Marcos Lopez de Prado's
book 'Advances in Financial Machine Learning' ETF trick class can generate ETF trick series either from .csv files
or from in memory pandas DataFrames
"""

# Imports
import warnings
import pandas as pd
import numpy as np


class ETFTrick:
    """
    Contains logic of vectorised ETF trick implementaion. Can used for both memory data frames (pd.DataFrame) and csv files.
    All data frames, files should be processed in a specific format, described in examples
    """

    def __init__(self, open_df, close_df, alloc_df, costs_df, rates_df=None, index_col=0):
        """
        Constructor

        Creates class object, for csv based files reads the first data chunk.
        :param open_df: (pd.DataFrame or string): open prices data frame or path to csv file,
         corresponds to o(t) from the book
        :param close_df: (pd.DataFrame or string): close prices data frame or path to csv file, corresponds to p(t)
        :param alloc_df: (pd.DataFrame or string): asset allocations data frame or path to csv file (in # of contracts),
         corresponds to w(t)
        :param costs_df: (pd.DataFrame or string): rebalance, carry and dividend costs of holding/rebalancing the
         position, corresponds to d(t)
        :param rates_df: (pd.DataFrame or string): dollar value of one point move of contract includes exchange rate,
         futures contracts multiplies). Corresponds to phi(t)
         For example, 1$ in VIX index, equals 1000$ in VIX futures contract value.
         If None then trivial (all values equal 1.0) is generated
        :param index_col: (int): positional index of index column. Used for to determine index column in csv files

        """

        warnings.warn("This a beta version of ETF trick. Please proof check the results", DeprecationWarning)
        self.index_col = index_col
        self.prev_k = 1.0  # Init with $1 as initial value

        # We need to track allocations vector change on previous step
        # Previous allocation change is needed for delta component calculation
        self.prev_allocs_change = False
        self.prev_h = None  # To find current etf_trick value we need previous h value

        self.data_dict = {}
        self.iter_dict = None  # Dictionary of csv files iterators, None for in_memory ETF trick calculation
        self.init_fields = None  # Dictionary of initial fields values, needed to call reset method

        if isinstance(alloc_df, str):
            # String values for open, close, alloc, costs and rates mean that we generate ETF trick from csv files
            # Remember constructor fields for possible reset() method call
            self.init_fields = {'open_df': open_df, 'close_df': close_df, 'alloc_df': alloc_df, 'costs_df': costs_df,
                                'rates_df': rates_df, 'index_col': index_col}
            self.iter_dict = dict.fromkeys(['open', 'close', 'alloc', 'costs', 'rates'], None)

            # Create file iterators
            self.iter_dict['open'] = pd.read_csv(open_df,
                                                 iterator=True,
                                                 index_col=self.index_col,
                                                 parse_dates=[self.index_col])
            self.iter_dict['close'] = pd.read_csv(close_df,
                                                  iterator=True,
                                                  index_col=self.index_col,
                                                  parse_dates=[self.index_col])
            self.iter_dict['alloc'] = pd.read_csv(alloc_df,
                                                  iterator=True,
                                                  index_col=self.index_col,
                                                  parse_dates=[self.index_col])
            self.iter_dict['costs'] = pd.read_csv(costs_df,
                                                  iterator=True,
                                                  index_col=self.index_col,
                                                  parse_dates=[self.index_col])

            if rates_df is not None:
                self.iter_dict['rates'] = pd.read_csv(rates_df,
                                                      iterator=True,
                                                      index_col=self.index_col,
                                                      parse_dates=[self.index_col])

            # Get headers(column names) from csv files (except index col) which correspond to security names
            self.securities = list(pd.read_csv(alloc_df, nrows=0, header=0, index_col=self.index_col))

        elif isinstance(alloc_df, pd.DataFrame):
            self.data_dict['open'] = open_df
            self.data_dict['close'] = close_df
            self.data_dict['alloc'] = alloc_df
            self.data_dict['costs'] = costs_df
            self.data_dict['rates'] = rates_df
            self.securities = self.data_dict['alloc'].columns  # Get all securities columns

            if rates_df is None:
                self.data_dict['rates'] = open_df.copy()
                # Set trivial(1.0) exchange rate if no data is provided
                self.data_dict['rates'][self.securities] = 1.0

            # Align all securities columns in one order
            for df_name in self.data_dict:
                self.data_dict[df_name] = self.data_dict[df_name][self.securities]

            self._index_check()
        else:
            raise TypeError('Wrong input to ETFTrick class. Either strings with paths to csv files, or pd.DataFrames')

        self.prev_allocs = np.array([np.nan for _ in range(0, len(self.securities))])  # Init weights with nan values

    def _append_previous_rows(self, cache):
        """
        Uses latest two rows from cache to append into current data. Used for csv based ETF trick, when the next
        batch is loaded and we need to recalculate K value which corresponds to previous batch.
        :param cache: (dict): dictionary which pd.DataFrames with latest 2 rows of open, close, alloc, costs, rates
        :return: (pd.DataFrame): data frame with close price differences (updates self.data_dict)
        """
        # Latest index from previous data_df(cache)
        max_prev_index = cache['open'].index.max()
        second_max_prev_index = cache['open'].index[-2]
        # Add the last row from previous data chunk to a new chunk
        for df_name in self.data_dict:
            temp_df = self.data_dict[df_name]
            temp_df.loc[max_prev_index, :] = cache[df_name].iloc[-1]
            self.data_dict[df_name] = temp_df

        # To recalculate latest row we need close price differences
        self.data_dict['close'].loc[second_max_prev_index, :] = cache['close'].loc[second_max_prev_index, :]
        # that is why close_df needs 2 previous chunk rows to omit first row nans

        for df_name in self.data_dict:
            self.data_dict[df_name].sort_index(inplace=True)  # Sort data frames after all appends
            self.data_dict[df_name] = self.data_dict[df_name][
                self.securities]  # Align all securities columns in one order

        # Get price diffs, take values from the second row (first is nan)
        price_diff = self.data_dict['close'].diff().iloc[1:]
        # Delete second max row from previous data chunk in close_df
        self.data_dict['close'] = self.data_dict['close'].iloc[1:]
        return price_diff

    def generate_trick_components(self, cache=None):
        """
        Calculates all etf trick operations which can be vectorised. Outputs multilevel pandas data frame.
        Generated components:
        'w': alloc_df
        'h_t': h_t/K value from ETF trick algorithm from the book. Which K to use is based on previous values and cannot be vectorised
        'close_open': close_df - open_df
        'price_diff': close price differences
        'costs': costs_df
        'rate': rates_df
        :param cache: (dict of pd.DataFrames): dictionary which contains latest 2 rows of open, close, rates, alloc, costs, rates data
        :return: (pd.DataFrame): pandas data frame with columns in a format: component_1/asset_name_1, component_1/asset_name_2, ..., component_6/asset_name_n
        """
        if cache:
            price_diff = self._append_previous_rows(cache)
        else:
            price_diff = self.data_dict['close'].diff()

        next_open_df = self.data_dict['open'].shift(-1)  # Generate next open prices
        close_open_diff = self.data_dict['close'].sub(self.data_dict['open'])  # close - open data frame
        self.data_dict['alloc']['abs_w_sum'] = self.data_dict['alloc'].abs().sum(
            axis=1)  # For each row generate absolute values sum for all assets

        # Allocations deleverage component
        delever_df = self.data_dict['alloc'].div(self.data_dict['alloc']['abs_w_sum'], axis='index')
        next_open_mul_rates_df = next_open_df.mul(
            self.data_dict['rates'], axis='index')  # o(t+1) * phi(t)

        # Generate calculated h_t values for each row
        # For complete h_t calculation multiplying by current K_t is needed(can't be vectorised)
        h_without_k = delever_df.div(next_open_mul_rates_df)

        weights_df = self.data_dict['alloc'][self.securities]  # Align all securities columns
        h_without_k = h_without_k[self.securities]
        close_open_diff = close_open_diff[self.securities]
        price_diff = price_diff[self.securities]

        return pd.concat(
            [weights_df, h_without_k, close_open_diff, price_diff, self.data_dict['costs'], self.data_dict['rates']],
            axis=1,
            keys=[
                'w', 'h_t', 'close_open', 'price_diff', 'costs',
                'rate'])  # Generate data frame with all precomputed info needed for ETF trick

    def _update_cache(self):
        """
        Updates cache (two previous rows) when new data batch is read into the memory. Cache is used to
        recalculate ETF trick value which corresponds to previous batch last row. That is why we need 2 previous rows
        for close price difference calculation
        :return: (dict): dictionary with open, close, alloc, costs and rates last 2 rows
        """
        cache_dict = {'open': self.data_dict['open'].iloc[-2:], 'close': self.data_dict['close'].iloc[-2:],
                      'alloc': self.data_dict['alloc'].iloc[-2:], 'costs': self.data_dict['costs'].iloc[-2:],
                      'rates': self.data_dict['rates'].iloc[-2:]}
        return cache_dict

    def _chunk_loop(self, data_df):
        """
        Single ETF trick iteration for currently stored(with needed components) data set in memory (data_df).
        For in-memory data set would yield complete ETF trick series, for csv based
        would generate ETF trick series for current batch.
        :return: (pd.Series): pandas Series with ETF trick values
        """
        etf_series = pd.Series()
        for index, row in zip(data_df.index, data_df.values):
            # Split row in corresponding values for ETF trick
            weights_arr, h_t, close_open, price_diff, costs, rate = np.array_split(row, 6)  # pylint: disable=unbalanced-tuple-unpacking
            # Replaces nan to zeros in allocations vector
            weights_arr = np.nan_to_num(weights_arr)

            # Convert np.bool to bool
            # Boolean flag of allocations vector change
            allocs_change = bool(
                ~(self.prev_allocs == weights_arr).all())  # Not(all elements in prev_w equal current_w)
            if self.prev_allocs_change is True:
                delta = close_open  # delta from book algorithm
            else:
                delta = price_diff
            if self.prev_h is None:
                # Previous h value is needed for current K calculation. The first iteration sets up prev_h value
                self.prev_h = h_t * self.prev_k
                # K is equal to 1 on the first iteration
                etf_series[index] = self.prev_k
            else:
                if self.prev_allocs_change is True:
                    # h_t is equal to previous h_t, or h_t * k if allocations vector changes (roll dates)
                    self.prev_h = h_t * self.prev_k

                k = self.prev_k + \
                    np.nansum(self.prev_h * rate * (delta + costs))
                etf_series[index] = k

                self.prev_k = k
                # Update previous allocation vector change
                self.prev_allocs_change = allocs_change
                self.prev_allocs = weights_arr
        return etf_series

    def _index_check(self):
        """
        Internal check for all price, rates and allocations data frames have the same index
        """
        # Check if all data frames have the same index
        for temp_df in self.data_dict.values():
            if self.data_dict['open'].index.difference(temp_df.index).shape[0] != 0 or \
                    self.data_dict['open'].shape != temp_df.shape:
                raise ValueError('DataFrames indices are different')

    def _get_batch_from_csv(self, batch_size):
        """
        Reads the next batch of data sets from csv files and puts them in class variable data_dict
        :param batch_size: number of rows to read
        """
        # Read the next batch
        self.data_dict['open'] = self.iter_dict['open'].get_chunk(batch_size)
        self.data_dict['close'] = self.iter_dict['close'].get_chunk(batch_size)
        self.data_dict['alloc'] = self.iter_dict['alloc'].get_chunk(batch_size)
        self.data_dict['costs'] = self.iter_dict['costs'].get_chunk(batch_size)

        if self.iter_dict['rates'] is not None:
            # If there is rates_df iterator, get the next chunk
            self.data_dict['rates'] = self.iter_dict['rates'].get_chunk(batch_size)
        else:
            # If no iterator is available, generate trivial rates_df (1.0 for all securities)
            self.data_dict['rates'] = self.data_dict['open'].copy()
            # Set trivial(1.0) exchange rate if no data is provided
            self.data_dict['rates'][self.securities] = 1.0

        # Align all securities columns in one order
        for df_name in self.data_dict:
            self.data_dict[df_name] = self.data_dict[df_name][self.securities]

        self._index_check()

    def _rewind_etf_trick(self, alloc_df, etf_series):
        """
        ETF trick uses next open price information, when we process csv file in batches the last row in batch will have
        next open price value as nan, that is why when new batch comes, we need to rewind ETF trick values
        one step back, recalculate ETF trick value for the last row from previous batch using open price from latest
        batch received. This function rewinds values needed for ETF trick calculation
        recalculate

        :param alloc_df: (pd.DataFrame): data frame with allocations vectors
        :param etf_series (pd.Series): current computed ETF trick series
        :return:
        """
        self.prev_k = etf_series.iloc[-2]  # Reset prev_k for previous row calculation
        self.prev_allocs = alloc_df.iloc[-2]  # Reset previous allocations vector
        self.prev_allocs_change = bool(~(self.prev_allocs == alloc_df.iloc[-3]).all())

    def _csv_file_etf_series(self, batch_size):
        """
        Csv based ETF trick series generation
        :return: (pd.Series): pandas Series with ETF trick values starting from 1.0
        """
        etf_series = pd.Series()
        self._get_batch_from_csv(batch_size)
        # Data frame which contains all precomputed info for etf trick
        data_df = self.generate_trick_components(cache=None)  # Cache is empty on the first batch run
        cache = self._update_cache()
        # Delete first nans (first row of close price difference is nan)
        data_df = data_df.iloc[1:]
        omit_last_row = False  # Drop last row value from previous batch (this row needs to be recalculated using new data)

        # Read data in batch until StopIteration exception is raised
        while True:
            try:
                chunk_etf_series = self._chunk_loop(data_df)
                if omit_last_row is True:
                    etf_series = etf_series.iloc[:-1]  # Delete last row (chunk_etf_series stores updated row value)
                etf_series = etf_series.append(chunk_etf_series)
                self._get_batch_from_csv(batch_size)
                self._rewind_etf_trick(data_df['w'], etf_series)  # Rewind etf series one step back
                data_df = self.generate_trick_components(cache)  # Update data_df for ETF trick calculation
                cache = self._update_cache()  # Update cache
                omit_last_row = True
            except StopIteration:
                return etf_series

    def _in_memory_etf_series(self):
        """
        In-memory based ETF trick series generation.

        :return: (pd.Series): pandas Series with ETF trick values starting from 1.0
        """
        data_df = self.generate_trick_components()  # Data frame which contains all precomputed info for etf trick

        # Delete first nans (first row of close price difference is nan)
        data_df = data_df.iloc[1:]
        return self._chunk_loop(data_df)

    def get_etf_series(self, batch_size=1e5):
        """
        External method which defines which etf trick method to use.

        :return: (pd.Series): pandas Series with ETF trick values starting from 1.0
        """
        if self.iter_dict is None:
            etf_trick_series = self._in_memory_etf_series()
        else:
            if batch_size < 3:
                # We use latest 2 elements from prev batch, so minimum batch is 3
                raise ValueError('Batch size should be >= 3')
            etf_trick_series = self._csv_file_etf_series(batch_size)
        return etf_trick_series

    def reset(self):
        """
        Reinits class object. This methods can be used to reset file iterators for multiple get_etf_trick() calls.
        """
        self.__init__(**self.init_fields)


def get_futures_roll_series(data_df, open_col, close_col, sec_col, current_sec_col, roll_backward=False):
    """
    Function for generating rolling futures series from data frame of multiple futures.

    :param data_df: (pd.DataFrame): pandas DataFrame containing price info, security name and current active futures column
    :param open_col: (string): open prices column name
    :param close_col: (string): close prices column name
    :param sec_col: (string): security name column name
    :param current_sec_col: (string): current active security column name. When value in this column changes it means rolling
    :param roll_backward: (boolean): True for subtracting final gap value from all values
    :return (pd.Series): futures roll close price series
    """
    # Filter out security data which is not used as current security
    filtered_df = data_df[data_df[sec_col] == data_df[current_sec_col]]
    filtered_df.sort_index(inplace=True)

    # Generate roll dates series based on current_sec column value change
    roll_dates = filtered_df[current_sec_col].drop_duplicates(keep='first').index
    gaps = filtered_df[close_col] * 0  # roll gaps series

    # On roll dates, gap equals open - close
    gaps.loc[roll_dates[1:]] = filtered_df[open_col].loc[roll_dates[1:]] - filtered_df[close_col].loc[
        roll_dates[1:]]
    gaps = gaps.cumsum()

    if roll_backward:
        gaps -= gaps.iloc[-1]  # Roll backward

    return gaps
