import numpy as np
import taichi as ti
from scipy.interpolate import interp1d


######################################################################################
# Aux functions.

def simpson(x:np.ndarray) -> np.ndarray:
    if x.shape[0]==1:
        return x[0]
    assert x.shape[0]%2==1
    return np.sum(x[:-2:2] + 4*x[1:-1:2] + x[2::2], axis=0)/6

######################################################################################
# Functions of wavelength.

class Function:
    def _normalize(self, x):
        x = np.asarray(x)
        if x.ndim==0:
            x = x[None]
        assert x.ndim == 1
        return x

    def compute(self, x: np.ndarray) -> np.ndarray:
        raise NotImplementedError
    
    def __call__(self, x: np.ndarray) -> np.ndarray:
        return self.compute(self._normalize(x))

class ConstantFunction(Function):
    def __init__(self, x):
        self.x = np.asarray(x)

    def compute(self, x: np.ndarray) -> np.ndarray:
        return np.repeat(self.x[None], x.shape[0], axis=0)

class GaussianFunction(Function):
    def __init__(self, center, width):
        self._center = np.asarray(center)
        self._width = np.asarray(width)
        self._2width2 = 2*self._width**2

    def compute(self, x):
        return np.exp(-(x[:,None]-self._center[None])**2/self._2width2)/self._width

class SimplifiedRGB(GaussianFunction):
    def __init__(self):
        super().__init__(
            center=[0.610,0.540,0.450],
            width=[0.050,0.040,0.030],
            )

class BlackbodyRadiation(Function):
    h = 6.62607015e-34 # J⋅Hz−1
    c = 299792458 # m/s
    k = 1.380649e-23 # J⋅K−1

    def __init__(self, temperature_in_K=[5800]):
        self.T = np.asarray(temperature_in_K)
        assert self.T.ndim==1, f"{self.T.shape=}"

    def compute(self, x, verbose=False):
        assert x.ndim==1, f"{x.shape=}"
        wavelength = x*1e-6 # meters
        b = 2.897771955e-3 # [m K], Wien displacement constant.
        if verbose:
            print(f"Radiation peak at {b/self.T[0]*1e6:.3f} um")
        result = (2*self.h*self.c**2/wavelength[:,None]**5)/( np.exp( self.h*self.c/(self.k*self.T[None,:]*wavelength[:,None]) ) - 1 )
        return result 

class FunctionTable(Function):
    def __init__(self, x, y):
        self.x = np.asarray(x)
        self.y = np.asarray(y)
        self.poly = interp1d(x=x, y=y, axis=0, bounds_error=False, fill_value=0)

    @classmethod
    def from_csv(cls, name:str) -> 'FunctionTable':
        data = np.loadtxt(fname=name, dtype=np.float32, delimiter=None)
        return cls(x=data[:,0], y=data[:,1:])

    def compute(self, x):
        return self.poly(x)


################################################################################################################################
# Multispectral cross-polarization.

def field_like(x):
    return ti.field(x.dtype, shape=x.shape)

class Quadrature:
    def __init__(self):
        pass

    def get_weights(self, x):
        raise NotImplementedError
    
class Trapezoid(Quadrature):
    def get_weights(self, x):
        dx = np.diff(x)/2
        w = np.zeros_like(x)
        w[:-1] = dx
        w[1:] += dx
        return w


def normalize_wavelengths(wavelengths, mono: bool) -> np.ndarray:
    wavelengths = np.asarray(wavelengths)
    if wavelengths.ndim != 1:
        raise ValueError(f"`wavelengths` must be a 1D array, got {wavelengths.shape}.")
    if mono:
        if wavelengths.shape[0] != 1:
            raise ValueError(
                f"`mono=True` requires exactly one wavelength, got shape {wavelengths.shape}."
            )
    elif wavelengths.shape[0] <= 1:
        raise ValueError(
            f"`mono=False` requires at least two wavelengths, got shape {wavelengths.shape}."
        )
    return wavelengths


def normalize_color(rgb):
    finite = np.isfinite(rgb)
    if np.any(finite):
        scale = np.max(rgb[finite])
        if scale > 0:
            n_rgb = rgb / scale
        else:
            n_rgb = np.zeros_like(rgb)
    else:
        n_rgb = np.zeros_like(rgb)
    return n_rgb


