# Path: ceviche/utils.py
# def make_sparse(entries, indices, shape):
# def get_entries_indices(csr_matrix):
# def transpose_indices(indices):
# def block_4(A, B, C, D):
# def make_IO_matrices(indices, N):
# def make_rand(N):
# def make_rand_complex(N):
# def make_rand_indeces(N, M):
# def make_rand_entries_indices(N, M):
# def make_rand_sparse(N, M):
# def make_rand_sparse_density(N, density=1):
# def der_num(fn, arg, index, delta):
# def grad_num(fn, arg, delta=1e-6):
# def jac_num(fn, arg, step_size=1e-7):
# def grid_center_to_xyz(Q_mid, averaging=True):
# def grid_xyz_to_center(Q_xx, Q_yy, Q_zz):
# def vec_zz_to_xy(info_dict, vec_zz, grid_averaging=True):
# def float_2_array(x):
# def reshape_to_ND(arr, N):
# def get_value(x):
# def get_shape(x):
# def vjp_maker_num(fn, arg_inds, steps):
#     def vjp_single_arg(ia):
#         def vjp_maker(fn_out, *args):
#             def vjp(v):
# def aniplot(F, source, steps, component='Ez', num_panels=10):
# def measure_fields(F, source, steps, probes, component='Ez'):
# def imarr(arr):
# def my_fft(x):    
# def fft_grad(g, ans, x):
# def get_spectrum(series, dt):
# def get_max_power_freq(series, dt):
# def get_spectral_power(series, dt):
# def plot_spectral_power(series, dt, f_top=2e14):
#     M = indices.shape[1]                                 # number of indices in the matrix
#     I = make_sparse(entries_1, indices_I, shape=(N, M))  # construct the I matrix
#     J = make_sparse(entries_1, indices_J, shape=(N, M))  # construct the J matrix
#     O = J.T                                              # make O = J^T matrix for consistency with my notes.
#     N = arg.size
#     ND = len(arr.shape)


import autograd.numpy as npa
from autograd.core import make_vjp, make_jvp
from autograd.wrap_util import unary_to_nary
from autograd.extend import vspace
from .utils import get_value, get_shape, get_value_arr, float_2_array

@unary_to_nary
def jacobian_reverse(fun, x):
    """ Compute jacobian of fun with respect to x using reverse mode differentiation"""
    vjp, ans = make_vjp(fun, x)
    grads = map(vjp, vspace(ans).standard_basis())
    m, n = _jac_shape(x, ans)
    return npa.reshape(npa.stack(grads), (n, m))


@unary_to_nary
def jacobian_forward(fun, x):
    """ Compute jacobian of fun with respect to x using forward mode differentiation"""
    jvp = make_jvp(fun, x)
    # ans = fun(x)
    val_grad = map(lambda b: jvp(b), vspace(x).standard_basis())
    vals, grads = zip(*val_grad)
    ans = npa.zeros((list(vals)[0].size,))  # fake answer so that dont have to compute it twice
    m, n = _jac_shape(x, ans)
    if _iscomplex(x):
        grads_real = npa.array(grads[::2])
        grads_imag = npa.array(grads[1::2])
        grads = grads_real - 1j * grads_imag
    return npa.reshape(npa.stack(grads), (m, n)).T


@unary_to_nary
def jacobian_numerical(fn, x, step_size=1e-7):
    """ numerically differentiate `fn` w.r.t. its argument `x` """
    in_array = float_2_array(x).flatten()

# Path: ceviche/utils.py
# def make_sparse(entries, indices, shape):
# def get_entries_indices(csr_matrix):
# def transpose_indices(indices):
# def block_4(A, B, C, D):
# def make_IO_matrices(indices, N):
# def make_rand(N):
# def make_rand_complex(N):
# def make_rand_indeces(N, M):
# def make_rand_entries_indices(N, M):
# def make_rand_sparse(N, M):
# def make_rand_sparse_density(N, density=1):
# def der_num(fn, arg, index, delta):
# def grad_num(fn, arg, delta=1e-6):
# def jac_num(fn, arg, step_size=1e-7):
# def grid_center_to_xyz(Q_mid, averaging=True):
# def grid_xyz_to_center(Q_xx, Q_yy, Q_zz):
# def vec_zz_to_xy(info_dict, vec_zz, grid_averaging=True):
# def float_2_array(x):
# def reshape_to_ND(arr, N):
# def get_value(x):
# def get_shape(x):
# def vjp_maker_num(fn, arg_inds, steps):
#     def vjp_single_arg(ia):
#         def vjp_maker(fn_out, *args):
#             def vjp(v):
# def aniplot(F, source, steps, component='Ez', num_panels=10):
# def measure_fields(F, source, steps, probes, component='Ez'):
# def imarr(arr):
# def my_fft(x):    
# def fft_grad(g, ans, x):
# def get_spectrum(series, dt):
# def get_max_power_freq(series, dt):
# def get_spectral_power(series, dt):
# def plot_spectral_power(series, dt, f_top=2e14):
#     M = indices.shape[1]                                 # number of indices in the matrix
#     I = make_sparse(entries_1, indices_I, shape=(N, M))  # construct the I matrix
#     J = make_sparse(entries_1, indices_J, shape=(N, M))  # construct the J matrix
#     O = J.T                                              # make O = J^T matrix for consistency with my notes.
#     N = arg.size
#     ND = len(arr.shape)


# Path: ceviche/jacobians.py
import autograd.numpy as npa
from autograd.core import make_vjp, make_jvp
from autograd.wrap_util import unary_to_nary
from autograd.extend import vspace
from .utils import get_value, get_shape, get_value_arr, float_2_array


@unary_to_nary
def jacobian_reverse(fun, x):
    """ Compute jacobian of fun with respect to x using reverse mode differentiation"""
    vjp, ans = make_vjp(fun, x)
    grads = map(vjp, vspace(ans).standard_basis())
    m, n = _jac_shape(x, ans)
    return npa.reshape(npa.stack(grads), (n, m))


@unary_to_nary
def jacobian_forward(fun, x):
    """ Compute jacobian of fun with respect to x using forward mode differentiation"""
    jvp = make_jvp(fun, x)
    # ans = fun(x)
    val_grad = map(lambda b: jvp(b), vspace(x).standard_basis())
    vals, grads = zip(*val_grad)
    ans = npa.zeros((list(vals)[0].size,))  # fake answer so that dont have to compute it twice
    m, n = _jac_shape(x, ans)
    if _iscomplex(x):
        grads_real = npa.array(grads[::2])
        grads_imag = npa.array(grads[1::2])
        grads = grads_real - 1j * grads_imag
    return npa.reshape(npa.stack(grads), (m, n)).T


@unary_to_nary
def jacobian_numerical(fn, x, step_size=1e-7):
    """ numerically differentiate `fn` w.r.t. its argument `x` """
    in_array = float_2_array(x).flatten()

    out_array = float_2_array(fn(x)).flatten()