# Path: ceviche/utils.py
# def reshape_to_ND(arr, N):
#     """ Adds dimensions to arr until it is dimension N
#     """
# 
#     ND = len(arr.shape)
#     if ND > N:
#         raise ValueError("array is larger than {} dimensional, given shape {}".format(N, arr.shape))
#     extra_dims = (N - ND) * (1,)
#     return arr.reshape(arr.shape + extra_dims)
# 
# def grid_center_to_xyz(Q_mid, averaging=True):
#     """ Computes the interpolated value of the quantity Q_mid felt at the Ex, Ey, Ez positions of the Yee latice
#         Returns these three components
#     """
# 
#     # initialize
#     Q_xx = copy.copy(Q_mid)
#     Q_yy = copy.copy(Q_mid)
#     Q_zz = copy.copy(Q_mid)
# 
#     # if averaging, set the respective xx, yy, zz components to the midpoint in the Yee lattice.
#     if averaging:
# 
#         # get the value from the middle of the next cell over
#         Q_x_r = npa.roll(Q_mid, shift=1, axis=0)
#         Q_y_r = npa.roll(Q_mid, shift=1, axis=1)
#         Q_z_r = npa.roll(Q_mid, shift=1, axis=2)
# 
#         # average with the two middle values
#         Q_xx = (Q_mid + Q_x_r)/2
#         Q_yy = (Q_mid + Q_y_r)/2
#         Q_zz = (Q_mid + Q_z_r)/2
# 
#     return Q_xx, Q_yy, Q_zz
# 
# def grid_xyz_to_center(Q_xx, Q_yy, Q_zz):
#     """ Computes the interpolated value of the quantitys Q_xx, Q_yy, Q_zz at the center of Yee latice
#         Returns these three components
#     """
# 
#     # compute the averages
#     Q_xx_avg = (Q_xx.astype('float') + npa.roll(Q_xx, shift=1, axis=0))/2
#     Q_yy_avg = (Q_yy.astype('float') + npa.roll(Q_yy, shift=1, axis=1))/2
#     Q_zz_avg = (Q_zz.astype('float') + npa.roll(Q_zz, shift=1, axis=2))/2
# 
#     return Q_xx_avg, Q_yy_avg, Q_zz_avg
# 
# Path: ceviche/derivatives.py
# def curl_E(axis, Ex, Ey, Ez, dL):
#     if axis == 0:
#         return (npa.roll(Ez, shift=-1, axis=1) - Ez) / dL - (npa.roll(Ey, shift=-1, axis=2) - Ey) / dL
#     elif axis == 1:
#         return (npa.roll(Ex, shift=-1, axis=2) - Ex) / dL - (npa.roll(Ez, shift=-1, axis=0) - Ez) / dL
#     elif axis == 2:
#         return (npa.roll(Ey, shift=-1, axis=0) - Ey) / dL - (npa.roll(Ex, shift=-1, axis=1) - Ex) / dL
# 
# def curl_H(axis, Hx, Hy, Hz, dL):
#     if axis == 0:
#         return (Hz - npa.roll(Hz, shift=1, axis=1)) / dL - (Hy - npa.roll(Hy, shift=1, axis=2)) / dL
#     elif axis == 1:
#         return (Hx - npa.roll(Hx, shift=1, axis=2)) / dL - (Hz - npa.roll(Hz, shift=1, axis=0)) / dL
#     elif axis == 2:
#         return (Hy - npa.roll(Hy, shift=1, axis=0)) / dL - (Hx - npa.roll(Hx, shift=1, axis=1)) / dL


import numpy as np
import autograd.numpy as npa
from copy import copy, deepcopy
from .constants import *
from .utils import reshape_to_ND, grid_center_to_xyz, grid_xyz_to_center
from .derivatives import curl_E, curl_H
    @npml.setter
    def npml(self, new_npml):
        """ Defines some attributes when npml is set. """
        self.__npml = new_npml
        self._compute_sigmas()

    @property
    def eps_r(self):
        """ Returns the relative permittivity grid """
        return self.__eps_r

    @eps_r.setter
    def eps_r(self, new_eps):
        """ Defines some attributes when eps_r is set. """
        self.__eps_r = new_eps
        self.eps_xx, self.eps_yy, self.eps_zz = grid_center_to_xyz(self.__eps_r)
        self.eps_arr = self.__eps_r.flatten()
        self.N = self.eps_arr.size
        self.grid_shape = self.Nx, self.Ny, self.Nz = self.__eps_r.shape
        self._compute_update_parameters()
        self.initialize_fields()

    def forward(self, Jx=None, Jy=None, Jz=None):
        """ one time step of FDTD """

        self.t_index += 1

        # get curls of E
        CEx = curl_E(0, self.Ex, self.Ey, self.Ez, self.dL)
        CEy = curl_E(1, self.Ex, self.Ey, self.Ez, self.dL)