def render_cp_multi(nn:np.ndarray, wavelengths:np.ndarray, ne:Function, no:Function, emission:Function, efficiency:Function, 
                    thickness:float, polarizer:float=0., deltafilter:float=np.pi/2, subdiv=1, normalize:bool=True, mono:bool=False
    ):
    nn = np.asarray(nn)
    if nn.ndim != 4 or nn.shape[3] != 3:
        raise ValueError(f"`nn` must have shape (sx, sy, sz, 3), got {nn.shape}.")

    wavelengths = normalize_wavelengths(wavelengths, mono=mono)

    n_ne = np.asarray(ne(wavelengths))
    if n_ne.shape != wavelengths.shape:
        raise ValueError(
            f"`ne(wavelengths)` must return one scalar per wavelength, got {n_ne.shape} for {wavelengths.shape}."
        )

    n_no = np.asarray(no(wavelengths))
    if n_no.shape != wavelengths.shape:
        raise ValueError(
            f"`no(wavelengths)` must return one scalar per wavelength, got {n_no.shape} for {wavelengths.shape}."
        )

    n_emission = np.asarray(emission(wavelengths))
    if n_emission.shape != wavelengths.shape + (1,):
        raise ValueError(
            f"`emission(wavelengths)` must return shape {wavelengths.shape + (1,)}, got {n_emission.shape}."
        )

    n_efficiency = np.asarray(efficiency(wavelengths))
    if n_efficiency.shape != wavelengths.shape + (3,):
        raise ValueError(
            f"`efficiency(wavelengths)` must return shape {wavelengths.shape + (3,)}, got {n_efficiency.shape}."
        )

    render = JonesCalculus(nn_shape=nn.shape, subdiv=subdiv, thickness=thickness)
    render.upload(nn)
    n_rgb = render.render(
        wavelengths=wavelengths,
        emission=emission,
        efficiency=efficiency,
        ne=ne,
        no=no,
        polarizer=polarizer,
        deltafilter=deltafilter,
        mono=mono,
    )
    n_rgb = np.asarray(n_rgb)
    if normalize:
        n_rgb=normalize_color(n_rgb)
    n_rgb = np.nan_to_num(n_rgb, nan=0.0, posinf=0.0, neginf=0.0)
    return np.clip(n_rgb, 0., 1.) # Save in linear color space.

def npunwrap(x):
    if isinstance(x, np.ndarray):
        if x.size!=1:
            raise ValueError(f"Not a single scalar array shape: {x.shape}")
        return x.ravel()[0]
    raise ValueError(f"Unsupported value {type(x)}")

