app/services/cfd_simulation.py

"""
CFD Simulation Service

Provides fluid flow analysis using the fluids library for validated calculations.
OpenFOAM is too large (~2GB) for Hugging Face Spaces, so we use analytical methods.

External dependencies:
- fluids: Industry-standard fluid dynamics library
- numpy: Numerical operations
"""
import math
import uuid
from typing import Dict, Any, Optional, List
from dataclasses import dataclass

import numpy as np
import fluids
from fluids import Reynolds, friction_factor, K_from_f, dP_from_K


@dataclass
class CFDInput:
    """Input parameters for CFD simulation."""
    pipe_diameter_mm: float
    pipe_length_mm: float
    flow_rate_lpm: float
    oil_viscosity_pa_s: float = 0.046  # ISO VG46 at 40°C
    oil_density_kg_m3: float = 870
    inlet_temperature_c: float = 40.0
    mesh_resolution: str = "medium"  # Kept for API compatibility


@dataclass
class CFDResult:
    """Results from CFD simulation."""
    success: bool
    message: str
    job_id: str
    # Flow results
    flow_regime: str
    reynolds_number: float
    average_velocity_ms: float
    max_velocity_ms: float
    pressure_drop_mpa: float
    # Comparison with analytical
    analytical_pressure_drop_mpa: float
    error_percent: float
    # Mesh info (zero for analytical)
    num_cells: int
    # Visualization data
    velocity_field: Optional[Dict] = None
    pressure_field: Optional[Dict] = None


def calculate_hagen_poiseuille(
    pipe_diameter_mm: float,
    pipe_length_mm: float,
    flow_rate_lpm: float,
    oil_viscosity_pa_s: float = 0.046,
    oil_density_kg_m3: float = 870
) -> Dict[str, float]:
    """
    Pipe flow analysis using the fluids library.

    Uses fluids.Reynolds and fluids.friction_factor for validated calculations.
    Handles laminar, transition, and turbulent flow regimes.

    Args:
        pipe_diameter_mm: Pipe inner diameter
        pipe_length_mm: Pipe length
        flow_rate_lpm: Flow rate in L/min
        oil_viscosity_pa_s: Dynamic viscosity
        oil_density_kg_m3: Oil density

    Returns:
        Dictionary with flow solution
    """
    # Convert units
    d = pipe_diameter_mm / 1000  # m
    L = pipe_length_mm / 1000  # m
    Q = flow_rate_lpm / 60000  # m³/s

    # Cross-sectional area
    A = math.pi * (d/2)**2

    # Average velocity
    v_avg = Q / A if A > 0 else 0

    # Use fluids library for Reynolds number calculation
    # fluids.Reynolds(V, D, rho, mu) - returns Reynolds number
    if d > 0 and oil_viscosity_pa_s > 0:
        Re = Reynolds(V=v_avg, D=d, rho=oil_density_kg_m3, mu=oil_viscosity_pa_s)
    else:
        Re = 0

    # Determine flow regime
    if Re < 2300:
        flow_regime = "laminar"
        v_max = 2 * v_avg  # Parabolic profile for laminar
    elif Re < 4000:
        flow_regime = "transition"
        v_max = 1.5 * v_avg  # Approximate
    else:
        flow_regime = "turbulent"
        v_max = 1.22 * v_avg  # Power-law profile

    # Use fluids library for friction factor
    # friction_factor(Re, eD) - eD is relative roughness (0 for smooth pipe)
    if Re > 0:
        f = friction_factor(Re=Re, eD=0)  # Smooth pipe
    else:
        f = 0

    # Pressure drop using Darcy-Weisbach equation
    # ΔP = f × (L/D) × (ρV²/2)
    if d > 0 and L > 0:
        delta_P = f * (L / d) * (0.5 * oil_density_kg_m3 * v_avg**2)
    else:
        delta_P = 0

    # Wall shear stress
    tau_wall = delta_P * d / (4 * L) if L > 0 else 0

    # Entry length (distance for flow to become fully developed)
    if flow_regime == "laminar":
        entry_length = 0.06 * Re * d
    else:
        entry_length = 4.4 * d * Re**(1/6) if Re > 0 else 0

    return {
        "flow_regime": flow_regime,
        "reynolds_number": round(Re, 0),
        "average_velocity_ms": round(v_avg, 4),
        "max_velocity_ms": round(v_max, 4),
        "pressure_drop_pa": round(delta_P, 2),
        "pressure_drop_mpa": round(delta_P / 1e6, 6),
        "friction_factor": round(f, 6),
        "wall_shear_stress_pa": round(tau_wall, 2),
        "entry_length_mm": round(entry_length * 1000, 1),
    }