# Path: ceviche/utils.py
# def reshape_to_ND(arr, N):
#     """ Adds dimensions to arr until it is dimension N
#     """
# 
#     ND = len(arr.shape)
#     if ND > N:
#         raise ValueError("array is larger than {} dimensional, given shape {}".format(N, arr.shape))
#     extra_dims = (N - ND) * (1,)
#     return arr.reshape(arr.shape + extra_dims)
# 
# def grid_center_to_xyz(Q_mid, averaging=True):
#     """ Computes the interpolated value of the quantity Q_mid felt at the Ex, Ey, Ez positions of the Yee latice
#         Returns these three components
#     """
# 
#     # initialize
#     Q_xx = copy.copy(Q_mid)
#     Q_yy = copy.copy(Q_mid)
#     Q_zz = copy.copy(Q_mid)
# 
#     # if averaging, set the respective xx, yy, zz components to the midpoint in the Yee lattice.
#     if averaging:
# 
#         # get the value from the middle of the next cell over
#         Q_x_r = npa.roll(Q_mid, shift=1, axis=0)
#         Q_y_r = npa.roll(Q_mid, shift=1, axis=1)
#         Q_z_r = npa.roll(Q_mid, shift=1, axis=2)
# 
#         # average with the two middle values
#         Q_xx = (Q_mid + Q_x_r)/2
#         Q_yy = (Q_mid + Q_y_r)/2
#         Q_zz = (Q_mid + Q_z_r)/2
# 
#     return Q_xx, Q_yy, Q_zz
# 
# def grid_xyz_to_center(Q_xx, Q_yy, Q_zz):
#     """ Computes the interpolated value of the quantitys Q_xx, Q_yy, Q_zz at the center of Yee latice
#         Returns these three components
#     """
# 
#     # compute the averages
#     Q_xx_avg = (Q_xx.astype('float') + npa.roll(Q_xx, shift=1, axis=0))/2
#     Q_yy_avg = (Q_yy.astype('float') + npa.roll(Q_yy, shift=1, axis=1))/2
#     Q_zz_avg = (Q_zz.astype('float') + npa.roll(Q_zz, shift=1, axis=2))/2
# 
#     return Q_xx_avg, Q_yy_avg, Q_zz_avg
# 
# Path: ceviche/derivatives.py
# def curl_E(axis, Ex, Ey, Ez, dL):
#     if axis == 0:
#         return (npa.roll(Ez, shift=-1, axis=1) - Ez) / dL - (npa.roll(Ey, shift=-1, axis=2) - Ey) / dL
#     elif axis == 1:
#         return (npa.roll(Ex, shift=-1, axis=2) - Ex) / dL - (npa.roll(Ez, shift=-1, axis=0) - Ez) / dL
#     elif axis == 2:
#         return (npa.roll(Ey, shift=-1, axis=0) - Ey) / dL - (npa.roll(Ex, shift=-1, axis=1) - Ex) / dL
# 
# def curl_H(axis, Hx, Hy, Hz, dL):
#     if axis == 0:
#         return (Hz - npa.roll(Hz, shift=1, axis=1)) / dL - (Hy - npa.roll(Hy, shift=1, axis=2)) / dL
#     elif axis == 1:
#         return (Hx - npa.roll(Hx, shift=1, axis=2)) / dL - (Hz - npa.roll(Hz, shift=1, axis=0)) / dL
#     elif axis == 2:
#         return (Hy - npa.roll(Hy, shift=1, axis=0)) / dL - (Hx - npa.roll(Hx, shift=1, axis=1)) / dL


# Path: ceviche/fdtd.py
import numpy as np
import autograd.numpy as npa
from copy import copy, deepcopy
from .constants import *
from .utils import reshape_to_ND, grid_center_to_xyz, grid_xyz_to_center
from .derivatives import curl_E, curl_H

    @npml.setter
    def npml(self, new_npml):
        """ Defines some attributes when npml is set. """
        self.__npml = new_npml
        self._compute_sigmas()

    @property
    def eps_r(self):
        """ Returns the relative permittivity grid """
        return self.__eps_r

    @eps_r.setter
    def eps_r(self, new_eps):
        """ Defines some attributes when eps_r is set. """
        self.__eps_r = new_eps
        self.eps_xx, self.eps_yy, self.eps_zz = grid_center_to_xyz(self.__eps_r)
        self.eps_arr = self.__eps_r.flatten()
        self.N = self.eps_arr.size
        self.grid_shape = self.Nx, self.Ny, self.Nz = self.__eps_r.shape
        self._compute_update_parameters()
        self.initialize_fields()

    def forward(self, Jx=None, Jy=None, Jz=None):
        """ one time step of FDTD """

        self.t_index += 1

        # get curls of E
        CEx = curl_E(0, self.Ex, self.Ey, self.Ez, self.dL)
        CEy = curl_E(1, self.Ex, self.Ey, self.Ez, self.dL)

        CEz = curl_E(2, self.Ex, self.Ey, self.Ez, self.dL)