@ti.data_oriented
class JonesCalculus:
    def __init__(self, nn_shape:tuple[int], subdiv:int, thickness:float, quadrature:Quadrature=None):
        self._nn_shape = nn_shape
        self._subdiv = subdiv
        self._thickness = thickness
        # Allocate memory.
        self._prepare_buffers()
        # Initialize integrator
        if quadrature is None:
            quadrature = Trapezoid()
        self._quadrature = quadrature
        # Compute parameters.
        sz = self._nn_shape[2]
        self._dz = self._thickness/(self._subdiv*(sz-1))

    def _prepare_buffers(self):
        assert len(self._nn_shape)==4 and self._nn_shape[3]==3, f"Wrong director field shape {self._nn_shape}"
        sx, sy, sz, _dim = self._nn_shape
        sx2, sy2 = sx*self._subdiv, sy*self._subdiv
        self._f_nn = ti.field(ti.math.vec3, shape=(sx, sy, sz))
        self._f_nn2 = ti.field(ti.math.vec3, shape=(sx2, sy2))
        self._f_ex = ti.field(ti.math.vec2, shape=(sx2, sy2)) 
        self._f_ey = ti.field(ti.math.vec2, shape=(sx2, sy2)) 
        self._f_acc = ti.field(ti.math.vec3, shape=(sx2,sy2))

    def upload(self, nn:np.ndarray):
        """
        Upload new director field. Array shape must match specified in the constructor. 
        """
        assert nn.shape==self._f_nn.shape+(3,), f"{nn.shape=} {self._f_nn.shape=}"
        self._f_nn.from_numpy(nn.astype(np.float32))

    def _weight_vector(self, wavelength, emission:Function, efficiency:Function, scale:float=1.0):
        weight = np.asarray(scale*emission(wavelength)*efficiency(wavelength), dtype=np.float32).reshape(3)
        return ti.math.vec3(weight[0], weight[1], weight[2])

    def render(self, wavelengths:np.ndarray, emission:Function, efficiency:Function, ne:Function, no:Function, polarizer:float=0., deltafilter:float=np.pi/2, mono:bool=False):
        """
        wavelength is in micrometers.
        """
        wavelengths = normalize_wavelengths(wavelengths, mono=mono)
        # min_lambda = np.min(wavelengths)
        # if min_lambda/self._dz<8:
        #     print(f"WARNING: Lattice constant {self._dz} is too large for wavelength {min_lambda}. Increase 'subdiv' argument in the constructor.")
        self._clear_output()
        if mono:
            wavelength = wavelengths[0]
            self._process(
                wavelength=wavelength,
                ne=npunwrap(ne(wavelength)),
                no=npunwrap(no(wavelength)),
                polarizer=polarizer,
            )
            self._accumulate(
                weight=self._weight_vector(wavelength, emission, efficiency),
                polarizer=polarizer+deltafilter,
            )
            return self._f_acc.to_numpy()

        ws = self._quadrature.get_weights(wavelengths)
        for w, wavelength in zip(ws, wavelengths):
            self._process(
                wavelength=wavelength, 
                ne=npunwrap(ne(wavelength)),
                no=npunwrap(no(wavelength)),
                polarizer=polarizer,
                )
            self._accumulate(
                weight=self._weight_vector(wavelength, emission, efficiency, scale=w), 
                polarizer=polarizer+deltafilter,
                )
        return self._f_acc.to_numpy()

    def _process(self, wavelength, polarizer, no, ne):
        self._incident_wave(polarizer=polarizer, wavelength=wavelength, no=no, ne=ne)
        self._propagate_through_exy(wavelength, no, ne)

    @ti.kernel
    def _propagate_through_exy(self, wavelength:ti.f32, ne:ti.f32, no:ti.f32):
        # Constants for given wavelength
        dzoverlambda = self._dz/wavelength
        nooverne2 = (no/ne)**2
        sz = self._f_nn.shape[2]
        # Process layer.
        for x,y in self._f_ex:
            ex = self._f_ex[x,y]
            ey = self._f_ey[x,y] 
            for layer in range(1, self._subdiv*(sz-1)+1):
                m = self._interpolate_n_1(x,y,layer)
                l2 = m.x**2+m.y**2
                ll = ti.math.sqrt(l2)+1e-15
                nx = m.x/ll
                ny = m.y/ll
                p = ex*nx+ey*ny
                neff = no/ti.math.sqrt( nooverne2*l2+m.z**2 )
                gammaz = np.pi*dzoverlambda*(neff-no)
                factor = ti.math.cmul(p, ti.math.vec2(1-ti.math.cos(2*gammaz), ti.math.sin(2*gammaz)))
                ex -= factor*nx
                ey -= factor*ny
            self._f_ex[x,y] = ex
            self._f_ey[x,y] = ey
            

    # def _process(self, wavelength, polarizer, no, ne):
    #     # First layer.
    #     self._incident_wave(polarizer=polarizer, wavelength=wavelength, no=no, ne=ne)
    #     self._interpolate_n(0.)
    #     # Process all layers.
    #     sz = self._f_nn.shape[2]
    #     for layer in range(1, self._subdiv*(sz-1)+1):
    #         self._interpolate_n(layer/self._subdiv)
    #         self._propagate_exy(wavelength, no, ne)

    # @ti.kernel
    # def _propagate_exy(self, wavelength:ti.f32, ne:ti.f32, no:ti.f32):
    #     # Constants for given wavelength
    #     dzoverlambda = self._dz/wavelength
    #     nooverne2 = (no/ne)**2
    #     # Process layer/
    #     for x,y in self._f_ex:
    #         m = self._f_nn2[x,y]
    #         l2 = m.x**2+m.y**2
    #         ll = ti.math.sqrt(l2)+1e-15
    #         nx = m.x/ll
    #         ny = m.y/ll
    #         p = self._f_ex[x,y]*nx+self._f_ey[x,y]*ny
    #         neff = no/ti.math.sqrt( nooverne2*l2+m.z**2 )
    #         gammaz = np.pi*dzoverlambda*(neff-no)
    #         factor = ti.math.cmul(p, ti.math.vec2(1-ti.math.cos(2*gammaz), ti.math.sin(2*gammaz)))
    #         self._f_ex[x,y] = self._f_ex[x,y]-factor*nx
    #         self._f_ey[x,y] = self._f_ey[x,y]-factor*ny

    @ti.kernel
    def _clear_output(self):
        for x,y in self._f_acc:
            self._f_acc[x,y] = 0

    @ti.kernel
    def _accumulate(self, weight:ti.math.vec3, polarizer:ti.f32):
        px, py = ti.math.cos(polarizer), ti.math.sin(polarizer)
        # Fields ex and ey are defined at the same points.
        for x, y in self._f_ex:
            self._f_acc[x, y] += weight*ti.math.length( self._f_ex[x,y]*px+self._f_ey[x,y]*py )

    @ti.kernel
    def _incident_wave(self, wavelength:ti.f32, polarizer:ti.f32, no:ti.f32, ne:ti.f32):
        """
        wavelength is in micrometers.
        """
        o = ti.math.vec2([1, 0])
        px, py = ti.math.cos(polarizer), ti.math.sin(polarizer)
        ex, ey = o*px, o*py 
        for x, y in self._f_ex:
            self._f_ex[x, y] = ex
            self._f_ey[x, y] = ey

    @ti.func
    def _interpolate_n_1(self, mx:ti.f32, my:ti.f32, mz:ti.f32) -> ti.math.vec3:
        subdiv = ti.static(self._subdiv)
        sx = ti.static(self._f_nn.shape[0])
        sy = ti.static(self._f_nn.shape[1])
        sz = ti.static(self._f_nn.shape[2])
        # Z
        x, y, z = mx/subdiv, my/subdiv, mz/subdiv
        ix, iy, iz = ti.int32(ti.floor(x)), ti.int32(ti.floor(y)), ti.int32(ti.floor(z))
        lx1, ly1, lz1 = x-ix, y-iy, z-iz
        lx0, ly0, lz0 = 1-lx1, 1-ly1, 1-lz1
        ix1, iy1, iz1 = (ix+1)%sx, (iy+1)%sy, iz+1 if iz<sz-1 else iz
        # Extract values at nearest grid points.
        n000 = self._f_nn[ix ,iy ,iz ]
        n001 = self._f_nn[ix ,iy ,iz1]
        n010 = self._f_nn[ix ,iy1,iz ]
        n011 = self._f_nn[ix ,iy1,iz1]
        n100 = self._f_nn[ix1,iy ,iz ]
        n101 = self._f_nn[ix1,iy ,iz1]
        n110 = self._f_nn[ix1,iy1,iz ]
        n111 = self._f_nn[ix1,iy1,iz1]
        # Interpolate
        na = n000*lx0*ly0*lz0+n001*lx0*ly0*lz1+n010*lx0*ly1*lz0+n011*lx0*ly1*lz1+n100*lx1*ly0*lz0+n101*lx1*ly0*lz1+n110*lx1*ly1*lz0+n111*lx1*ly1*lz1
        return ti.math.normalize(na)

    # @ti.kernel
    # def _interpolate_n(self, lz:ti.f32):
    #     """
    #     z is given in grid size of nn units.
    #     """
    #     subdiv = ti.static(self._subdiv)
    #     sx = ti.static(self._f_nn.shape[0])
    #     sy = ti.static(self._f_nn.shape[1])
    #     sz = ti.static(self._f_nn.shape[2])
    #     iz = ti.int32(ti.floor(lz))
    #     iz1 = iz+1 if iz<sz-1 else iz
    #     lza1 = lz-iz
    #     lza0 = 1-lza1
    #     sx2 = ti.static(self._f_nn2.shape[0])
    #     sy2 = ti.static(self._f_nn2.shape[1])
    #     for mx in range(sx2):
    #         x = mx/(subdiv)
    #         ix = ti.int32(ti.floor(x))
    #         lx1 = x-ix
    #         lx0 = 1-lx1
    #         ix1 = (ix+1)%sx
    #         for my in range(sy2):
    #             y = my/(subdiv)
    #             iy = ti.int32(ti.floor(y))
    #             ly1 = y-iy
    #             ly0 = 1-ly1
    #             iy1 = (iy+1)%sy
    #             # Extract values at nearest grid points.
    #             n000 = self._f_nn[ix,iy,iz]
    #             n001 = self._f_nn[ix,iy,iz1]
    #             n010 = self._f_nn[ix,iy1,iz]
    #             n011 = self._f_nn[ix,iy1,iz1]
    #             n100 = self._f_nn[ix1,iy,iz]
    #             n101 = self._f_nn[ix1,iy,iz1]
    #             n110 = self._f_nn[ix1,iy1,iz]
    #             n111 = self._f_nn[ix1,iy1,iz1]
    #             # Interpolate
    #             na = n000*lx0*ly0*lza0+n001*lx0*ly0*lza1+n010*lx0*ly1*lza0+n011*lx0*ly1*lza1+n100*lx1*ly0*lza0+n101*lx1*ly0*lza1+n110*lx1*ly1*lza0+n111*lx1*ly1*lza1
    #             self._f_nn2[mx,my] = ti.math.normalize(na)



