"""
DensitySpace.VoxelSpace.py
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def get_meshgrid(N):
    """
    Create a meshgrid for the electron density space.

    Parameters
    ----------
    N : int
        The size of the grid along each axis.

    Returns
    -------
    tuple
        A tuple containing the meshgrid arrays (x, y, z).
    """
    x = y = z = np.arange(N)
    return np.meshgrid(x, y, z)

def draw_voxles_as_dots_impl(ax, image, cmap=cm.plasma, min_value=0, colorbar=False, limits=None, alpha=1, s=30):
    """
    Draw voxels as dots in a 3D plot.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        The axes to plot on.
    image : np.ndarray
        The 3D density array.
    min_value : float
        The minimum density value to consider for plotting.

    colorbar : bool
        Whether to display a colorbar.
    limits : tuple or None
        The limits for the axes in the form (xlim, ylim, zlim). If
        None, the limits will be set to the size of the image.
    alpha : float
        The alpha value for the dots.
    s : float
        The size of the dots.

    Returns
    -------
    None
    """
    
    w = np.where(image > min_value)
    wi = np.array(w, dtype=int).T
    xyz = wi
    v = image[w]
    sc = ax.scatter3D(xyz[:,0], xyz[:,1], xyz[:,2], c=v, cmap=cmap, alpha=alpha, s=s)
    if colorbar:
        ax.get_figure().colorbar(sc, ax=ax)
    if limits is None:
        ax.set_xlim(0, image.shape[0])
        ax.set_ylim(0, image.shape[1])
        ax.set_zlim(0, image.shape[2])

def add_center_plane(ax, N, k=None, alpha=0.2, color='gray'):
    """
    Add a plane at z=k to the 3D axes.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        The axes to plot on.
    N : int
        The size of the grid along each axis.
    k : int or None
        The z value for the plane. If None, it will be set to N//2
    alpha : float
        The alpha value for the plane.
    """
    if k is None:
        k = N // 2
    # Create grid for the plane
    xx, yy = np.meshgrid(np.arange(N), np.arange(N))
    zz = np.full_like(xx, k)
    ax.plot_surface(xx, yy, zz, color=color, alpha=alpha, edgecolor='none')

def compute_rg_impl(rho, min_value=0):
    """
    Compute the radius of gyration for a 3D density array.
    (This code was generated by Copilot)

    Parameters
    ----------
    rho : np.ndarray
        3D density array.
    min_value : float
        Minimum density value to include in calculation.

    Returns
    -------
    float
        Radius of gyration.
    """
    # Get indices of voxels above threshold
    w = np.where(rho > min_value)
    x, y, z = w
    m = rho[w]
    total_mass = np.sum(m)
    if total_mass == 0:
        return 0.0

    # Center of mass
    x_cm = np.sum(x * m) / total_mass
    y_cm = np.sum(y * m) / total_mass
    z_cm = np.sum(z * m) / total_mass

    # Squared distances from center of mass
    r2 = (x - x_cm)**2 + (y - y_cm)**2 + (z - z_cm)**2
    rg2 = np.sum(m * r2) / total_mass
    return np.sqrt(rg2)

class VoxelSpace:
    """
    VoxelSpace class to handle voxel-based density spaces.
    It can be initialized with a shape or a density array.
    
    Attributes
    ----------
    rho : np.ndarray
        3D density array.
    """
    def __init__(self, N, shape=None, density=None):
        self.rho = np.zeros((N, N, N))
        if shape is None:
            if density is None:
                raise ValueError("Either shape or density must be provided.")
            else:
                self.rho = density
        else:
            xx, yy, zz = get_meshgrid(N)
            shape_condition = shape.get_condition(xx, yy, zz)
            self.rho[shape_condition] = 1

    def plot_as_dots(self, ax=None, cmap=cm.plasma, min_value=0, colorbar=False, alpha=1, s=30):
        """
        Plot the voxel space.

        Parameters
        ----------
        ax : matplotlib.axes.Axes
            The axes to plot on.
        min_value : float
            The minimum density value to consider for plotting.
        """
        # from learnsaxs import draw_voxles_as_dots
        if ax is None:
            fig, ax = plt.subplots(subplot_kw={'projection': '3d'})

        draw_voxles_as_dots_impl(ax, self.rho, cmap=cmap, min_value=min_value, colorbar=colorbar, alpha=alpha, s=s)

    def plot_with_density(self, axes=None, cmap=cm.plasma, min_value=1e-4, colorbar=True, alpha=0.3, s=1, k=None):
        """
        Plot the voxel space.

        Parameters
        ----------
        axes : matplotlib.axes.Axes
            The axes to plot on.
        min_value : float
            The minimum density value to consider for plotting.
        colorbar : bool
            Whether to display a colorbar.
        alpha : float
            The alpha value for the plots.
        s : float
            The size of the dots.
        k : int or None
            The z value for the cross-sectional plane. If None, it will be set to N//2

        Returns
        -------
        None
        """
        from molass.PlotUtils.MatrixPlot import simple_plot_3d
        if axes is None:
            fig, axes = plt.subplots(ncols=2, figsize=(10, 5), subplot_kw={'projection': '3d'})
        ax1, ax2 = axes
        ax1.set_title("3D View of the Density")
        draw_voxles_as_dots_impl(ax1, self.rho, cmap=cmap, min_value=min_value, colorbar=colorbar, alpha=alpha, s=s)
        N = self.rho.shape[0]
        if k is None:
            k = N // 2
        add_center_plane(ax1, N, k=k, alpha=alpha)

        ax2.set_title(f"Density on the Cross Section with z={k}")
        simple_plot_3d(ax2, self.rho[:,:,k], cmap=cmap, colorbar=colorbar, alpha=alpha)

    def compute_rg(self, min_value=0):
        """
        Compute the radius of gyration of the voxel space.

        Parameters
        ----------
        min_value : float
            Minimum density value to include in calculation.

        Returns
        -------
        float
            The radius of gyration.
        """
        return compute_rg_impl(self.rho)