import numpy as np
import matplotlib.pyplot as plt

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 = 10 
H_max = 20  
n = 1.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

# Create a single figure with 3 subplots (1 row, 3 columns)
fig, axes = plt.subplots(1, 3, figsize=(24, 8))

# First subplot: Conical Tree
ax1 = axes[0]
c1 = ax1.pcolormesh(R, H, L, cmap='seismic', shading='auto')
ax1.set_xlabel('Radius (r)')
ax1.set_ylabel('Height (h)')
ax1.set_title(f'Conical Tree Light\nLength for {n} Loops')
fig.colorbar(c1, ax=ax1)

ax2 = axes[1]
c2 = ax2.pcolormesh(R, H, L_cyl, cmap='seismic', shading='auto')
ax2.set_xlabel('Radius (r)')
ax2.set_ylabel('Height (h)')
ax2.set_title(f'Cylindrical Tree Light\nLength for {n} Loops')
fig.colorbar(c2, ax=ax2)

ax3 = axes[2]
c3 = ax3.pcolormesh(R, H, L_diff, cmap='seismic', shading='auto')
ax3.set_xlabel('Radius (r)')
ax3.set_ylabel('Height (h)')
ax3.set_title('Difference in Length of Lights')
fig.colorbar(c3, ax=ax3)

plt.tight_layout()

plt.show()
plt.close()