def _generate_velocity_profile(
    pipe_diameter_mm: float,
    pipe_length_mm: float,
    flow_rate_lpm: float,
    flow_regime: str
) -> Dict[str, Any]:
    """
    Generate velocity profile data for visualization.
    """
    import numpy as np

    d = pipe_diameter_mm / 1000  # m
    R = d / 2
    Q = flow_rate_lpm / 60000  # m³/s
    A = math.pi * R**2
    v_avg = Q / A if A > 0 else 0

    # Generate radial positions and velocities
    n_radial = 20
    n_axial = 10

    r_vals = np.linspace(0, R, n_radial)
    z_vals = np.linspace(0, pipe_length_mm, n_axial)

    positions = []
    velocities = []
    magnitudes = []

    for z in z_vals:
        for r in r_vals:
            for theta in [0, math.pi/2, math.pi, 3*math.pi/2]:
                x = r * math.cos(theta) * 1000  # Convert back to mm
                y = r * math.sin(theta) * 1000

                positions.extend([x, y, z])

                # Velocity profile (parabolic for laminar)
                if flow_regime == "laminar":
                    v = 2 * v_avg * (1 - (r/R)**2) if R > 0 else 0
                else:
                    # Power-law profile for turbulent (1/7 power law)
                    v = v_avg * 1.22 * (1 - r/R)**(1/7) if R > 0 and r < R else 0

                # Velocity is axial (z-direction)
                velocities.extend([0, 0, v])
                magnitudes.append(v)

    return {
        "positions": positions,
        "vectors": velocities,
        "magnitudes": magnitudes,
    }


def _generate_pressure_field(
    pipe_diameter_mm: float,
    pipe_length_mm: float,
    pressure_drop_pa: float
) -> Dict[str, Any]:
    """
    Generate pressure field data for visualization.
    Pressure varies linearly along the pipe.
    """
    import numpy as np

    n_axial = 20

    z_vals = np.linspace(0, pipe_length_mm, n_axial)

    positions = []
    values = []

    inlet_pressure = pressure_drop_pa  # Gauge pressure

    for z in z_vals:
        # Linear pressure drop
        p = inlet_pressure * (1 - z / pipe_length_mm)

        # Add points at center
        positions.extend([0, 0, z])
        values.append(p)

    return {
        "positions": positions,
        "values": values,
        "range": {
            "min": 0,
            "max": inlet_pressure,
        }
    }


def run_cfd_simulation(params: CFDInput) -> CFDResult:
    """
    Run CFD analysis using analytical Hagen-Poiseuille solution.

    Args:
        params: CFD input parameters

    Returns:
        CFDResult with flow analysis results
    """
    job_id = str(uuid.uuid4())[:8]

    # Calculate analytical solution
    analytical = calculate_hagen_poiseuille(
        params.pipe_diameter_mm,
        params.pipe_length_mm,
        params.flow_rate_lpm,
        params.oil_viscosity_pa_s,
        params.oil_density_kg_m3
    )

    # Generate visualization data
    velocity_field = _generate_velocity_profile(
        params.pipe_diameter_mm,
        params.pipe_length_mm,
        params.flow_rate_lpm,
        analytical["flow_regime"]
    )

    pressure_field = _generate_pressure_field(
        params.pipe_diameter_mm,
        params.pipe_length_mm,
        analytical["pressure_drop_pa"]
    )

    return CFDResult(
        success=True,
        message=f"Flow analysis completed using analytical {analytical['flow_regime']} solution",
        job_id=job_id,
        flow_regime=analytical["flow_regime"],
        reynolds_number=analytical["reynolds_number"],
        average_velocity_ms=analytical["average_velocity_ms"],
        max_velocity_ms=analytical["max_velocity_ms"],
        pressure_drop_mpa=analytical["pressure_drop_mpa"],
        analytical_pressure_drop_mpa=analytical["pressure_drop_mpa"],
        error_percent=0,  # No error since using analytical directly
        num_cells=0,  # No mesh for analytical
        velocity_field=velocity_field,
        pressure_field=pressure_field,
    )