######################################################################################
# Marching wave optics computation.

# @ti.data_oriented
# class MarchingWave:
#     def __init__(self, nn:np.ndarray, thickness:float, subdiv:int, niter:int=5, quadrature:Quadrature=None):
#         self.eps0 = 8.8541878188e-12
#         self.mu0 = 1.25663706127e-6
#         self.c = (self.mu0*self.eps0)**(-0.5)
#         self._niter = niter
#         # Initialize integrator
#         if quadrature is None:
#             quadrature = Trapezoid()
#         self._quadrature = quadrature
#         # Initialize sizes
#         shp = nn if isinstance(nn, tuple) else nn.shape[:3] # Director field shape
#         assert len(shp)==3, f"3D field shape is expected, received {shp}"
#         assert shp[2]>1 
#         self._thickness = thickness
#         self._a = thickness/(shp[2]-1) # Director field grid lattice constant 
#         self._subdiv = int(subdiv)
#         assert self._subdiv>=1
#         self._dz = self._a/self._subdiv
#         # Allocate buffer
#         self._prepare_buffers(shp)
#         if not isinstance(nn,tuple):
#             self.upload(nn)

#     def upload(self, nn:np.ndarray):
#         """
#         Upload new director field. Array shape must match specified in the constructor. 
#         """
#         assert nn.shape==self._f_nn.shape
#         self._f_nn.from_numpy(nn.astype(np.float32))

