# AUTOGENERATED! DO NOT EDIT! File to edit: README.ipynb (unless otherwise specified).

__all__ = ['SAD', 'rect_loss_p', 'argmin_int', 'rect_loss_l', 'min_path_int', 'rect_match_arr_min_path',
           'make_rect_match_arr_min_path', 'argmin_sub', 'min_path_sub', 'make_min_path_int_dp', 'interp',
           'make_min_path_sub_dp', 'rect_match_pyr', 'isclose', 'clip', 'line_loop', 'rect_loss_arr', 'callback_sgm',
           'rect_match_arr_sgm', 'make_rect_match_arr_sgm', 'RectMatch']

# Cell
import numba
import numpy as np

from camera_calib.utils import *

# Cell
@numba.jit(nopython=True)
def SAD(arr1, arr2):
    l = 0
    for i in range(arr1.shape[0]):
        for j in range(arr2.shape[1]):
            l += abs(arr1[i,j] - arr2[i,j])
    return l

# Cell
@numba.jit(nopython=True)
def rect_loss_p(arr1, arr2, x, y, min_disp, max_disp, hw, loss, buf_loss):
    h_arr, w_arr = arr1.shape

    l_t, t_t, r_t, b_t = max(x-hw, 0), max(y-hw, 0), min(x+hw, w_arr-1), min(y+hw, h_arr-1)
    h_t, w_t = b_t-t_t+1, r_t-l_t+1
    for j in range(min_disp, max_disp+1):
        if (j+l_t >= 0) and (j+r_t < w_arr): # Template in bounds
            buf_loss[j-min_disp] = loss(arr1[t_t:t_t+h_t, l_t:l_t+w_t], arr2[t_t:t_t+h_t, j+l_t:j+l_t+w_t])
        else:                                # Template out of bound
            buf_loss[j-min_disp] = np.inf

# Cell
@numba.jit(nopython=True)
def argmin_int(arr): return np.argmin(arr)

# Cell
@numba.jit(nopython=True)
def rect_loss_l(arr1, arr2, y, r_disp, hw, loss, buf_loss, arr_disp_init=None):
    h_arr, w_arr = arr1.shape

    for i in range(w_arr):
        disp_init = 0 if arr_disp_init is None else arr_disp_init[y, i]
        min_disp, max_disp = [disp + disp_init for disp in r_disp]
        rect_loss_p(arr1, arr2, i, y, min_disp, max_disp, hw, loss, buf_loss[i])

# Cell
@numba.jit(nopython=True)
def min_path_int(arr_loss, buf_path):
    for i in range(len(arr_loss)):
        buf_path[i] = argmin_int(arr_loss[i])

# Cell
@numba.jit(nopython=True, parallel=True)
def rect_match_arr_min_path(arr1, arr2, r_disp, hw, loss, min_path, arr_disp_init=None):
    h_arr, w_arr = arr1.shape[0], arr1.shape[1]

    arr_disp = np.empty((h_arr, w_arr))
    for i in numba.prange(h_arr):
        buf_loss = np.empty((arr1.shape[1], r_disp[1]-r_disp[0]+1))
        rect_loss_l(arr1, arr2, i, r_disp, hw, loss, buf_loss, arr_disp_init)
        min_path(buf_loss, arr_disp[i]) # range offset and initial disparity need to be applied after
        arr_disp[i] += r_disp[0]
        if arr_disp_init is not None: arr_disp[i] += arr_disp_init[i]
    return arr_disp

# Cell
def make_rect_match_arr_min_path(r_disp, hw, loss, min_path):
    @numba.jit(nopython=True)
    def rect_match_arr(arr1, arr2, arr_disp_init=None):
        return rect_match_arr_min_path(arr1, arr2, r_disp, hw, loss, min_path, arr_disp_init)
    return rect_match_arr

# Cell
@numba.jit(nopython=True)
def argmin_sub(arr):
    idx_min = argmin_int(arr)
    if 1 <= idx_min <= len(arr)-2:
        delta_idx = ((arr[idx_min+1]-arr[idx_min-1])/2)/(arr[idx_min+1]-2*arr[idx_min]+arr[idx_min-1])
        if np.isnan(delta_idx): delta_idx =  0
        if delta_idx < -1:      delta_idx = -1
        if delta_idx >  1:      delta_idx =  1
        idx_min = idx_min - delta_idx
    return idx_min

