# 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/diag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, 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
from cubical.kernels import diag_complex
import numpy as np
import cubical.kernels
    """

    n_dir = jhr.shape[0]
    n_tim = jhr.shape[1]
    n_fre = jhr.shape[2]
    n_ant = jhr.shape[3]

    for aa in prange(n_ant):
        for t in range(n_tim):
            for f in range(n_fre):
                for d in range(n_dir):
                    upd[d,t,f,aa,0,0] = jhjinv[d,t,f,aa,0,0]*jhr[d,t,f,aa,0,0]
                    upd[d,t,f,aa,0,1] = upd[d,t,f,aa,1,0] = 0
                    upd[d,t,f,aa,1,1] = jhjinv[d,t,f,aa,1,1]*jhr[d,t,f,aa,1,1]

# Remaining kernel functions are reused.

# J^H is computed assuming diagonal gains.
compute_jh = diag_complex.compute_jh

# J^H.J inverse is computed assuming diagonal blocks.
compute_jhjinv = generics.compute_diag_inverse

# Residuals computed assuming diagonal gains.
compute_residual = diag_complex.compute_residual

# Corrected visibilities computed assuming diagonal gains.
compute_corrected = diag_complex.compute_corrected

# Gains applied as diagonal.

# 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/diag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, 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/phase_only.py
from builtins import range
from numba import jit, prange
from cubical.kernels import generics
from cubical.kernels import full_complex
from cubical.kernels import diag_complex
import numpy as np
import cubical.kernels

    """

    n_dir = jhr.shape[0]
    n_tim = jhr.shape[1]
    n_fre = jhr.shape[2]
    n_ant = jhr.shape[3]

    for aa in prange(n_ant):
        for t in range(n_tim):
            for f in range(n_fre):
                for d in range(n_dir):
                    upd[d,t,f,aa,0,0] = jhjinv[d,t,f,aa,0,0]*jhr[d,t,f,aa,0,0]
                    upd[d,t,f,aa,0,1] = upd[d,t,f,aa,1,0] = 0
                    upd[d,t,f,aa,1,1] = jhjinv[d,t,f,aa,1,1]*jhr[d,t,f,aa,1,1]

# Remaining kernel functions are reused.

# J^H is computed assuming diagonal gains.
compute_jh = diag_complex.compute_jh

# J^H.J inverse is computed assuming diagonal blocks.
compute_jhjinv = generics.compute_diag_inverse

# Residuals computed assuming diagonal gains.
compute_residual = diag_complex.compute_residual

# Corrected visibilities computed assuming diagonal gains.
compute_corrected = diag_complex.compute_corrected

# Gains applied as diagonal.

apply_gains = diag_complex.apply_gains