# -*- coding: utf-8 -*-
"""
Created on Tue Aug  3 18:20:20 2021

@author: bchan
"""
import numpy as np
from numba import jit

@jit
def calc_elec_currz(cond, pot, img_shape, dz=1, EPS=1e-2):
    nz = img_shape[0]; ny = img_shape[1]; nx = img_shape[2];
    nxy = nx*ny
    y = np.zeros(nx*ny*nz)

    for k in range(nz):
        K = nx*ny*k
        for j in range(ny):
            for i in range(nx):
                if(cond[K] < EPS):
                    y[K] = 0.
                elif k < nz-1 and cond[K + nxy] > EPS:
                    y[K] = -(pot[K+nxy] - pot[K])*cond[K]/dz
                elif k == nz-1 and cond[K-nxy] > EPS:
                    y[K] = -(pot[K] - pot[K - nxy])*cond[K]/dz
                else:
                    y[K] = 0.
                K+=1
    
    return y