def cfd_result_to_dict(result: CFDResult) -> Dict[str, Any]:
    """Convert CFDResult to dictionary for JSON response."""
    return {
        "success": result.success,
        "message": result.message,
        "job_id": result.job_id,
        "flow_analysis": {
            "flow_regime": result.flow_regime,
            "reynolds_number": int(result.reynolds_number),
            "average_velocity_ms": round(result.average_velocity_ms, 4),
            "max_velocity_ms": round(result.max_velocity_ms, 4),
            "pressure_drop_mpa": round(result.pressure_drop_mpa, 6),
        },
        "analytical_comparison": {
            "analytical_pressure_drop_mpa": round(result.analytical_pressure_drop_mpa, 6),
            "cfd_vs_analytical_error_percent": round(result.error_percent, 2),
        },
        "mesh_info": {
            "num_cells": result.num_cells,
        },
        "visualization": {
            "velocity_field": result.velocity_field,
            "pressure_field": result.pressure_field,
        } if result.velocity_field else None,
    }


def calculate_hydraulic_loss_coefficient(
    fitting_type: str,
    pipe_diameter_mm: float,
    flow_rate_lpm: float,
    oil_density_kg_m3: float = 870
) -> Dict[str, float]:
    """
    Calculate pressure loss for common hydraulic fittings.

    Uses fluids.fittings library where possible for validated K values,
    with fallback to standard handbook values.

    Args:
        fitting_type: Type of fitting ("elbow_90", "tee", "valve", etc.)
        pipe_diameter_mm: Pipe diameter
        flow_rate_lpm: Flow rate
        oil_density_kg_m3: Oil density

    Returns:
        Dictionary with loss coefficient and pressure drop
    """
    d = pipe_diameter_mm / 1000  # m
    Q = flow_rate_lpm / 60000  # m³/s
    A = math.pi * (d/2)**2
    v = Q / A if A > 0 else 0

    # Try fluids library fittings first
    K = None
    source = "fluids"

    try:
        if fitting_type == "elbow_90":
            K = fluids.fittings.bend_rounded(Di=d, angle=90, fd=0.02)
        elif fitting_type == "elbow_90_long":
            K = fluids.fittings.bend_rounded(Di=d, angle=90, fd=0.02, bend_diameters=1.5)
        elif fitting_type == "elbow_45":
            K = fluids.fittings.bend_rounded(Di=d, angle=45, fd=0.02)
        elif fitting_type == "elbow_180":
            K = fluids.fittings.bend_rounded(Di=d, angle=180, fd=0.02)
        elif fitting_type == "sudden_expansion":
            K = fluids.fittings.contraction_sharp(Di1=d*0.8, Di2=d, fd=0.02)
        elif fitting_type == "sudden_contraction":
            K = fluids.fittings.contraction_sharp(Di1=d, Di2=d*0.8, fd=0.02)
        elif fitting_type == "pipe_entrance_sharp":
            K = fluids.fittings.entrance_sharp()
        elif fitting_type == "pipe_exit":
            K = fluids.fittings.exit_normal()
    except Exception:
        K = None

    # Fallback to handbook values if fluids library doesn't have the fitting
    if K is None:
        source = "handbook"
        K_VALUES = {
            "elbow_90": 0.9,
            "elbow_90_long": 0.6,
            "elbow_45": 0.4,
            "elbow_180": 1.5,
            "tee_through": 0.3,
            "tee_branch": 1.0,
            "gate_valve_open": 0.2,
            "gate_valve_half": 5.6,
            "ball_valve_open": 0.1,
            "ball_valve_half": 5.5,
            "check_valve": 2.5,
            "filter": 3.0,
            "strainer": 2.0,
            "sudden_expansion": 1.0,
            "sudden_contraction": 0.5,
            "gradual_expansion": 0.3,
            "gradual_contraction": 0.2,
            "pipe_entrance_sharp": 0.5,
            "pipe_entrance_rounded": 0.04,
            "pipe_exit": 1.0,
        }
        K = K_VALUES.get(fitting_type, 1.0)

    # Use fluids library for pressure drop calculation
    # dP_from_K(K, rho, V) - returns pressure drop in Pa
    if v > 0:
        delta_P = dP_from_K(K=K, rho=oil_density_kg_m3, V=v)
    else:
        delta_P = 0

    return {
        "fitting_type": fitting_type,
        "loss_coefficient_K": round(K, 4),
        "velocity_ms": round(v, 4),
        "pressure_loss_pa": round(delta_P, 2),
        "pressure_loss_mpa": round(delta_P / 1e6, 6),
        "source": source,
    }


