class VoxelSphere:
    mH = 1.00784  
    mHe = 4.002602  
    R = 8.3145e7
    R = 1

    _3DDiffKernels = make_3d_kernels()

    def __init__(self, radius, resolution, t0=0, X=0.75, Y=0.25, Z=0, D=1, k=1):
        self._radius = radius
        self._resolution = resolution
        self._t = t0
        self._X, self._Y, self._Z = X, Y, Z
        self._effectiveMolarMass = self._X * self.mH + self._Y * self.mHe

        self._create_voxel_sphere()
        
        self._D = D
        self._k = k
        
    @property
    def radius(self):
        return self._radius

    @property
    def resolution(self):
        return self._resolution

    def _create_voxel_sphere(self):
        gridSize = np.ceil(self.radius * self.resolution) * 2 + 1
        
        self.x = np.linspace(-self.radius, self.radius, int(gridSize))
        self.y = np.linspace(-self.radius, self.radius, int(gridSize))
        self.z = np.linspace(-self.radius, self.radius, int(gridSize))
        xGrid, yGrid, zGrid = np.meshgrid(self.x, self.y, self.z, indexing='ij')
        self._dr = np.linalg.norm(np.array([xGrid[0,0,0] - xGrid[0,1,0], yGrid[0,0,0] - yGrid[1,0,0], zGrid[0,0,0] - zGrid[0,0,1]]))
        self._voxelVolume = np.abs((self.x[1] - self.x[0]) * (self.y[1] - self.y[0]) * (self.z[1] - self.z[0]))
        self._voxelFaceArea = np.abs((self.x[1] - self.x[0]) * (self.y[1] - self.y[0]))

        self._temperature = CallbackNDArray(np.ones_like(xGrid), callback=self._recompute_energy)
        self._density = CallbackNDArray(np.ones_like(xGrid), callback=self._recompute_energy)
        self._pressure = CallbackNDArray(np.ones_like(xGrid), callback=self._recompute_energy)
        
        distanceFromCenter = np.sqrt(xGrid**2 + yGrid**2 + zGrid**2)
        self._voxelSphere = distanceFromCenter <= self.radius
        self._recompute_energy()        

    def _recompute_energy(self):
        self._n = (self._pressure * self._voxelVolume)/self._effectiveMolarMass
        self._energy = (3/2) * self._n * self.R * self._temperature
        self._energy[~self._voxelSphere] = 0

    def _recompute_temperature(self):
        self._n = (self._pressure * self._voxelVolume)/self._effectiveMolarMass
        self._temperature = (2/3) * self._energy/(self._n*self.R)
        self._temperature[~self._voxelSphere] = 0

    @property
    def temperature(self):
        return self._temperature

    def _temperature_differences(self):
        diffs = np.zeros(shape=(26, *self._temperature.shape))
        for kernelID, kernel in enumerate(self._3DDiffKernels):
            correlation = correlate(self._temperature, kernel, mode='constant', cval=0.0)
            diffs[kernelID] = correlation
        return diffs

    def timestep(self, dt, ect = 0.01):
        totalEnergyInit = self._energy.sum()
        dTs = self._temperature_differences()
        Q = np.zeros_like(dTs)
        for orientationID, dT in enumerate(dTs): 
            orientationQ = -self._k*self._voxelFaceArea*dT/(self._dr)
        self._energy = self._energy + dE
        pretty_print_3d_array(self._energy)
        totalEnergyFinal = self._energy.sum()
        fchange = abs((totalEnergyFinal - totalEnergyInit)/totalEnergyInit)
        if fchange > ect:
            raise EnergyConservationError(f"Total Energy varied more than allowed minimum fraction ({ect*100:0.3f}%) {fchange*100:0.3f}%")
        self._recompute_temperature()
        
    def plot_voxel_region(self, voxelSphere, voxelColors='red', **kwargs):
        fig, ax = plt.subplots(1, 1, figsize=(10, 7), subplot_kw=dict(projection='3d'))
        
        filled = np.zeros(voxelSphere.shape, dtype=bool)
        filled[np.where(voxelSphere)] = True
        ax.voxels(filled, facecolors=voxelColors, edgecolor='k', **kwargs)
        
        ax.set_xlabel('X Axis')
        ax.set_ylabel('Y Axis')
        ax.set_zlabel('Z Axis')
        
        return fig, ax

    def _mk_voxel_color_array(self, array):
        Norm = Normalize(vmin=array.min(), vmax=array.max())
        colors = viridis(Norm(array).flatten()).reshape(*array.shape, 4)
        return colors

    def _pcparam_to_color(self, pcparam):
        if pcparam == None:
            colors = np.zeros(shape=(*self._voxelSphere.shape, 4))
            colors[:,:,:] = [1,1,1,1]
        elif pcparam == 'temperature':
            colors = self._mk_voxel_color_array(self._temperature)
        elif pcparam == 'pressure':
            colors = self._mk_voxel_color_array(self._pressure)
        elif pcparam == 'energy':
            colors = self._mk_voxel_color_array(self._energy)
        return colors

    def plot_voxel_sphere(self, pcparam=None, cmap='viridis', **kwargs):
        colors = self._pcparam_to_color(pcparam)
        return self.plot_voxel_region(self._voxelSphere, voxelColors=colors, **kwargs)