%%bash
export LC_ALL=C.UTF-8
export LANG=C.UTF-8
pip uninstall -y cupy cupy-cuda11x cupy-cuda12x

# !pip uninstall -y cupy cupy-cuda11x cupy-cuda12x

!pip install cupy-cuda12x

!nvidia-smi

# CPU version (you can skip this cell and directly run the GPU one first if you'd like)
# The script below will run a Monte Carlo simulation of an ion beam etching process on palladium
# The palladium is set to a thickness of 50 nm

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
from mpl_toolkits.mplot3d import Axes3D  # for 3D plotting
import cv2
import matplotlib as mpl
from scipy.ndimage import gaussian_filter  # for smoothing

# Increase plot quality and animation embed limit for Colab.
mpl.rcParams['figure.dpi'] = 100
mpl.rcParams['animation.embed_limit'] = 100


# --- Ion Beam Etching Parameters ---

incidence_angle = 75          # degrees from surface normal
ion_beam_voltage = 200        # V
ion_beam_current = 70         # mA (slow recipe)
acceleration_voltage = 24     # V
argon_flow = 13               # sccm
palladium_thickness = 50      # nm (film to be etched)

# Sputter yield parameters (simplified, per-ion approach)
yield_peak = 0.3              # nm removal at the center of the impact per ion (for Pd)
sigma = 2.0                   # nm, baseline width of the removal profile
# Effective sigma to mimic the grazing (75°) incidence (wider impact footprint)
sigma_eff = sigma / np.cos(np.radians(incidence_angle))


# --- Simplified Time-Stepped Etch Model ---

etch_time = 600               # total etching time in seconds
vert_rate_nm_per_s = 0.083    # nm/s (example rate; 5 nm/min)

# --- Conversions and grid setup ---
pixel_size_nm = 5.0           # each pixel represents 5 nm in the simulation grid
pixel_um_conv = 1e-3 / (pixel_size_nm)  # conversion factor: nm to µm
n_points = 600                # grid size in pixels

# Precompute smoothing sigmas in pixel units.
# We assume the beam direction is along the x-axis:
sigma_x_pixels = sigma_eff / pixel_size_nm   # larger spread along beam direction
sigma_y_pixels = sigma / pixel_size_nm         # baseline spread perpendicular

# --- Create Circular Mask (Protected Region) ---

circle_mask = np.zeros((n_points, n_points), dtype=np.uint8)
center = (n_points // 2, n_points // 2)
radius = n_points // 3
cv2.circle(circle_mask, center, radius, 255, thickness=-1)
mask_binary = (circle_mask > 0).astype(int)

# Create an RGB overlay for visualization
rgb_mask = np.stack((circle_mask,) * 3, axis=-1)
contours, _ = cv2.findContours(circle_mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(rgb_mask, contours, -1, (0, 255, 0), 2)


# --- Simulation Grid ---

total_size_nm = n_points * pixel_size_nm
total_size_um = total_size_nm * 1e-3
x_axis = np.linspace(0, total_size_um, n_points)
y_axis = np.linspace(0, total_size_um, n_points)
xv, yv = np.meshgrid(x_axis, y_axis)


# --- Time Steps ---
t_start = 0
t_end = etch_time
t_step = 5                    # seconds per simulation frame
time_steps = np.arange(t_start, t_end, t_step)

# Initialize topography (etch depth) in nm; 0 means unetched.
topo = np.zeros((len(time_steps), n_points, n_points))


# --- Etch Depth Evolution with Smoothing ---
for i, t in enumerate(time_steps):
    if i == 0:
        topo[i] = 0.0  # initial flat surface
    else:
        # Calculate vertical etch per time step (nm)
        delta_nm = vert_rate_nm_per_s * t_step
        # Etch only outside the mask: mask_binary==1 protects the film.
        delta_masked = delta_nm * (1 - mask_binary)
        # Update etch depth (accumulating negative values)
        topo[i] = topo[i - 1] - delta_masked
        # Apply an anisotropic Gaussian smoothing to simulate the taper effect.
        # The smoothing is stronger in the x-direction (beam direction).
        topo[i] = gaussian_filter(topo[i], sigma=(sigma_y_pixels, sigma_x_pixels))

# Convert etch depth from nm to µm for plotting (1 nm = 0.001 µm)
topo_um = topo * 1e-3

# Define color scale limits: maximum etch depth expected.
max_etch_um = (vert_rate_nm_per_s * etch_time) * 1e-3
vmin = -max_etch_um
vmax = 0.0
cmap = 'gnuplot'


# --- Visualization: 3D Surface + 2D Contour Animation ---
fig = plt.figure(figsize=(12, 6))
ax3d = fig.add_subplot(121, projection='3d')
ax2d = fig.add_subplot(122)

def update(frame_index):
    ax3d.cla()
    ax2d.cla()

    # 3D Surface Plot
    surf = ax3d.plot_surface(xv, yv, topo_um[frame_index],
                             cmap=cmap, vmin=vmin, vmax=vmax,
                             rstride=2, cstride=2, linewidth=0, antialiased=False)
    ax3d.set_zlim(vmin, vmax)
    ax3d.set_xlabel("X (µm)")
    ax3d.set_ylabel("Y (µm)")
    ax3d.set_zlabel("Etch Depth (µm)")
    elapsed_time = t_start + frame_index * t_step
    ax3d.set_title(f"3D Etch Profile at t = {elapsed_time} s\n"
                   f"(Ion Voltage={ion_beam_voltage} V, Current={ion_beam_current} mA)")
    ax3d.view_init(elev=30, azim=frame_index)  # dynamic rotation for visual effect

    # 2D Contour Plot
    cont = ax2d.contourf(xv, yv, topo_um[frame_index],
                         500, cmap=cmap, vmin=vmin, vmax=vmax)
    ax2d.set_title("Top-view Etch Contour")
    ax2d.set_xlabel("X (µm)")
    ax2d.set_ylabel("Y (µm)")
    # Overlay the mask (semi-transparent)
    ax2d.imshow(mask_binary, extent=[0, total_size_um, 0, total_size_um],
                origin='lower', cmap='gray', alpha=0.3)

    return surf,

anim = FuncAnimation(fig, update, frames=topo_um.shape[0], interval=200, blit=False)
HTML(anim.to_jshtml())

# Look for the animation play button after running this script in a Jupyter Notebook