#!/usr/bin/env python3

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

plt.rcParams.update({
    "font.family": "serif",
    "font.size": 12
})

def generate_surface():

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(111, projection='3d')

    x = np.linspace(-1.1,1.1,350)
    y = np.linspace(-1.1,1.1,350)

    X,Y = np.meshgrid(x,y)

    R2 = X**2 + Y**2

    # instability boundary
    mask = R2 < 1

    Z = np.zeros_like(X)

    Z[mask] = 1/(1 - R2[mask])

    # gentle saturation for plotting
    Z = np.tanh(Z/4)*4

    # mask outside region
    Z[~mask] = np.nan

    surf = ax.plot_surface(
        X,Y,Z,
        cmap="plasma",
        linewidth=0,
        antialiased=True
    )

    # stability boundary circle
    theta = np.linspace(0,2*np.pi,500)
    xb = np.cos(theta)
    yb = np.sin(theta)

    ax.plot(xb,yb,0,color="black",linewidth=2.5)

    # contours on base plane
    ax.contour(
        X,Y,Z,
        levels=12,
        offset=0,
        cmap="plasma",
        linewidths=0.6
    )

    ax.set_xlabel(r"$dI/dt$")
    ax.set_ylabel(r"$dB/dt$")
    ax.set_zlabel("Rotational response")

    ax.set_title("MTDE Torque–Inertia Stability Surface")

    ax.view_init(elev=30, azim=32)

    plt.tight_layout()

    plt.savefig("mtde_stability_surface.pdf")

    plt.close()

if __name__ == "__main__":
    generate_surface()