# Cell
@numba.jit(nopython=True)
def min_path_sub(arr_loss, buf_path):
    for i in range(len(arr_loss)):
        buf_path[i] = argmin_sub(arr_loss[i, :])

# Cell
@numba.jit(nopython=True)
def _min_path_int_dp(arr_loss, buf_path, r_disp, max_change, penalty_disp):
    buf_route     = np.zeros(arr_loss.shape)
    buf_move      = np.empty((2*max_change+1, r_disp[1]-r_disp[0]+1))
    buf_loss_prev = arr_loss[-1].copy() # Going backwards, this is initial optimal loss
    for i in range(len(arr_loss)-2, -1, -1):
        # Get loss of each move
        buf_move[:] = np.inf
        for j in range(-max_change, max_change+1):
            idx_minl, idx_maxl = max( j,0), min(arr_loss.shape[1]+j,arr_loss.shape[1])
            idx_minm, idx_maxm = max(-j,0), min(arr_loss.shape[1]-j,arr_loss.shape[1])
            buf_move[j+max_change, idx_minm:idx_maxm] = buf_loss_prev[idx_minl:idx_maxl] + abs(j)*penalty_disp
        # Get optimal move and store it
        for j in range(buf_move.shape[1]):
            idx_min = np.argmin(buf_move[:,j])
            buf_route[i,j] = idx_min - max_change
            buf_loss_prev[j] = arr_loss[i,j] + buf_move[idx_min, j] # total loss = loss + optimal move
    # Gather path
    buf_path[0] = np.argmin(buf_loss_prev)
    for i in range(1, len(buf_route)):
        buf_path[i] = buf_path[i-1] + buf_route[i-1, int(buf_path[i-1])]

# Cell
def make_min_path_int_dp(r_disp, max_change, penalty_disp):
    @numba.jit(nopython=True)
    def min_path(arr_loss, buf_path):
        return _min_path_int_dp(arr_loss, buf_path, r_disp, max_change, penalty_disp)
    return min_path

# Cell
@numba.jit(nopython=True)
def interp(arr, idx):
    if idx < 0 or idx > len(arr)-1: val = np.nan
    else:
        idx_f = np.floor(idx)
        if idx == idx_f: val = arr[int(idx_f)]
        else:            val = (idx_f+1-idx)*arr[int(idx_f)] + (idx-idx_f)*arr[int(idx_f)+1]
    return val

# Cell
@numba.jit(nopython=True)
def _min_path_sub_dp(arr_loss, buf_path, r_disp, max_change, penalty_disp):
    buf_route     = np.zeros(arr_loss.shape)
    buf_move      = np.empty((2*max_change+1, r_disp[1]-r_disp[0]+1))
    buf_loss_prev = arr_loss[-1].copy() # Going backwards, this is initial optimal loss
    for i in range(len(arr_loss)-2, -1, -1):
        # Get loss of each move
        buf_move[:] = np.inf
        for j in range(-max_change, max_change+1):
            idx_minl, idx_maxl = max( j,0), min(arr_loss.shape[1]+j,arr_loss.shape[1])
            idx_minm, idx_maxm = max(-j,0), min(arr_loss.shape[1]-j,arr_loss.shape[1])
            buf_move[j+max_change, idx_minm:idx_maxm] = buf_loss_prev[idx_minl:idx_maxl] + abs(j)*penalty_disp
        # Get optimal move and store it
        for j in range(buf_move.shape[1]):
            idx_min = argmin_sub(buf_move[:,j])
            buf_route[i,j] = idx_min - max_change
            buf_loss_prev[j] = arr_loss[i,j] + interp(buf_move[:, j], idx_min)
    # Gather path
    buf_path[0] = argmin_sub(buf_loss_prev)
    for i in range(1, len(buf_route)):
        buf_path[i] = buf_path[i-1] + interp(buf_route[i-1, :], buf_path[i-1])

# Cell
def make_min_path_sub_dp(r_disp, max_change, penalty_disp):
    @numba.jit(nopython=True)
    def min_path(arr_loss, buf_path):
        return _min_path_sub_dp(arr_loss, buf_path, r_disp, max_change, penalty_disp)
    return min_path

