"""
Module for decomposing the prediction ANN into piecewise linear functions
via Taylor Series decomposition
"""

import neuro_mine.model as model
import numpy as np
import tensorflow as tf
from numba import njit
from typing import List, Tuple


# def _dCa_dR_relu(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> np.ndarray:
#     """
#     For the output unit of the ANN with relu activations computes the vector of first derivatives
#     with respect to each regressor input
#     :param mdl: The model for which to compute the derivative
#     :param reg_input: The input at which to compute the derivative
#     :returns: (input_length x nRegressors) long vector of partial first derivatives
#     """
#     # For notes see tanh implementation. The only difference here is the calculation of the outer derivatives
#     # in the chain rule: This derivative for ReLu is 0 whenever the activation is < 0, 1 whenever the activation
#     # is > 0 and undefined at 0 - to avoid NaN's we do what everyone incorrectly does and set it to 0 at 0
#
#     # CONVOLUTIONAL LAYER #
#     c = mdl.conv_layer(reg_input)  # (1 x 1 x N_c)
#     c = mdl.flatten(c).numpy()  # (1 x N_c)
#     weights_r_c = mdl.conv_layer.weights[0].numpy()  # (input_length x n_channels x N_c)
#     # linearize convolutional weights with respect to channel dimension so we later don't have
#     # to deal with matrices of partial derivatives but rather a big vector where groups of input_length
#     # elements belong to one given regressor
#     weights_r_c = weights_r_c.reshape((weights_r_c.shape[0]*weights_r_c.shape[1], weights_r_c.shape[2]))
#     dc_dr = weights_r_c  # (N_r x N_c)
#
#     # FIRST DEEP LAYER #
#     x = mdl.deep_1(c).numpy()
#     weights_c_x = mdl.deep_1.weights[0].numpy()  # (N_c x N_x)
#     d_relu_x = np.diag((x.ravel() > 0).astype(np.float32))
#     dx_dr = np.matmul(dc_dr, np.matmul(d_relu_x, weights_c_x.T).T)  # (N_r x N_x)
#
#     # SECOND DEEP LAYER #
#     y = mdl.deep_2(x).numpy()
#     weights_x_y = mdl.deep_2.weights[0].numpy()  # (N_x x N_y)
#     d_relu_y = np.diag((y.ravel() > 0).astype(np.float32))
#     dy_dr = np.matmul(dx_dr, np.matmul(d_relu_y, weights_x_y.T).T)  # (N_r x N_y)
#
#     # OUTPUT NODE #
#     weights_y_z = mdl.out.weights[0].numpy()  # (N_y x 1)
#     dz_dr = np.matmul(dy_dr, weights_y_z)  # (N_r x 1)
#     return dz_dr.ravel()
#
#
# def _ddCa_dRdR_relu(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
#     """
#     For the output unit of the ANN with relu activations computes the matrix of second derivatives
#     with respect to each regressor input
#     :param mdl: The model for which to compute the derivative
#     :param reg_input: The input at which to compute the derivative
#     :returns:
#         [0] (input_length x nRegressors) long vector of partial first derivatives
#         [1] (input_length x nRegressors) x (input_length x nRegressors) matrix of partial second derivatives
#     """
#     # With relu activation all second partial derivatives are 0 (or technically undefined for
#     # inputs into a node = 0)
#     return _dCa_dR_relu(mdl, reg_input), np.zeros((reg_input.size, reg_input.size), np.float32)
#
#
# def _dCa_dR_tanh(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> np.ndarray:
#     """
#     For the output unit of the ANN with tanh activations computes the vector of first derivatives
#     with respect to each regressor input
#     :param mdl: The model for which to compute the derivative
#     :param reg_input: The input at which to compute the derivative
#     :returns: (input_length x nRegressors) long vector of partial first derivatives
#     """
#     # r: Regressor inputs
#     # c: Convolutional layer output (Note: This layer is linear)
#     # x: First deep layer output
#     # y: Second deep layer output
#     # z: Output (one node, linear)
#     # N_r: Total length of regressor input = =n_regressors x input_length
#     # N_c: Number of convolutional nodes
#     # N_x: Number of deep layer 1 nodes
#     # N_y: Number of deep layer 2 nodes
#     # weigths_r_c: weights from layer R to layer C, etc.
#     # dc_dr: derivative of convolutional layer units with respect to regressor inputs, etc.
#     # shape notes the number of rows followed by number of columns in brackets
#
#     # CONVOLUTIONAL LAYER #
#     c = mdl.conv_layer(reg_input)  # (1 x 1 x N_c)
#     c = mdl.flatten(c).numpy()  # (1 x N_c)
#     weights_r_c = mdl.conv_layer.weights[0].numpy()  # (input_length x n_channels x N_c)
#     # linearize convolutional weights with respect to channel dimension so we later don't have
#     # to deal with matrices of partial derivatives but rather a big vector where groups of n_channels
#     # elements belong to one given timepoint
#     weights_r_c = weights_r_c.reshape((weights_r_c.shape[0]*weights_r_c.shape[1], weights_r_c.shape[2]))
#     dc_dr = weights_r_c  # (N_r x N_c)
#
#     # FIRST DEEP LAYER #
#     x = mdl.deep_1(c).numpy()  # already has nonlinearity applied!!!
#     # this is the activation before the nonlinearity
#     xa = tf.add(tf.matmul(c, mdl.deep_1.weights[0]), mdl.deep_1.weights[1]).numpy()
#     weights_c_x = mdl.deep_1.weights[0].numpy()  # (N_c x N_x)
#     d_tanh_x = np.diag(((1 / np.cosh(xa)) ** 2).ravel()) since dtanh/dx = sech2(x) and sech(x) = 1/cosh(x) (N_x x N_x)
#     dx_dr = np.matmul(dc_dr, np.matmul(d_tanh_x, weights_c_x.T).T)  # (N_r x N_x)
#
#     # SECOND DEEP LAYER #
#     ya = tf.add(tf.matmul(x, mdl.deep_2.weights[0]), mdl.deep_2.weights[1]).numpy()
#     weights_x_y = mdl.deep_2.weights[0].numpy()  # (N_x x N_y)
#     d_tanh_y = np.diag(((1 / np.cosh(ya)) ** 2).ravel())  # (N_y x N_y)
#     dy_dr = np.matmul(dx_dr, np.matmul(d_tanh_y, weights_x_y.T).T)  # (N_r x N_y)
#
#     # OUTPUT NODE #
#     weights_y_z = mdl.out.weights[0].numpy()  # (N_y x 1)
#     dz_dr = np.matmul(dy_dr, weights_y_z)  # (N_r x 1)
#     return dz_dr.ravel()
#
#
# @njit
# def _compute_part_tensor(projection: np.ndarray, n: int):
#     """
#     Computes intermediate tensor for second derivative computations
#     :param projection: Projection of previous layer derivative by current layer weights
#     :param n: Total number of input elements (n_timepoints x n_regressors)
#     :return: The intermediate tensor for computation
#     """
#     part_tensor = np.empty((n, projection.shape[1], n), dtype=np.float32)
#     for r1 in range(n):
#         for r2 in range(r1, n):
#             val = projection[r1, :] * projection[r2, :]
#             part_tensor[r1, :, r2] = val
#             part_tensor[r2, :, r1] = val
#     return part_tensor
#
#
# def _ddCa_dRdR_tanh(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
#     """
#     For the output unit of the ANN with tanh activations computes the matrix of second and vector of
#     first derivatives with respect to each regressor input
#     :param mdl: The model for which to compute the derivative
#     :param reg_input: The input at which to compute the derivative
#     :returns:
#         [0] (input_length x nRegressors) long vector of partial first derivatives
#         [1] (input_length x nRegressors) x (input_length x nRegressors) matrix of partial second derivatives
#     """
#     # CONVOLUTIONAL LAYER #
#     c = mdl.conv_layer(reg_input)  # (1 x 1 x N_c)
#     c = mdl.flatten(c).numpy()  # (1 x N_c)
#     weights_r_c = mdl.conv_layer.weights[0].numpy()  # (input_length x n_channels x N_c)
#     # linearize convolutional weights with respect to channel dimension so we later don't have
#     # to deal with matrices of partial derivatives but rather a big vector where groups of N-regressor
#     # elements belong to one given timepoint
#     weights_r_c = weights_r_c.reshape((weights_r_c.shape[0]*weights_r_c.shape[1], weights_r_c.shape[2]))
#     dc_dr = weights_r_c  # (N_r x N_c)
#     # second derivative is 0 in a linear layer
#     # ddc_drr = np.zeros((weights_r_c.shape[0], c.size, weights_r_c.shape[0]), dtype=np.float32)  # (N_r x N_c x N_r)
#
#     # FIRST DEEP LAYER #
#     # the following is the activation before the nonlinearity
#     weights_c_x = mdl.deep_1.weights[0].numpy()  # (N_c x N_x)
#     bias_x = mdl.deep_1.weights[1].numpy()
#     xa = c@weights_c_x + bias_x  # activation before the nonlinearity
#     x = np.tanh(xa)  # output of the neurons
#     d_tanh_x = np.diag((1 - x**2).ravel())  # since dtanh/dx = sech2(x) & sech2(x) = 1-tanh2(x) (N_x x N_x)
#     dd_tanh_x = -2 * d_tanh_x * np.diag(x.ravel())  # since ddtanh/dx = -2 * dtanh/dx * tanh (N_x x N_x)
#     dx_dr = dc_dr@(d_tanh_x@weights_c_x.T).T  # (N_r x N_x)
#     dc_dr_projection = dc_dr@weights_c_x  # (N_r x N_x)
#     # NOTE: The following part_tensor is independent of the input, in other words it can be calculated once per model
#     if mdl.part_tensor_1 is not None:
#         part_tensor_1 = mdl.part_tensor_1
#     else:
#         part_tensor_1 = _compute_part_tensor(dc_dr_projection, weights_r_c.shape[0])
#         mdl.part_tensor_1 = part_tensor_1
#     # Note: for symmetry we could compute the first term here, but it is by definition 0
#     # ddx_drr = np.matmul(d_tanh_x, np.matmul(weights_c_x.T, ddc_drr)) + np.matmul(dd_tanh_x, part_tensor)
#     ddx_drr = dd_tanh_x@part_tensor_1
#
#     # SECOND DEEP LAYER #
#     weights_x_y = mdl.deep_2.weights[0].numpy()  # (N_x x N_y)
#     bias_y = mdl.deep_2.weights[1].numpy()
#     ya = x@weights_x_y + bias_y
#     y = np.tanh(ya)
#     d_tanh_y = np.diag((1 - y**2).ravel())  # (N_y x N_y)
#     dd_tanh_y = -2 * d_tanh_y * np.diag(y.ravel())  # (N_x x N_x)
#     dy_dr = dx_dr@(d_tanh_y@weights_x_y.T).T  # (N_r x N_y)
#     dx_dr_projection = dx_dr@weights_x_y  # (N_r x N_y)
#     part_tensor_2 = _compute_part_tensor(dx_dr_projection, weights_r_c.shape[0])
#     ddy_drr = d_tanh_y@(weights_x_y.T@ddx_drr) + dd_tanh_y@part_tensor_2
#
#     # OUTPUT NODE #
#     weights_y_z = mdl.out.weights[0].numpy()  # (N_y x 1)
#     dz_dr = dy_dr@weights_y_z  # (N_r x 1)
#     ddz_drr = weights_y_z.T@ddy_drr  # (N_r x 1 x N_r)
#
#     return dz_dr.ravel(), ddz_drr[:, 0, :]


