import os
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # needed for 3D plotting
from matplotlib import rcParams

# -------------------------------
# Parametric surface functions
# -------------------------------
def x_func(t, u):
    """Parametric x(t,u)"""
    return 4 * t * np.cos(u)

def y_func(t, u):
    """Parametric y(t,u)"""
    return 4 * t * np.sin(u)

def z_func(t, u):
    """Parametric z(t,u)"""
    return 10 * t

def z_func_neg(t, u):
    """Optional lower branch"""
    return -10 * t

# -------------------------------
# Generate mesh
# -------------------------------
def generate_mesh(t_range=(0, 5), u_range=(0, 2*np.pi), n_t=200, n_u=400):
    """High-resolution parametric mesh"""
    t = np.linspace(t_range[0], t_range[1], n_t)
    u = np.linspace(u_range[0], u_range[1], n_u)
    T, U = np.meshgrid(t, u)
    X = x_func(T, U)
    Y = y_func(T, U)
    Z1 = z_func(T, U)
    Z2 = z_func_neg(T, U)
    return X, Y, Z1, Z2

# -------------------------------
# Plot surface with high-quality settings
# -------------------------------
def plot_surface():
    # Generate high-res mesh
    X, Y, Z1, Z2 = generate_mesh()

    # Create figure with high DPI
    fig = plt.figure(figsize=(12, 10), dpi=600)
    ax = fig.add_subplot(111, projection="3d")

    # Smooth shading, anti-aliasing
    surf1 = ax.plot_surface(
        X, Y, Z1,
        cmap="viridis",
        edgecolor="none",
        rstride=1,
        cstride=1,
        antialiased=True,
        linewidth=0,
        alpha=0.95
    )

    surf2 = ax.plot_surface(
        X, Y, Z2,
        cmap="viridis",
        edgecolor="none",
        rstride=1,
        cstride=1,
        antialiased=True,
        linewidth=0,
        alpha=0.95
    )

    # Axis labels and title
    ax.set_xlabel("X", labelpad=15)
    ax.set_ylabel("Y", labelpad=15)
    ax.set_zlabel("Z", labelpad=15)
    ax.set_title("Parametric Surface", pad=20)

    # Optional: adjust viewing angle
    ax.view_init(elev=30, azim=45)

    # Add colorbar
    fig.colorbar(surf1, shrink=0.5, aspect=10, pad=0.1)

    # Save at maximum quality
    parent_dir = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
    filename_png = f"{parent_dir}.png"

    plt.savefig(filename_png, format="png", dpi=600, bbox_inches='tight')  # ultra-high-res raster
    print(f"and PNG as {filename_png}")

    # Show interactive window (optional)

# -------------------------------
# Main
# -------------------------------
if __name__ == "__main__":
    plot_surface()