# Cell
def rect_match_pyr(arr1, arr2, rect_match_arr, steps=3):
    if not np.all(shape(arr1) % 2**steps == 0): raise RuntimeError('Shape must be divisible by 2^steps')

    def _get_pyr(arr):
        arr_pyr = [arr]
        for idx in range(steps-1):
            arr_pyr.append(imresize(arr_pyr[-1], shape(arr_pyr[-1])/2))
        return arr_pyr
    arr1_pyr, arr2_pyr = [_get_pyr(arr) for arr in [arr1, arr2]]

    arr_disp = None
    for idx in range(steps-1,-1,-1):
        arr1, arr2 = arr1_pyr[idx], arr2_pyr[idx]
        if arr_disp is not None:
            arr_disp = imresize(2*arr_disp, 2*shape(arr_disp)) # Remember to multiply disparities by 2
            arr_disp = np.round(arr_disp).astype(np.long)      # Must be integer
        arr_disp = rect_match_arr(arr1, arr2, arr_disp)
    return arr_disp

# Cell
@numba.jit(nopython=True)
def isclose(x, y, atol=1e-8): return abs(x-y) < atol

# Cell
@numba.jit(nopython=True)
def clip(x, min_x, max_x): return max(min(x, max_x), min_x)

# Cell
@numba.jit(nopython=True)
def line_loop(arr, theta, callback):
    h, w = arr.shape[0:2]

    # Get dx, dy
    dx, dy = np.cos(theta), np.sin(theta)
    if    isclose(dx, 0): dx, dy = 0, np.sign(dy)*h
    elif  isclose(dy, 0): dx, dy = np.sign(dx)*w, 0
    else:
        sf = min(abs(dx), abs(dy))
        dx, dy = np.round(dx/sf), np.round(dy/sf)
        dx, dy = clip(dx, 1-w, w-1), clip(dy, 1-h, h-1)
    dx, dy = int(dx), int(dy)

    # Get increments
    if abs(dx) > abs(dy): l, dx_l, dy_l, dx_c, dy_c = abs(dx), np.sign(dx), 0, 0, np.sign(dy)
    else:                 l, dx_l, dy_l, dx_c, dy_c = abs(dy), 0, np.sign(dy), np.sign(dx), 0

    # Get initial p0
    if   0 <= dx <   w and 0 <  dy <=  h: x0, y0 = w-1,   0
    elif 0 <  dx <=  w and 0 >= dy >  -h: x0, y0 =   0,   0
    elif 0 >= dx >  -w and 0 >  dy >= -h: x0, y0 =   0, h-1
    elif 0 >  dx >= -w and 0 <= dy <   h: x0, y0 = w-1, h-1
    else: raise RuntimeError('Invalid dx, dy')

    # Do line iterations
    it_p, num_p = 0, h*w              # Iterate until it_p == num_p
    while it_p < num_p:
        x, y, started = x0, y0, False # Start tracing line at p0
        while True:
            for i in range(l):
                if (0 <= x < w) and (0 <= y < h):
                    if not started: callback.start_line(arr, x, y); started=True
                    else:           callback.in_line(arr, x, y)
                    it_p += 1         # Point increment
                x += dx_l; y += dy_l  # Line increment
            x += dx_c; y += dy_c      # Change increment
            if not ((0 <= x < w) and (0 <= y < h)): break # Lines goes out of array, so end this line
        # Shift start of line based on current p0
        if   x0 >= w-1 and y0 >=   1: y0 -= max(abs(dy), 1)
        elif x0 >=   1 and y0 <=   0: x0 -= max(abs(dx), 1)
        elif x0 <=   0 and y0 <= h-2: y0 += max(abs(dy), 1)
        elif x0 <= w-2 and y0 >= h-1: x0 += max(abs(dx), 1)
        else: raise RuntimeError('Invalid x0, y0')

# Cell
@numba.jit(nopython=True, parallel=True)
def rect_loss_arr(arr1, arr2, r_disp, hw, loss, buf_loss, arr_disp_init=None):
    h_arr, w_arr = arr1.shape
    for i in numba.prange(h_arr):
        rect_loss_l(arr1, arr2, i, r_disp, hw, loss, buf_loss[i], arr_disp_init)

# Cell
@numba.experimental.jitclass([('buf_accum', numba.float64[:,:,:]),
                              ('buf_move', numba.float64[:,:]),
                              ('x_prev', numba.int32),
                              ('y_prev', numba.int32),
                              ('max_change', numba.int32),
                              ('penalty_disp', numba.int32)])
