import numpy as np
from scipy.special import comb
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

m = 10
C = 10
r = 2
def P(x, y, z, m=10, C=10, r=2):
    """
    Function to calculate P(x, y, z)
    """
    if x < m + C - r or y < m - 2 or y > m + 2 or z < r - 2 or z > r + 2:
        return 0  # Outside the domain, return 0
    p = (comb(y, z) * comb(x - y, C - z) * z * (z - 1) / (comb(x, C) * y * C)
        * (-0.25 * abs(y - m) + 0.5) * (-0.25 * abs(z - r) + 0.5))
    return p


# Define domain ranges and step size
x_min, x_max = m + C - r, 200
y_min, y_max = m - 2, m + 2
z_min, z_max = r - 2, r + 2
step = 0.5

# Create meshgrid
X, Y, Z = np.mgrid[x_min:x_max:step, y_min:y_max:step, z_min:z_max:step]

# Calculate P(x, y, z) values
P_values = np.vectorize(P)(X, Y, Z)

# Remove zero values
X, Y, Z, P_values = X[P_values != 0], Y[P_values != 0], Z[P_values != 0], P_values[P_values != 0]

# Create 3D scatter plot with color representing P(x, y, z)
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
cmap = plt.cm.viridis  # Choose a colormap for visualization
scatter = ax.scatter(X, Y, Z, c=P_values, cmap=cmap, s=1)
fig.colorbar(scatter, label="$P(N_x|m_y,C,r_z)$")

# Set labels and title
ax.set_xlabel("$N_x$")
ax.set_ylabel("$m_y$")
ax.set_zlabel("$r_z$")
plt.title("Four Dimensional Probability Distribution")

# Show the plot
plt.show()