#     def render(self, wavelengths:np.ndarray, emission:Function, efficiency:Function, ne:Function, no:Function, polarizer:float=0., deltafilter:float=np.pi/2):
#         """
#         wavelength is in micrometers.
#         """
#         min_lambda = np.min(wavelengths)
#         if min_lambda/self._dz<8:
#             print(f"WARNING: Lattice constant {self._dz} is too large for wavelength {min_lambda}. Increase 'subdiv' argument in the constructor.")
#         self._clear_output()
#         ws = self._quadrature.get_weights(wavelengths)
#         for w, wavelength in zip(ws, wavelengths):
#             self._process(
#                 wavelength=wavelength, 
#                 ne=ne(wavelength),
#                 no=no(wavelength),
#                 polarizer=polarizer,
#                 )
#             self._accumulate(
#                 weight=w*emission(wavelength)*efficiency(wavelength), 
#                 polarizer=polarizer+deltafilter
#                 )
#         return self._f_acc.to_numpy()
    
#     ###

#     @ti.kernel
#     def _clear_output(self):
#         for x,y in self._f_acc:
#             self._f_acc[x,y] = 0

#     @ti.kernel
#     def _accumulate(self, weight, polarizer):
#         px, py = ti.math.cos(polarizer), ti.math.sin(polarizer)
#         sx, sy = self._f_ex_1.shape
#         # Field ex is defined at 1/2,0
#         #       ey            at   0,1/2
#         #     |        |
#         # ----#------- ex ----
#         #     |        |
#         # -- ey -------# -----
#         #     |        |   
#         for x, y in self._f_ex_1:
#             xp, xm = (x+1)%sx, (x-1)%sx
#             yp, ym = (y+1)%sy, (y-1)%sy
#             # Point 0,0
#             self._f_acc[2*x, 2*y] += weight*ti.math.length( (self._f_ex_1[x,y]+self._f_ex_1[xm,y])/2*px+(self._f_ey_1[x,y]+self._f_ey_1[x,ym])/2*py )
#             # Point 1/2,0
#             self._f_acc[2*x+1, 2*y] += weight*ti.math.length( self._f_ex_1[x,y]*px+(self._f_ey_1[x,y]+self._f_ey_1[x,ym]+self._f_ey_1[xp,y]+self._f_ey_1[xp,ym])/4*py )
#             # Point 0,1/2
#             self._f_acc[2*x, 2*y+1] += weight*ti.math.length( (self._f_ex_1[x,y]+self._f_ex_1[xm,y]+self._f_ex_1[x,yp]+self._f_ex_1[xm,yp])/2*px+self._f_ey_1[x,y]*py )
#             # Point 1/2,1/2
#             self._f_acc[2*x+1, 2*y+1] += weight*ti.math.length( (self._f_ex_1[x,y]+self._f_ex_1[x,yp])/2*px+(self._f_ey_1[x,y]+self._f_ey_1[xp,y])/2*py )

