# src/friction.py
"""
Darcy-Weisbach friction factor calculation.

Supports:
  - Laminar:   f = 64 / Re                           (Re < 2300)
  - Turbulent: Colebrook-White implicit equation      (Re > 4000)
  - Transition: linear blend between laminar & turbulent (2300 <= Re <= 4000)
"""

import numpy as np


def _colebrook_white(Re, eps_D, n_iter=10):
    """
    Solve the Colebrook-White equation for Darcy friction factor f:
        1/sqrt(f) = -2 log10( eps_D/3.7 + 2.51/(Re*sqrt(f)) )

    Uses fixed-point iteration seeded with the Swamee-Jain approximation.
    """
    # Swamee-Jain initial guess (explicit approximation)
    A = eps_D / 3.7
    B = 2.51 / Re
    f = 0.25 / (np.log10(A + B / np.sqrt(0.02))) ** 2

    for _ in range(n_iter):
        f = 0.25 / (np.log10(A + B / np.sqrt(f))) ** 2

    return f


def darcy_friction_factor(Re, eps_D):
    """
    Compute Darcy-Weisbach friction factor for a given Reynolds number
    and relative roughness eps/D.

    Parameters
    ----------
    Re : float or ndarray
        Reynolds number (ρ|u|D/μ). Values <= 0 return 0 (no flow).
    eps_D : float
        Relative roughness ε/D (dimensionless).

    Returns
    -------
    f : same shape as Re
        Darcy friction factor.
    """
    Re = np.asarray(Re, dtype=float)
    scalar = Re.ndim == 0
    Re = np.atleast_1d(Re)

    f = np.zeros_like(Re)

    lam = Re < 2300
    turb = Re > 4000
    trans = ~lam & ~turb  # 2300 <= Re <= 4000

    # Laminar: f = 64/Re (avoid division by zero for Re~0)
    Re_lam = np.where(Re > 1e-12, Re, 1e-12)
    f[lam] = 64.0 / Re_lam[lam]

    # Turbulent: Colebrook-White
    if np.any(turb):
        f[turb] = _colebrook_white(Re[turb], eps_D)

    # Transition: linear blend
    if np.any(trans):
        f_lam = 64.0 / Re_lam[trans]
        f_turb = _colebrook_white(Re[trans], eps_D)
        alpha = (Re[trans] - 2300.0) / 1700.0  # 0 at Re=2300, 1 at Re=4000
        f[trans] = (1.0 - alpha) * f_lam + alpha * f_turb

    return float(f[0]) if scalar else f