@tf.function
def d2ca_dr2(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> Tuple[tf.Tensor, tf.Tensor]:
    x = tf.convert_to_tensor(reg_input)
    with tf.GradientTape() as t2:
        t2.watch(x)
        with tf.GradientTape() as t1:
            t1.watch(x)
            # NOTE: The following is slightly faster than ca = mdl(x) presumably due to skipping of dropout layers
            # However it is not compatible with resizing the model!!
            # c = mdl.conv_layer(x)
            # d1 = mdl.deep_1(mdl.flatten(c))
            # d2 = mdl.deep_2(d1)
            # ca = mdl.out(d2)
            ca = mdl(x)
        jacobian = t1.gradient(ca, x)
    hessian = t2.jacobian(jacobian, x)
    return jacobian, hessian


@tf.function
def dca_dr(mdl: model.ActivityPredictor, reg_input: np.ndarray) -> tf.Tensor:
    x = tf.convert_to_tensor(reg_input)
    with tf.GradientTape() as t1:
        t1.watch(x)
        # NOTE: The following is slightly faster than ca = mdl(x) presumably due to skipping of dropout layers
        # but not compatible with change layer number in the model
        # c = mdl.conv_layer(x)
        # d1 = mdl.deep_1(mdl.flatten(c))
        # d2 = mdl.deep_2(d1)
        # ca = mdl.out(d2)
        ca = mdl(x)
    jacobian = t1.gradient(ca, x)
    return jacobian


def taylor_predict(mdl: model.ActivityPredictor, regressors: np.ndarray, use_d2: bool, take_every: int,
                   predict_ahead=1) -> Tuple[np.ndarray, np.ndarray]:
    """
    For each time t in regressors, evaluates the model, computes the selected derivatives and
    then attempts to predict the model response at time t+1
    :param mdl: The model to use for predictions
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param use_d2: If set to False only the first derivative will be used for the prediction
    :param take_every: Only form predictions every n frames to save time
    :param predict_ahead: The number of frames to predict ahead with the taylor expansion
    :returns:
        [0]: (n_timesteps-input_length-predict_ahead)/n long timeseries of taylor predictions
        [1]: (n_timesteps-input_length-predict_ahead)/n long timeseries of actual network outputs
    """
    if predict_ahead < 1:
        raise ValueError("predict_ahead has to be integer >= 1")
    inp_length = mdl.input_length
    mdl_output = []
    taylor_prediction = []
    t = inp_length - 1
    while t < regressors.shape[0]-predict_ahead:
        cur_regs = regressors[None, t - inp_length + 1: t + 1, :]
        next_regs = regressors[None, t - inp_length + 1 + predict_ahead: t + 1 + predict_ahead, :]
        cur_mod_out = mdl.get_output(cur_regs)
        next_mod_out = mdl.get_output(next_regs)
        mdl_output.append(next_mod_out)
        if use_d2:
            d1, d2 = d2ca_dr2(mdl, cur_regs)
            d1 = d1.numpy().ravel()
            d2 = np.reshape(d2.numpy(), (regressors.shape[1] * mdl.input_length,
                                         regressors.shape[1] * mdl.input_length))
        else:
            d1 = dca_dr(mdl, cur_regs)
            d1 = d1.numpy().ravel()
            d2 = None
        tay_pred = _taylor_predict(cur_regs, next_regs, cur_mod_out, d1, d2)
        taylor_prediction.append(tay_pred)
        t += take_every
    return np.hstack(taylor_prediction), np.hstack(mdl_output)


def _taylor_predict(reg_fix_point: np.ndarray, reg_test: np.ndarray, ann_fix: float, d1: np.ndarray,
                    d2: np.ndarray) -> float:
    """
    Computes the taylor prediction about a point for another test point nearby
    :param reg_fix_point: The regressor input at the fix point where derivatives have been calculated
    :param reg_test: The regressor input for which to predict the ann response
    :param ann_fix: The output of the ann at reg_fix_point
    :param d1: The set of first order partial derivatives at reg_fix_point
    :param d2: The matrix of second order partial derivatives at reg_fix_point
    """
    diff = (reg_test - reg_fix_point).ravel()
    if d2 is None:
        return ann_fix + np.dot(diff, d1)
    return ann_fix + np.dot(diff, d1) + 0.5 * np.sum(np.dot(diff[:, None], diff[None, :]) * d2)


@njit
def _compute_taylor_d2(reg_diff: np.ndarray, d2: np.ndarray, nregs: int, inp_length: int) -> np.ndarray:
    """
    Computes responses belonging to the second derivative, rearranging by regressor
    instead of by time
    :param reg_diff: The difference in regressors as 2D (1 x (n_regs*n_timepoints)) vector
    :param d2: The hessian
    :param nregs: The number of regressors
    :param inp_length: The timelength of each regressor input
    :return: The second derivative contribution ((n_regs*n_timepoints) x (n_regs*n_timepoints))
    """
    taylor_d2_temp = 0.5 * np.dot(reg_diff, reg_diff.T) * d2  # this matrix is organized by time not by regressor
    taylor_d2 = np.empty_like(taylor_d2_temp, dtype=np.float32)
    for row in range(taylor_d2_temp.shape[0]):
        regnum = row % nregs
        time = row // nregs
        row_ix = regnum * inp_length + time
        for col in range(taylor_d2_temp.shape[1]):
            regnum = col % nregs
            time = col // nregs
            col_ix = regnum * inp_length + time
            taylor_d2[row_ix, col_ix] = taylor_d2_temp[row, col]
    return taylor_d2


@njit
def _compute_by_reg(taylor_d1: np.ndarray, taylor_d2: np.ndarray, nregs: int, inp_length: int) -> np.ndarray:
    """
    Aggregates derivative contributions by regressor
    :param taylor_d1: The first partial derivative contributions (by regressor and time vector)
    :param taylor_d2: The second partial derivative contributions (by regressor and time square matrix)
    :param nregs: The number of regressors
    :param inp_length: The number of timepoints
    :return: Contribution aggregated by regressor as array (1 xnregs x nregs) to account for possible interactions
    """
    by_reg = np.full((1, nregs, nregs), 0.0, dtype=np.float32)
    for r1 in range(nregs):
        for r2 in range(nregs):
            if r1 == r2:
                # these are the non-interacting parts which need to take d1 into account
                by_reg[0, r1, r2] += np.sum(taylor_d1[r1 * inp_length:(r1 + 1) * inp_length])
            by_reg[0, r1, r2] += np.sum(
                taylor_d2[r1 * inp_length:(r1 + 1) * inp_length, r2 * inp_length:(r2 + 1) * inp_length])
    return by_reg


def taylor_decompose(mdl: model.ActivityPredictor, regressors: np.ndarray, take_every: int, predict_ahead: int,
                     use_d2=True) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Uses taylor decomposition to predict changes in network output around chosen point using the
    all information as well as only the information corresponding to each regressor and their
    interactions terms
    :param mdl: The model to use for predictions
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param take_every: Only form predictions every n frames to save time
    :param predict_ahead: The number of frames to predict ahead with the taylor expansion
    :param use_d2: If set to false only the first derivative will be used in the taylor expansion
    :returns:
        [0]: The true change for each timepoint by going predict_ahead frames into the future
                (n_timesteps-input_length-predict_ahead)/n long vector
        [1]: The predicted change for the whole taylor series
                (n_timesteps-input_length-predict_ahead)/n long vector
        [2]: Array of predicted changes by regressors and their interactions
                (n_timesteps-input_length-predict_ahead)/n x n_regressors x n_regressors
    """
    if predict_ahead < 1:
        raise ValueError("predict_ahead has to be integer >= 1")
    inp_length = mdl.input_length
    mdl_out_change = []
    full_tp_change = []
    by_reg_tp_change = []
    t = inp_length - 1
    nregs = regressors.shape[1]
    while t < regressors.shape[0]-predict_ahead:
        cur_regs = regressors[None, t - inp_length + 1: t + 1, :]
        next_regs = regressors[None, t - inp_length + 1 + predict_ahead: t + 1 + predict_ahead, :]
        reg_diff = (next_regs - cur_regs).ravel().astype(np.float64)
        cur_mod_out = mdl.get_output(cur_regs)
        next_mod_out = mdl.get_output(next_regs)
        mdl_out_change.append(next_mod_out - cur_mod_out)
        # unfortunately, none of the following are contiguous by regressor but rather by timepoint
        # we therefore need to reshape taylor_d1 and taylor_d2 to allow for easy indexing below
        if use_d2:
            d1, d2 = d2ca_dr2(mdl, cur_regs)
            d1 = d1.numpy().ravel()
            d2 = np.reshape(d2.numpy(), (regressors.shape[1] * mdl.input_length,
                                         regressors.shape[1] * mdl.input_length))
            taylor_d1 = reg_diff * d1.astype(np.float64)
            taylor_d1 = np.reshape(taylor_d1, (inp_length, nregs)).T.ravel()
            taylor_d2 = _compute_taylor_d2(reg_diff[:, None], d2.astype(np.float64), nregs, inp_length)
        else:
            d1 = dca_dr(mdl, cur_regs)
            d1 = d1.numpy().ravel()
            taylor_d1 = reg_diff * d1.astype(np.float64)
            taylor_d1 = np.reshape(taylor_d1, (inp_length, nregs)).T.ravel()
            taylor_d2 = np.zeros((taylor_d1.size, taylor_d1.size))
        full_tp_change.append(np.sum(taylor_d1) + np.sum(taylor_d2))
        by_reg = _compute_by_reg(taylor_d1, taylor_d2, nregs, inp_length)
        by_reg_tp_change.append(by_reg)
        t += take_every
    return np.hstack(mdl_out_change), np.hstack(full_tp_change), np.vstack(by_reg_tp_change)


def avg_directional_curvature(mdl: model.ActivityPredictor, regressors: np.ndarray, delta_distance: int,
                              take_every: int) -> Tuple[float, List[float]]:
    """
    Along a trajectory through input space defined by regressors computes the average
    curvature of the function approximated by the ANN
    :param mdl: The model to use for predictions
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param delta_distance: How many timesteps to look ahead to compute direction in input space
    :param take_every: Only curvature every n frames to save time
    :return:
        [0]: Scalar, the average curvature along the path
        [1]: List of all point-wise curvatures
    """
    if delta_distance < 1:
        raise ValueError("delta_distance has to be integer >= 1")
    inp_length = mdl.input_length
    t = inp_length - 1
    all_curvatures = []
    while t < regressors.shape[0]-delta_distance:
        cur_regs = regressors[None, t - inp_length + 1: t + 1, :]
        next_regs = regressors[None, t - inp_length + 1 + delta_distance: t + 1 + delta_distance, :]
        d2 = d2ca_dr2(mdl, cur_regs)[1]
        if d2 is None:
            d2 = np.zeros((regressors.shape[1] * mdl.input_length,
                           regressors.shape[1] * mdl.input_length))
        else:
            d2 = np.reshape(d2.numpy(), (regressors.shape[1] * mdl.input_length,
                                         regressors.shape[1] * mdl.input_length))
        dir_curv = _directional_curvature(cur_regs, next_regs, d2)
        all_curvatures.append(dir_curv)
        t += take_every
    avg_curve = float(np.mean(all_curvatures))
    return avg_curve, all_curvatures


def _directional_curvature(reg_fix_point: np.ndarray, reg_ahead: np.ndarray, d2: np.ndarray) -> float:
    """
    Computes the local curvature in direction of movement in the input space
    from reg_fix_point to reg_ahead
    :param reg_fix_point: The regressor input at the current timepoint
    :param reg_ahead: The regressor input at the next point on the path through input space
    :param d2: The hessian at the current timepoint
    :return: The absolute curvature at fix-point in path direction = ||(reg_ahead-reg_fix_point)d2||2
    """
    diff = (reg_ahead - reg_fix_point).ravel()[None, :]
    nm = np.linalg.norm(diff)
    if nm == 0:
        nm = 10
    return np.linalg.norm(np.dot(diff, d2)) / nm


def all_decomposition_metrics(mdl: model.ActivityPredictor, regressors: np.ndarray, take_every: int,
                              look_ahead: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float, np.ndarray]:
    """
    Computes both the taylor decomposition and non-linearity metric to save time by calculating the needed
    derivatives only once
    :param mdl: The model to use for predictions
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param take_every: Only compute metrics every n frames to save time
    :param look_ahead: How many timesteps to look ahead to compute direction in input space
    :return:
        [0]: The true change for each timepoint by going predict_ahead frames into the future
                (n_timesteps-input_length-predict_ahead)/n long vector
        [1]: The predicted change for the whole taylor series
                (n_timesteps-input_length-predict_ahead)/n long vector
        [2]: Array of predicted changes by regressors and their interactions
                (n_timesteps-input_length-predict_ahead)/n x n_regressors x n_regressors
        [3]: Scalar, the average non-linearity metric along the path
        [4]: Array of point-wise non-linearity metrics
    """
    if look_ahead < 1:
        raise ValueError("look_ahead has to be integer >= 1")
    inp_length = mdl.input_length
    mdl_out_change: List[float] = []
    full_tp_change: List[float] = []
    by_reg_tp_change: List[np.ndarray] = []
    all_curvatures: List[float] = []
    t: int = inp_length - 1
    nregs: int = regressors.shape[1]
    while t < regressors.shape[0]-look_ahead:
        cur_regs = regressors[None, t - inp_length + 1: t + 1, :]
        next_regs = regressors[None, t - inp_length + 1 + look_ahead: t + 1 + look_ahead, :]
        reg_diff = (next_regs - cur_regs).ravel().astype(np.float64)
        cur_mod_out = mdl.get_output(cur_regs)
        next_mod_out = mdl.get_output(next_regs)
        mdl_out_change.append(next_mod_out - cur_mod_out)
        d1, d2 = d2ca_dr2(mdl, cur_regs)
        d1 = d1.numpy().ravel()
        d2 = np.reshape(d2.numpy(), (regressors.shape[1] * mdl.input_length,
                                     regressors.shape[1] * mdl.input_length))
        # compute non-linearity metric
        dir_curv = _directional_curvature(cur_regs, next_regs, d2)
        all_curvatures.append(dir_curv)
        # compute taylor decomposition - note, since the derivatives are not orderd by regressor but by time
        # they need to be re-arranged
        taylor_d1 = reg_diff * d1.astype(np.float64)
        taylor_d1 = np.reshape(taylor_d1, (inp_length, nregs)).T.ravel()
        taylor_d2 = _compute_taylor_d2(reg_diff[:, None], d2.astype(np.float64), nregs, inp_length)
        full_tp_change.append(np.sum(taylor_d1) + np.sum(taylor_d2))
        by_reg = _compute_by_reg(taylor_d1, taylor_d2, nregs, inp_length)
        by_reg_tp_change.append(by_reg)
        t += take_every
    avg_curve = float(np.mean(all_curvatures))
    return np.hstack(mdl_out_change), np.hstack(full_tp_change), np.vstack(by_reg_tp_change), avg_curve,\
        np.hstack(all_curvatures)


def data_mean_prediction(mdl: model.ActivityPredictor, x_bar, j_x_bar, h_x_bar, regressors: np.ndarray, take_every: int,
                         return_1storder=False):
    """
    Computes the prediction of responses based on a fixed Taylor expansion of the network around a specific point
    in our case taken to be the data mean
    :param mdl: The CNN model
    :param x_bar: The data mean (or any arbitrary fix point)
    :param j_x_bar: The jacobian of the model at x_bar
    :param h_x_bar: The hession of the model at x_bar
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param take_every: Only compute metrics every n frames to save time
    :param return_1storder: If set to true first order model predictions will be returned as well
    :return:
        [0]: The prediction of the CNN model
        [1]: The prediction of the fixed-point expansion
        [2]: If return_1storder is set to true the prediction of a linear fixed-point expansion
    """
    # NOTE: We should speed this up - create 2nd-order design matrix from input and then directly treat this as a
    # regression problem, replacing the ugly while-loop with a matrix multiplication
    f_x_bar = mdl(x_bar)
    mean_prediction: List[float] = []
    mean_prediction_lin: List[float] = []
    mdl_out: List[float] = []
    inp_length = mdl.input_length
    t: int = inp_length - 1
    # prepare our first and second derivatives at the data mean
    d1 = j_x_bar.numpy().ravel()
    d2 = np.reshape(h_x_bar.numpy(), (regressors.shape[1] * mdl.input_length, regressors.shape[1] * mdl.input_length))
    while t < regressors.shape[0]:
        cur_regs = regressors[None, t - inp_length + 1: t + 1, :]
        # compute our difference to the data mean
        reg_diff = (cur_regs - x_bar).ravel().astype(np.float64)
        # get the actual model prediction at the current point and add to our return
        cur_mod_out = mdl.get_output(cur_regs)
        mdl_out.append(cur_mod_out)
        # compute taylor decomposition around data mean
        mp_lin = f_x_bar + np.dot(reg_diff, d1)
        mp = mp_lin + 0.5 * np.sum(np.dot(reg_diff[:, None], reg_diff[None, :]) * d2)
        mean_prediction.append(mp)
        if return_1storder:
            mean_prediction_lin.append(mp_lin)
        t += take_every
    if return_1storder:
        return np.hstack(mdl_out), np.hstack(mean_prediction), np.hstack(mean_prediction_lin)
    return np.hstack(mdl_out), np.hstack(mean_prediction)


def complexity_scores(mdl: model.ActivityPredictor, x_bar, j_x_bar, h_x_bar, regressors: np.ndarray, take_every: int):
    """
    Computes complexity scores - the squared correlation of a linear and a squared model around the data mean
    :param mdl: The CNN model
    :param x_bar: The data mean (or any arbitrary fix point)
    :param j_x_bar: The jacobian of the model at x_bar
    :param h_x_bar: The hession of the model at x_bar
    :param regressors: The 2D regressor matrix, n_timesteps x m_regressors
    :param take_every: Only compute metrics every n frames to save time
    :return:
        [0]: The R2 (coefficient of determination) of the linear 1st order approximation
        [1]: The R2 of the 2nd order approximation
    """
    true_model, order_2, order_1 = data_mean_prediction(mdl, x_bar, j_x_bar, h_x_bar, regressors, take_every, True)
    ss_tot = np.sum((true_model - np.mean(true_model))**2)
    lin_score = 1 - np.sum((true_model - order_1)**2)/ss_tot
    sq_score = 1 - np.sum((true_model - order_2)**2)/ss_tot
    return lin_score, sq_score


if __name__ == "__main__":
    print("Module for ANN Taylor decomposition")
