import jax
import jax.numpy as jnp

class Groups:
    """
    JAX-accelerated suite of dimensionless groups for scaling Physics AI models.
    Inputs may have leading batch dimensions; all operations broadcast and support
    both single-point and batched evaluation.
    """

    @staticmethod
    @jax.jit
    def re(u, L, rho, mu):
        """
        Reynolds Number (Re): Ratio of Inertial to Viscous forces.
        
        :param u: Flow velocity [m/s].
        :param L: Characteristic length scale (e.g., pipe diameter) [m].
        :param rho: Fluid density [kg/m^3].
        :param mu: Dynamic viscosity [Pa*s].
        :return: Dimensionless Reynolds number.
        
        SciML: Scales the diffusive term in Navier-Stokes loss functions.
        """
        return (rho * u * L) / mu

    @staticmethod
    @jax.jit
    def pr(mu, cp, k):
        """
        Prandtl Number (Pr): Ratio of Momentum to Thermal diffusivity.
        
        :param mu: Dynamic viscosity [Pa*s].
        :param cp: Specific heat capacity at constant pressure [J/kg*K].
        :param k: Thermal conductivity [W/m*K].
        :return: Dimensionless Prandtl number.
        
        SciML: Bridges the velocity and temperature fields in multi-physics PINNs.
        """
        return (cp * mu) / k

    @staticmethod
    @jax.jit
    def nu(h, L, k):
        """
        Nusselt Number (Nu): Ratio of Convective to Conductive heat transfer.
        
        :param h: Convective heat transfer coefficient [W/m^2*K].
        :param L: Characteristic length [m].
        :param k: Fluid thermal conductivity [W/m*K].
        :return: Dimensionless Nusselt number.
        
        SciML: Acts as a target metric for AI-driven thermal optimization.
        """
        return (h * L) / k

    @staticmethod
    @jax.jit
    def gr(beta, dT, L, nu):
        """
        Grashof Number (Gr): Ratio of Buoyancy to Viscous forces.
        
        :param beta: Volumetric thermal expansion coefficient [1/K].
        :param dT: Temperature difference (T_surface - T_ambient) [K].
        :param L: Characteristic length [m].
        :param nu: Kinematic viscosity [m^2/s].
        :return: Dimensionless Grashof number.
        
        SciML: Drives the buoyancy source term in natural convection models.

        Note: ``g`` is fixed at 9.81 m/s² (not a function argument). π-constant audit
        recipes in ``moju.monitor.pi_constant_recipes`` rescale only ``beta``, ``dT``,
        ``L``, and ``nu``.
        """
        g = 9.81
        return (g * beta * dT * L**3) / (nu**2)

    @staticmethod
    @jax.jit
    def ma(u, a):
        """
        Mach Number (Ma): Ratio of Flow velocity to speed of sound.
        
        :param u: Flow velocity [m/s].
        :param a: Local speed of sound [m/s].
        :return: Dimensionless Mach number.
        """
        return u / a

    @staticmethod
    @jax.jit
    def we(rho, u, L, sigma):
        """
        Weber Number (We): Ratio of Inertia to Surface Tension.
        
        :param rho: Fluid density [kg/m^3].
        :param u: Flow velocity [m/s].
        :param L: Characteristic length (e.g., droplet diameter) [m].
        :param sigma: Surface tension [N/m].
        :return: Dimensionless Weber number.
        """
        return (rho * u**2 * L) / sigma

    @staticmethod
    @jax.jit
    def pe(re, pr):
        """
        Peclet Number (Pe): Ratio of Advection to Diffusion.
        
        :param re: Reynolds number (dimensionless).
        :param pr: Prandtl number (dimensionless).
        :return: Dimensionless Peclet number.
        
        Use case: Heat and mass transfer in flowing fluids.
        SciML: Determines if the AI should prioritize the advective or diffusive term.
        """
        return re * pr

    @staticmethod
    @jax.jit
    def st(f, u, L):
        """
        Strouhal Number (St): Ratio of Oscillatory to Inertial mechanisms.
        
        :param f: Characteristic frequency (e.g., vortex shedding) [Hz].
        :param u: Flow velocity [m/s].
        :param L: Characteristic length [m].
        :return: Dimensionless Strouhal number.
        
        Use case: Vortex shedding and unsteady aerodynamics.
        SciML: Helps AI learn temporal frequencies in periodic flow signatures.
        """
        return (f * L) / u

    @staticmethod
    @jax.jit
    def bi(h, L, k_solid):
        """
        Biot Number (Bi): Internal vs External thermal resistance.
        
        :param h: Convective heat transfer coefficient [W/m^2*K].
        :param L: Characteristic length (e.g., half-thickness of solid) [m].
        :param k_solid: Thermal conductivity of the solid [W/m*K].
        :return: Dimensionless Biot number.
        
        Use case: Transient conduction in solids.
        SciML: Decides if a solid's internal temperature gradient can be ignored (lumped mass).
        """
        return (h * L) / k_solid

    @staticmethod
    @jax.jit
    def fo(alpha, t, L):
        """
        Fourier Number (Fo): Dimensionless time for heat conduction.
        
        :param alpha: Thermal diffusivity of the solid [m^2/s].
        :param t: Elapsed time [s].
        :param L: Characteristic length [m].
        :return: Dimensionless Fourier number.
        
        Use case: Measuring how far heat has penetrated a solid.
        SciML: Provides a normalized time-scale for unsteady thermal PINNs.
        """
        return (alpha * t) / L**2

    @staticmethod
    @jax.jit
    def sc(nu, D):
        """
        Schmidt Number (Sc): Momentum vs Mass diffusivity.
        
        :param nu: Kinematic viscosity [m^2/s].
        :param D: Mass diffusivity [m^2/s].
        :return: Dimensionless Schmidt number.
        
        Use case: Mass transfer and chemical vapor deposition.
        SciML: Controls the coupling between velocity and species concentration fields.
        """
        return nu / D

    @staticmethod
    @jax.jit
    def le(sc, pr):
        """
        Lewis Number (Le): Thermal vs Mass diffusivity.
        
        :param sc: Schmidt number (dimensionless).
        :param pr: Prandtl number (dimensionless).
        :return: Dimensionless Lewis number.
        
        Use case: Combined heat and mass transfer (e.g., combustion).
        SciML: Regulates the relative widths of thermal and concentration boundary layers.
        """
        return sc / pr

    @staticmethod
    @jax.jit
    def kn(lambda_mfp, L):
        """
        Knudsen Number (Kn): Mean free path vs Length scale.
        
        :param lambda_mfp: Mean free path of the gas [m].
        :param L: Characteristic length scale [m].
        :return: Dimensionless Knudsen number.
        
        Use case: Rarefied gas dynamics (microchannels, high-altitude flight).
        SciML: Indicates when Navier-Stokes breaks down (AI may need kinetic models).
        """
        return lambda_mfp / L

    @staticmethod
    @jax.jit
    def bo(rho, g, L, sigma):
        """
        Bond Number (Bo): Gravity vs Surface tension.
        
        :param rho: Fluid density [kg/m^3].
        :param g: Gravitational acceleration [m/s^2].
        :param L: Characteristic length [m].
        :param sigma: Surface tension [N/m].
        :return: Dimensionless Bond number.
        
        Use case: Shape of bubbles and capillary action.
        SciML: Critical for AI-driven microfluidics or ink-jet simulation.
        """
        return (rho * g * L**2) / sigma

    @staticmethod
    @jax.jit
    def ca(mu, u, sigma):
        """
        Capillary Number (Ca): Viscous forces vs Surface tension.
        
        :param mu: Dynamic viscosity [Pa*s].
        :param u: Flow velocity [m/s].
        :param sigma: Surface tension [N/m].
        :return: Dimensionless Capillary number.
        
        Use case: Porous media flow or thin-film coating.
        SciML: Scales the pressure-jump condition across interfaces in the AI loss.
        """
        return (mu * u) / sigma

    @staticmethod
    @jax.jit
    def eu(dp, rho, u):
        """
        Euler Number (Eu): Pressure vs Inertial forces.
        
        :param dp: Pressure difference [Pa].
        :param rho: Fluid density [kg/m^3].
        :param u: Characteristic flow velocity [m/s].
        :return: Dimensionless Euler number.
        
        Use case: Flow through valves, orifices, or pumps.
        SciML: Used to check the balance of the pressure-gradient term in momentum residuals.
        """
        return dp / (rho * u**2)

    @staticmethod
    @jax.jit
    def da(K, L):
        """
        Darcy Number (Da): Permeability vs Characteristic area.
        
        :param K: Permeability of the porous medium [m^2].
        :param L: Characteristic length scale [m].
        :return: Dimensionless Darcy number.
        
        Use case: Flow through sponges, soils, or filters.
        SciML: Scales the Darcy resistance term in porous-media-informed networks.
        """
        return K / L**2

    @staticmethod
    @jax.jit
    def ec(u, cp, dT):
        """
        Eckert Number (Ec): Kinetic energy vs Enthalpy difference.
        
        :param u: Flow velocity [m/s].
        :param cp: Specific heat capacity at constant pressure [J/kg*K].
        :param dT: Characteristic temperature difference [K].
        :return: Dimensionless Eckert number.
        
        Use case: High-speed flow where viscous dissipation causes heating.
        SciML: Controls the strength of the viscous dissipation term in the energy equation.
        """
        return u**2 / (cp * dT)

    @staticmethod
    @jax.jit
    def fo_mass(D, t, L):
        """
        Fourier Number for mass diffusion (Fo_mass): Dimensionless time D·t/L^2.

        :param D: Mass diffusivity [m^2/s].
        :param t: Elapsed time [s].
        :param L: Characteristic length [m].
        :return: Dimensionless Fourier number for mass.

        Use case: Fick diffusion; unsteady mass transfer.
        SciML: Normalized time-scale for species/concentration PINNs.
        """
        return (D * t) / L**2

    @staticmethod
    @jax.jit
    def wavenumber(k, L):
        """
        Dimensionless wavenumber (kL): k·L for Helmholtz/wave problems.

        :param k: Wavenumber [1/m].
        :param L: Characteristic length [m].
        :return: Dimensionless wavenumber.

        Use case: Helmholtz equation; frequency-domain wave scaling.
        SciML: Scales the k^2·phi term in Helmholtz residuals.
        """
        return k * L

    @staticmethod
    @jax.jit
    def pe_mass(re, sc):
        """
        Peclet Number for species (Pe_mass): Re·Sc = advection / mass diffusion.

        :param re: Reynolds number (dimensionless).
        :param sc: Schmidt number (dimensionless).
        :return: Dimensionless Peclet number for mass transfer.

        Use case: Species advection-diffusion; use with Laws.advection_diffusion for mass.
        SciML: Symmetry with Groups.pe(re, pr) for heat.
        """
        return re * sc

    @staticmethod
    @jax.jit
    def st_wave(omega, L, c):
        """
        Wave Strouhal (St_wave): omega·L/c for frequency-domain wave scaling.

        :param omega: Angular frequency [rad/s].
        :param L: Characteristic length [m].
        :param c: Wave speed [m/s].
        :return: Dimensionless wave Strouhal number.

        Use case: Wave equation; acoustics; resonance.
        SciML: Normalized frequency in wave-equation residuals.
        """
        return (omega * L) / c

    @staticmethod
    @jax.jit
    def poisson_rhs_pi(L, source, epsilon, phi_ref):
        """
        Dimensionless Poisson source strength for a reference-scaled potential.

        With spatial coordinate :math:`x^* = x/L`, a common nondimensional form is
        :math:`\\nabla^{*2}\\phi^* + \\Pi\\, s^* = 0` where :math:`\\Pi = L^2 s / (\\varepsilon \\phi_\\mathrm{ref})`
        when the dimensional equation is :math:`\\nabla^2\\phi + s/\\varepsilon = 0`
        and :math:`\\phi^* = \\phi/\\phi_\\mathrm{ref}`, :math:`s^*` is the scaled source.

        :param L: Characteristic length [m].
        :param source: Source term (same physical meaning as in the dimensional Poisson equation).
        :param epsilon: Permittivity or field constant (same as dimensional equation).
        :param phi_ref: Reference scale for the potential [same units as ``phi``].
        :return: Dimensionless group :math:`\\Pi` (broadcasts like inputs).
        """
        return (L**2) * source / (epsilon * phi_ref)

    @staticmethod
    @jax.jit
    def schrodinger_kinetic_length_squared(m, L, h_bar):
        """
        Kinetic index :math:`2 m L^2 / \\hbar^2` for the nondimensional steady Schrödinger equation.

        :func:`moju.piratio.Laws.schrodinger_steady` uses
        :math:`\\psi_\\mathrm{lap} = L^2\\nabla^2\\psi` and residual
        :math:`-\\psi_\\mathrm{lap} + K\\,(V-E)\\psi = 0` with :math:`K = 2mL^2/\\hbar^2`
        (equivalent to scaling the dimensional TISE by :math:`2mL^2/\\hbar^2`).

        :param m: Particle mass [kg].
        :param L: Characteristic length [m].
        :param h_bar: Reduced Planck constant [J*s].
        :return: Dimensionless scalar :math:`K` (broadcasts like inputs).
        """
        return 2.0 * m * (L**2) / (h_bar**2)