#     @ti.kernel
#     def _incident_wave(self, wavelength, polarizer, no, ne):
#         """
#         wavelength is in micrometers.
#         """
#         omega = wavelength*1e-6
#         k = omega/self.c*no
#         phi = -k*self._dz/2
#         o = ti.math.vec2([1, 0])
#         s = ti.math.vec2([ti.math.cos(phi), ti.math.sin(phi)])
#         px, py = ti.math.cos(polarizer), ti.math.sin(polarizer)
#         r = k/(omega*self.mu0)
#         ex, ey = o*px, o*py 
#         hx, hy = r*py*s, -r*px*s
#         for x, y in self._f_ex_1:
#             self._f_ex_1[x, y] = ex
#             self._f_ey_1[x, y] = ey
#             self._f_hx_1[x, y] = hx
#             self._f_hy_1[x, y] = hy


#     def _process(self, wavelength, polarizer, no, ne):
#         # First layer.
#         self._interpolate_n(0.)
#         self._incident_wave(polarizer=polarizer, wavelength=wavelength, )
#         # Process all layers.
#         sz = self._f_nn.shape[2]
#         for layer in range(1, self._subdiv*(sz-1)+1):
#             self._swap()
#             self._interpolate_n(layer/self._subdiv)
#             self._initial_exy()
#             for _ in range(self._niter):
#                 self._propagate_ez(wavelength, no, ne)
#                 self._propagate_exy(wavelength, no, ne)
#             self._propagate_hz(wavelength, no, ne)
#             self._propagate_hxy(wavelength, no, ne)



#     def _prepare_buffers(self, shp):
#         sx, sy, sz = shp
#         sx2, sy2 = sx*self._subdiv, sy*self._subdiv
#         self._f_nn = ti.field(ti.math.vec3, shape=(sx, sy, sz))

#         self._f_nn2_0 = ti.field(ti.math.vec3, shape=(2*sx2, 2*sy2, 2)) # Two layers on distance self._dz/2 of 3d director vectors.
#         self._f_ex_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1
#         self._f_ey_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1
#         self._f_ez_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1
#         self._f_hx_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1
#         self._f_hy_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1
#         self._f_hz_0 = ti.field(ti.math.vec2, shape=(sx2, sy2)) # C1

#         self._f_nn2_1 = field_like(self._f_nn2_0)
#         self._f_ex_1  = field_like(self._f_ex_0)
#         self._f_ey_1  = field_like(self._f_ey_0)
#         self._f_ez_1  = field_like(self._f_ez_0)
#         self._f_hx_1  = field_like(self._f_hx_0)
#         self._f_hy_1  = field_like(self._f_hy_0)
#         self._f_hz_1  = field_like(self._f_hz_0)

#         self._f_acc = ti.field(ti.float32, shape=(2*sx2, 2*sy2)) # Accumulated intensity.

