import numpy as np
import matplotlib.pyplot as plt


def plotCCD(elev=20, azim=30, figsize=(10, 8), filename=None, title="Central Composite Design (CCD) for k=3") -> None:
    """
    Plots a Central Composite Design (CCD) for k=3 with customizable viewpoint.

    Args:
        elev (int):
            Elevation angle in the z-plane for the 3D plot.
        azim (int):
            Azimuthal angle in the x,y-plane for the 3D plot.
        figsize (tuple):
            Size of the figure (width, height) in inches. Defaults to (10, 8).
        filename (str, optional):
            The name of the file to save the plot. If None, the plot will be shown instead. Supported formats: 'pdf', 'png'.
        title (str):
            Title of the plot. Defaults to "Central Composite Design (CCD) for k=3".

    Returns:
        None

    Examples:
        ```{python}
        from spotdesirability.plot.ccd import plotCCD
        import matplotlib.pyplot as plt
        # Simple plot
        plotCCD(elev=20, azim=30, title="CCD Plot Example")
        ```
    """
    # Define the number of factors (k) and axial distance (alpha)
    k = 3
    alpha = np.sqrt(k)  # Rotatable CCD

    # Generate factorial points (2^k design)
    factorial_points = np.array([[x1, x2, x3] for x1 in [-1, 1] for x2 in [-1, 1] for x3 in [-1, 1]])

    # Generate axial (star) points
    axial_points = []
    for i in range(k):
        point_positive = [0] * k
        point_negative = [0] * k
        point_positive[i] = alpha
        point_negative[i] = -alpha
        axial_points.append(point_positive)
        axial_points.append(point_negative)
    axial_points = np.array(axial_points)

    # Center point
    center_point = np.array([[0, 0, 0]])

    # Plot the CCD
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(111, projection="3d")

    # Plot factorial points
    ax.scatter(factorial_points[:, 0], factorial_points[:, 1], factorial_points[:, 2], c="blue", label="Factorial Points (2^k)", s=50)

    # Plot axial points
    ax.scatter(axial_points[:, 0], axial_points[:, 1], axial_points[:, 2], c="red", label="Axial Points (±α)", s=50)

    # Plot center point
    ax.scatter(center_point[:, 0], center_point[:, 1], center_point[:, 2], c="green", label="Center Point", s=100)

    # Connect edges of the cube (factorial points)
    for i in range(len(factorial_points)):
        for j in range(i + 1, len(factorial_points)):
            # Check if points differ by only one coordinate
            if np.sum(np.abs(factorial_points[i] - factorial_points[j])) == 2:
                ax.plot(
                    [factorial_points[i, 0], factorial_points[j, 0]], [factorial_points[i, 1], factorial_points[j, 1]], [factorial_points[i, 2], factorial_points[j, 2]], color="black", linewidth=0.5
                )

    # Add axes through the origin
    ax.plot([-alpha, alpha], [0, 0], [0, 0], color="gray", linestyle="--", label="X1 Axis")
    ax.plot([0, 0], [-alpha, alpha], [0, 0], color="gray", linestyle="--", label="X2 Axis")
    ax.plot([0, 0], [0, 0], [-alpha, alpha], color="gray", linestyle="--", label="X3 Axis")

    # Set plot labels
    ax.set_xlabel("X1")
    ax.set_ylabel("X2")
    ax.set_zlabel("X3")
    if title is not None:
        ax.set_title(title)

    # Set the viewpoint
    ax.view_init(elev=elev, azim=azim)

    # Add legend outside the plot
    ax.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0)
    plt.tight_layout()
    # Save or show the plot
    if filename:
        plt.savefig(filename, bbox_inches="tight")
        plt.close()
    else:
        plt.show()