def calculate_system_pressure_drop(
    pipes: List[Dict],
    fittings: List[Dict],
    flow_rate_lpm: float,
    oil_viscosity_pa_s: float = 0.046,
    oil_density_kg_m3: float = 870
) -> Dict[str, Any]:
    """
    Calculate total pressure drop for a hydraulic system.

    Args:
        pipes: List of pipe segments [{diameter_mm, length_mm}, ...]
        fittings: List of fittings [{type, diameter_mm}, ...]
        flow_rate_lpm: System flow rate
        oil_viscosity_pa_s: Oil viscosity
        oil_density_kg_m3: Oil density

    Returns:
        Dictionary with total pressure drop and breakdown
    """
    total_pipe_loss = 0
    total_fitting_loss = 0
    pipe_details = []
    fitting_details = []

    # Calculate pipe losses
    for pipe in pipes:
        result = calculate_hagen_poiseuille(
            pipe["diameter_mm"],
            pipe["length_mm"],
            flow_rate_lpm,
            oil_viscosity_pa_s,
            oil_density_kg_m3
        )
        total_pipe_loss += result["pressure_drop_pa"]
        pipe_details.append({
            "diameter_mm": pipe["diameter_mm"],
            "length_mm": pipe["length_mm"],
            "pressure_drop_mpa": result["pressure_drop_mpa"],
            "flow_regime": result["flow_regime"],
        })

    # Calculate fitting losses
    for fitting in fittings:
        result = calculate_hydraulic_loss_coefficient(
            fitting["type"],
            fitting["diameter_mm"],
            flow_rate_lpm,
            oil_density_kg_m3
        )
        total_fitting_loss += result["pressure_loss_pa"]
        fitting_details.append({
            "type": fitting["type"],
            "diameter_mm": fitting["diameter_mm"],
            "pressure_drop_mpa": result["pressure_loss_mpa"],
            "K_factor": result["loss_coefficient_K"],
        })

    total_loss = total_pipe_loss + total_fitting_loss

    return {
        "total_pressure_drop_mpa": round(total_loss / 1e6, 4),
        "pipe_losses_mpa": round(total_pipe_loss / 1e6, 4),
        "fitting_losses_mpa": round(total_fitting_loss / 1e6, 4),
        "pipe_loss_fraction": round(total_pipe_loss / total_loss, 2) if total_loss > 0 else 0,
        "fitting_loss_fraction": round(total_fitting_loss / total_loss, 2) if total_loss > 0 else 0,
        "pipes": pipe_details,
        "fittings": fitting_details,
    }