import numpy as np
from scipy.interpolate import interp1d

def calculate_reynolds(Q, D):
    """Calculate Reynolds number"""
    A = np.pi * (D/2)**2
    v = Q / A
    return (v * D) / (1.005e-6)

def friction_factor(Re):
    """Calculate friction factor for smooth pipe using Blasius equation"""
    if Re < 2300:
        return 64/Re
    else:
        return 0.316 * Re**(-0.25)  # Blasius equation for turbulent flow

def head_loss(Q):
    """Calculate total head loss for given flow rate"""
    # Convert Q from l/s to m³/s
    Q = Q / 1000
    
    # Pipe dimensions
    D_large = 0.028  # m
    D_small = 0.0136  # m
    
    # Lengths
    L_main = 7.5  # m
    L_large_fig2 = (0.56 + 0.24 + 0.09 + 0.095 + 0.8)  # m
    L_small = (0.4 + 0.4)  # m
    
    # Calculate velocities
    A_large = np.pi * (D_large/2)**2
    A_small = np.pi * (D_small/2)**2
    v_large = Q / A_large
    v_small = Q / A_small
    
    # Reynolds numbers
    Re_large = calculate_reynolds(Q, D_large)
    Re_small = calculate_reynolds(Q, D_small)
    
    # Friction factors
    f_large = friction_factor(Re_large)
    f_small = friction_factor(Re_small)
    
    # Major losses
    h_major_main = f_large * (L_main + L_large_fig2) * v_large**2 / (2 * 9.81 * D_large)
    h_major_small = f_small * L_small * v_small**2 / (2 * 9.81 * D_small)
    
    # Minor losses
    # Sum of K values from Fig 2.1 = 2
    # Additional K values from Fig 2.2: bends and area changes
    K_bends = 4 * 0.5  # Assuming K ≈ 0.5 for 90° bends with r/D = 1
    K_contractions = 0.5  # Approximate value for sudden contraction
    K_expansions = 1.0  # Approximate value for sudden expansion
    
    h_minor_large = (2 + K_bends) * v_large**2 / (2 * 9.81)
    h_minor_small = (K_contractions + K_expansions) * v_small**2 / (2 * 9.81)
    
    return h_major_main + h_major_small + h_minor_large + h_minor_small

# Pump curve data (Q in l/s, H in m)
Q_pump = np.array([0, 1, 2, 3, 4])
H_pump = np.array([6.8, 6.2, 5.4, 4.2, 2.8])

# Create interpolation function for pump curve
pump_curve = interp1d(Q_pump, H_pump, kind='linear')

# Find intersection point
Q_range = np.linspace(0, 4, 1000)
H_system = [head_loss(q) for q in Q_range]
H_pump_range = pump_curve(Q_range)

# Find where difference is minimal
intersection_idx = np.argmin(np.abs(H_system - H_pump_range))
Q_operating = Q_range[intersection_idx]
H_operating = H_system[intersection_idx]

print(f"Operating point:")
print(f"Q = {Q_operating:.2f} l/s")
print(f"H = {H_operating:.2f} m")