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

def christmas_light_length(r, h, n):
    rho = r / h 
    L1 = (h / 2) * np.sqrt(1 + rho**2 * (1 + 4 * np.pi**2 * n**2))
    L2 = (h * (1 + rho**2)) / (4 * np.pi * n * rho) * np.arcsinh((2 * np.pi * n * rho) / np.sqrt(1 + rho**2))
    return L1 + L2
    
R_max = 5 
H_max = 20  
n = 5  
radius = np.linspace(1, R_max, 100)  
height = np.linspace(1, H_max, 100)  

R, H = np.meshgrid(radius, height)

L = christmas_light_length(R, H, n)

# Cylindrical Christmas tree
def cyl_christmas_light_length(r, h, n):
    L_cyl = np.sqrt(h**2 + 4 * np.pi**2 * r**2 * n**2)
    return L_cyl

L_cyl = cyl_christmas_light_length(R, H, n)

L_diff = L_cyl - L

fig = plt.figure(figsize=(24, 8))

ax1 = fig.add_subplot(131, projection='3d')
surf1 = ax1.plot_surface(R, H, L, cmap='seismic', edgecolor='none')
ax1.set_xlabel('Radius (r)')
ax1.set_ylabel('Height (h)')
ax1.set_zlabel('Christmas light length $(L_{cone})$')
ax1.set_title(f'Conical Tree\nLength for {n} Loops')
fig.colorbar(surf1, ax=ax1)

ax2 = fig.add_subplot(132, projection='3d')
surf2 = ax2.plot_surface(R, H, L_cyl, cmap='seismic', edgecolor='none')
ax2.set_xlabel('Radius (r)')
ax2.set_ylabel('Height (h)')
ax2.set_zlabel('Christmas light length ($L_{cyl}$)')
ax2.set_title(f'Cylindrical Tree\nLength for {n} Loops')
fig.colorbar(surf2, ax=ax2)

ax3 = fig.add_subplot(133, projection='3d')
surf3 = ax3.plot_surface(R, H, L_diff, cmap='seismic', edgecolor='none')
ax3.set_xlabel('Radius (r)')
ax3.set_ylabel('Height (h)')
ax3.set_zlabel('Difference $(L_{cyl}-L_{cone})$')
ax3.set_title('Difference in Length of Lights')
fig.colorbar(surf3, ax=ax3)

plt.tight_layout()
#plt.savefig('christmaslight.png', dpi=300)
plt.show()
plt.close()