class callback_sgm(object):
    def __init__(self, sz, r_disp, max_change, penalty_disp):
        self.buf_accum = np.empty((sz[0], sz[1], r_disp[1]-r_disp[0]+1))
        self.buf_move = np.empty((2*max_change+1, r_disp[1]-r_disp[0]+1))
        self.x_prev, self.y_prev = -1, -1
        self.max_change, self.penalty_disp = max_change, penalty_disp

    def start_line(self, arr, x, y):
        self.buf_accum[y, x] = arr[y, x] # Initialize
        self.x_prev, self.y_prev = x, y

    def in_line(self, arr, x, y):
        # Get loss of each move
        self.buf_move[:] = np.inf
        for j in range(-self.max_change, self.max_change+1):
            idx_minl, idx_maxl = max( j,0), min(self.buf_move.shape[1]+j,self.buf_move.shape[1])
            idx_minm, idx_maxm = max(-j,0), min(self.buf_move.shape[1]-j,self.buf_move.shape[1])
            self.buf_move[j+max_change, idx_minm:idx_maxm] = \
                self.buf_accum[self.y_prev, self.x_prev, idx_minl:idx_maxl] + abs(j)*self.penalty_disp

        # Get optimal cost and store it
        for j in range(self.buf_move.shape[1]):
            self.buf_accum[y, x, j] = arr[y, x, j] + np.min(self.buf_move[:, j])

        self.x_prev, self.y_prev = x, y

# Cell
@numba.jit(nopython=True, parallel=True)
def rect_match_arr_sgm(arr1, arr2, r_disp, hw, loss, max_change, penalty_disp, thetas, arr_disp_init=None):
    h_arr, w_arr = arr1.shape[0], arr1.shape[1]

    # Precompute losses
    buf_loss = np.empty((h_arr, w_arr, r_disp[1]-r_disp[0]+1))
    rect_loss_arr(arr1, arr2, r_disp, hw, loss, buf_loss, arr_disp_init)

    # Do SGM accumulation
    callback = callback_sgm(arr1.shape, r_disp, max_change, penalty_disp)
    buf_accum = np.zeros((h_arr, w_arr, r_disp[1]-r_disp[0]+1))
    for theta in thetas:
        line_loop(buf_loss, theta, callback)
        buf_accum += callback.buf_accum

    # Get disparity map
    arr_disp = np.empty((h_arr, w_arr))
    for i in range(arr1.shape[0]):
        for j in range(arr2.shape[1]):
            arr_disp[i, j] = np.argmin(buf_accum[i, j]) + r_disp[0]
    if arr_disp_init is not None: arr_disp += arr_disp_init
    return arr_disp

# Cell
def make_rect_match_arr_sgm(r_disp, hw, loss, max_change, penalty_disp, thetas):
    @numba.jit(nopython=True)
    def rect_match_arr(arr1, arr2, arr_disp_init=None):
        return rect_match_arr_sgm(arr1, arr2, r_disp, hw, loss, max_change, penalty_disp, thetas, arr_disp_init)
    return rect_match_arr

# Cell
class RectMatch:
    def __init__(self, type_rect_match, hw=15, r_disp=(-15,15), loss=SAD, max_change=3, penalty_disp=2, steps=1):
        if type_rect_match in ['int', 'sub', 'int_dp', 'sub_dp']:
            if   type_rect_match == 'int':
                min_path = min_path_int
            elif type_rect_match == 'sub':
                min_path = min_path_sub
            elif type_rect_match == 'int_dp':
                min_path = make_min_path_int_dp(r_disp, max_change, penalty_disp)
            elif type_rect_match == 'sub_dp':
                min_path = make_min_path_sub_dp(r_disp, max_change, penalty_disp)
            rect_match_arr = make_rect_match_arr_min_path(r_disp, hw, loss, min_path)
        elif type_rect_match == 'sgm':
            rect_match_arr = make_rect_match_arr_sgm(r_disp, hw, loss, max_change, penalty_disp, thetas)
        else:
            raise RuntimeError(f'Unrecognized min path type: {type_rect_match}')
        self.rect_match_arr, self.steps = rect_match_arr, steps

    def __call__(self, arr1, arr2):
        return rect_match_pyr(arr1, arr2, self.rect_match_arr, self.steps)