import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import fsolve

# Colebrook-White equation to solve for Darcy friction factor (f)
def colebrook_white_equation(f, Re, epsilon_D):
    left_side = -2 * np.log10((epsilon_D / 3.7) + (2.51 / (Re * np.sqrt(f))))
    right_side = 1 / np.sqrt(f)
    return left_side - right_side

# Function to calculate Darcy friction factor (f) using Colebrook-White equation
def colebrook_white_friction_factor(Re, epsilon_D):
    f_guess = 0.02  # Initial guess for f
    f_solution = fsolve(colebrook_white_equation, f_guess, args=(Re, epsilon_D))
    return f_solution[0]

# Reynolds numbers range
Re_values = np.logspace(2, 8, 100)

# Range of roughness ratios (epsilon/D)
epsilon_D_range = np.linspace(0.00005, 0.05, 10)

# Plotting the Moody diagram using Colebrook-White equation for each roughness ratio
plt.figure(figsize=(8, 6))

for epsilon_D in epsilon_D_range:
    # Calculate corresponding Darcy friction factors using Colebrook-White equation
    f_values = [colebrook_white_friction_factor(Re, epsilon_D) for Re in Re_values]
   
    # Plot each curve
    plt.loglog(Re_values, f_values, label=f'Roughness Ratio: {epsilon_D:.5f}')

# Customize the plot
plt.title('Moody Diagram using Colebrook-White Equation (Range of Roughness Ratios)')
plt.xlabel('Reynolds Number')
plt.ylabel('Darcy Friction Factor (f)')
plt.grid(True, which='both', linestyle='--', linewidth=0.5)
plt.legend()
plt.show()