# Path: cubical/kernels/generics.py
# def compute_2x2_inverse(x, xinv, flags, eps, flagbit):
# def compute_diag_inverse(x, xinv, flags, eps, flagbit):
# def compute_chisq(r, chisq):
# def compute_chisq_diag(r, chisq):
# def compute_chisq_offdiag(r, chisq):
# 
# Path: cubical/kernels/full_complex.py
# def allocate_vis_array(shape, dtype, zeros=False):
# def allocate_gain_array(shape, dtype, zeros=False):
# def allocate_flag_array(shape, dtype, zeros=False):
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):


from builtins import range
from numba import jit, prange
from cubical.kernels import generics
from cubical.kernels import full_complex
import numpy as np
import cubical.kernels



use_parallel = True if cubical.kernels.num_omp_threads > 1 else False
use_cache = cubical.kernels.use_cache

# defines memory layout of model-like arrays (axis layout is NDxNMxNTxNFxNAxNAxNCxNC)
_model_axis_layout = [4,5,1,2,3,0,6,7]    # layout is AAMTFD

# defines memory layout of gain-like arrays  (axis layout is NDxNTxNFxNAxNCxNC)
_gain_axis_layout = [3,1,2,0,4,5]   # layout is ATFD

# defines memory layout of flag-like arrays  (axis layout is NTxNFxNAxNA)
_flag_axis_layout = [2,3,0,1]      # layout is AATF

def allocate_vis_array(shape, dtype, zeros=False):
    """Allocates a visibility array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _model_axis_layout, zeros=zeros)

def allocate_gain_array(shape, dtype, zeros=False):
    """Allocates a gain array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _gain_axis_layout, zeros=zeros)

def allocate_flag_array(shape, dtype, zeros=False):
    """Allocates a flag array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _flag_axis_layout, zeros=zeros)

allocate_param_array = allocate_gain_array

# compute_residual is identical to the general full complex case.

# Path: cubical/kernels/generics.py
# def compute_2x2_inverse(x, xinv, flags, eps, flagbit):
# def compute_diag_inverse(x, xinv, flags, eps, flagbit):
# def compute_chisq(r, chisq):
# def compute_chisq_diag(r, chisq):
# def compute_chisq_offdiag(r, chisq):
# 
# Path: cubical/kernels/full_complex.py
# def allocate_vis_array(shape, dtype, zeros=False):
# def allocate_gain_array(shape, dtype, zeros=False):
# def allocate_flag_array(shape, dtype, zeros=False):
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):


# Path: cubical/kernels/full_W_complex.py
from builtins import range
from numba import jit, prange
from cubical.kernels import generics
from cubical.kernels import full_complex
import numpy as np
import cubical.kernels




use_parallel = True if cubical.kernels.num_omp_threads > 1 else False
use_cache = cubical.kernels.use_cache

# defines memory layout of model-like arrays (axis layout is NDxNMxNTxNFxNAxNAxNCxNC)
_model_axis_layout = [4,5,1,2,3,0,6,7]    # layout is AAMTFD

# defines memory layout of gain-like arrays  (axis layout is NDxNTxNFxNAxNCxNC)
_gain_axis_layout = [3,1,2,0,4,5]   # layout is ATFD

# defines memory layout of flag-like arrays  (axis layout is NTxNFxNAxNA)
_flag_axis_layout = [2,3,0,1]      # layout is AATF

def allocate_vis_array(shape, dtype, zeros=False):
    """Allocates a visibility array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _model_axis_layout, zeros=zeros)

def allocate_gain_array(shape, dtype, zeros=False):
    """Allocates a gain array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _gain_axis_layout, zeros=zeros)

def allocate_flag_array(shape, dtype, zeros=False):
    """Allocates a flag array of the desired shape, laid out in preferred order"""
    return cubical.kernels.allocate_reordered_array(shape, dtype, _flag_axis_layout, zeros=zeros)

allocate_param_array = allocate_gain_array

# compute_residual is identical to the general full complex case.

compute_residual = full_complex.compute_residual