#     @ti.kernel
#     def _interpolate_n(self, z):
#         """
#         z is given in grid size of nn units.
#         """
#         subdiv = ti.static(self._subdiv)
#         sx = ti.static(self._f_nn.shape[0])
#         sy = ti.static(self._f_nn.shape[1])
#         sz = ti.static(self._f_nn.shape[2])
#         iz = ti.int32(ti.floor(z))
#         iz1 = iz+1 if iz<sz-1 else iz
#         lza1 = z-iz
#         lza0 = 1-lza1
#         lzb1 = lza1+0.5/subdiv
#         lzb0 = 1-lzb1
#         sx2 = ti.static(self._f_nn2_1.shape[0])
#         sy2 = ti.static(self._f_nn2_1.shape[1])
#         for mx in range(sx2):
#             x = mx/(2*sx2)
#             ix = ti.int32(ti.floor(x))
#             lx1 = x-ix
#             lx0 = 1-lx1
#             ix1 = (ix+1)%sx
#             for my in range(sy2):
#                 y = my/(2*sy2)
#                 iy = ti.int32(ti.floor(y))
#                 ly1 = y-iy
#                 ly0 = 1-ly1
#                 iy1 = (iy+1)%sy
#                 # Extract values at nearest grid points.
#                 n000 = self._f_nn[ix,iy,iz]
#                 n001 = self._f_nn[ix,iy,iz1]
#                 n010 = self._f_nn[ix,iy1,iz]
#                 n011 = self._f_nn[ix,iy1,iz1]
#                 n100 = self._f_nn[ix1,iy,iz]
#                 n101 = self._f_nn[ix1,iy,iz1]
#                 n110 = self._f_nn[ix1,iy1,iz]
#                 n111 = self._f_nn[ix1,iy1,iz1]
#                 # Interpolate
#                 na = n000*lx0*ly0*lza0+n001*lx0*ly0*lza1+n010*lx0*ly1*lza0+n011*lx0*ly1*lza1+n100*lx1*ly0*lza0+n101*lx1*ly0*lza1+n110*lx1*ly1*lza0+n111*lx1*ly1*lza1
#                 nb = n000*lx0*ly0*lzb0+n001*lx0*ly0*lzb1+n010*lx0*ly1*lzb0+n011*lx0*ly1*lzb1+n100*lx1*ly0*lzb0+n101*lx1*ly0*lzb1+n110*lx1*ly1*lzb0+n111*lx1*ly1*lzb1
#                 self._f_nn2_1[mx,my,0] = ti.math.normalize(na)
#                 self._f_nn2_1[mx,my,1] = ti.math.normalize(nb)

#     def _swap(self):
#         self._f_inn_0, self._f_inn_1 = self._f_inn_1, self._f_inn_0
#         self._f_ex_0, self._f_ex_1 = self._f_ex_1, self._f_ex_0
#         self._f_ey_0, self._f_ey_1 = self._f_ey_1, self._f_ey_0
#         self._f_ez_0,  self._f_ez_1  = self._f_ez_1,  self._f_ez_0 
#         self._f_hx_0, self._f_hx_1 = self._f_hx_1, self._f_hx_0
#         self._f_hy_0, self._f_hy_1 = self._f_hy_1, self._f_hy_0
#         self._f_hz_0,  self._f_hz_1  = self._f_hz_1,  self._f_hz_0 

#     @ti.kernel
#     def _initial_exy(self):
#         # Copy previous state.
#         for x, y in self._f_ex_1:
#             self._f_ex_1[x,y] = self._f_ex_0[x,y]
#         for x, y in self._f_ey_1:
#             self._f_ey_1[x,y] = self._f_ey_0[x,y]

#     @ti.kernel
#     def _propagate_ez(self, omega, no, ne):
#         eps0 = self.eps0
#         eps_o = eps0*no**2
#         eps_d = eps0*(ne**2-no**2)
#         sx, sy = self._f_ex_0.shape
#         ia = ti.math.vec2(0,1/self._dz*omega)
#         for x,y in self._f_ez_1:
#             xm = (x-1)%sx
#             ym = (y-1)%sy
#             n = self._f_nn2_0[2*x, 2*y, 1]
#             eps_zx = eps_o + eps_d*n[2]*n[0] 
#             eps_zy = eps_o + eps_d*n[2]*n[1]
#             eps_zz = eps_o + eps_d*n[2]*n[2]
#             self._f_ez_1[x,y]=(
#                 ti.math.cmul(ia, 
#                     (self._f_hy_0[x,y]-self._f_hy_0[xm,y])-
#                     (self._f_hx_0[x,y]-self._f_hx_0[x,ym])
#                 ) 
#                 + eps_zx*(self._f_ex_0[x,y]+self._f_ex_1[x,y]+self._f_ex_0[xm,y]+self._f_ex_1[xm,y])/4
#                 + eps_zy*(self._f_ey_0[x,y]+self._f_ey_1[x,y]+self._f_ey_0[x,ym]+self._f_ey_1[x,ym])/4
#             )/eps_zz

#     @ti.kernel        
#     def _propagate_hxy(self, omega, no, ne):
#         eps0 = self.eps0
#         eps_o = eps0*no**2
#         eps_d = eps0*(ne**2-no**2)
#         sx, sy = self._f_ex_0.shape
#         a = ti.math.vec2(0,self._dz*omega)
#         for x,y in self._f_hx_1:
#             xm = (x-1)%sx
#             ym = (y-1)%sy
#             xp = (x+1)%sx
#             yp = (y+1)%sy

#             nx = self._f_nn2_0[2*x, 2*y+1, 1]
#             eps_xx = eps_o + eps_d*nx[2]*nx[0] 
#             eps_xy = eps_o + eps_d*nx[2]*nx[1]
#             eps_xz = eps_o + eps_d*nx[2]*nx[2]

#             ny = self._f_nn2_0[2*x+1, 2*y, 1]
#             eps_yx = eps_o + eps_d*ny[2]*ny[0] 
#             eps_yy = eps_o + eps_d*ny[2]*ny[1]
#             eps_yz = eps_o + eps_d*ny[2]*ny[2]

#             self._f_hx_1[x,y]=(
#                 self._f_hx_0[x,y]+(self._f_hz_1[x,y]-self._f_hz_1[xm,y])
#                 -ti.math.cmul(a,
#                     + eps_yx*(self._f_ex_1[x,y]+self._f_ex_1[xm,y]+self._f_ex_1[x,yp]+self._f_ex_1[xm,yp])/4
#                     + eps_yy*(self._f_ey_1[x,y])
#                     + eps_yz*(self._f_ez_1[x,y]+self._f_ez_1[x,yp]+self._f_ez_0[x,y]+self._f_ez_0[x,yp])/4
#                 ) 
#             )
#             self._f_hy_1[x,y]=(
#                 self._f_hy_0[x,y]+(self._f_hz_1[x,y]-self._f_hz_1[x,ym])
#                 +ti.math.cmul(a,
#                     + eps_xx*(self._f_ex_1[x,y])
#                     + eps_xy*(self._f_ey_1[x,y]+self._f_ey_1[xp,y]+self._f_ey_1[x,ym]+self._f_ey_1[xp,ym])/4
#                     + eps_xz*(self._f_ez_1[x,y]+self._f_ez_1[xp,y]+self._f_ez_0[x,y]+self._f_ez_0[xp,y])/2
#                 ) 
#             )


#     @ti.kernel        
#     def _propagate_exy(self, omega, no, ne):
#         sx, sy = self._f_ex_0.shape
#         b = ti.math.vec2(0,self._dz*omega*self.mu0)
#         for x,y in self._f_ex_1:
#             xp = (x+1)%sx
#             yp = (y+1)%sy
#             self._f_ey_1[x,y]=(
#                 -ti.math.cmul(b,self._f_hx_0[x,y])
#                 +self._f_ey_0[x,y]
#                 -(self._f_ez_0[x,yp]-self._f_ez_0[x,y])
#             )
#             self._f_ex_1[x,y]=(
#                 ti.math.cmul(b,self._f_hy_0[x,y])
#                 +self._f_ex_0[x,y]
#                 -(self._f_ez_0[xp,y]-self._f_ez_0[x,y])
#             )
            

#     @ti.kernel        
#     def _propagate_hz(self, omega, no, ne):
#         sx, sy = self._f_ex_0.shape
#         ib = ti.math.vec2(0,-1/(self._dz*omega*self.mu0))
#         for x,y in self._f_hz_1:
#             xp = (x+1)%sx
#             yp = (y+1)%sy
#             self._f_hz_1[x,y]=ti.math.cmul(ib, 
#                 (self._f_ey_0[xp,y]-self._f_ey_0[x,y])
#                 -(self._f_ex_0[x,yp]-self._f_ex_0[x,y])
#             )




