"""GPU RANS backend constants."""

from __future__ import annotations

FULL_GPU_RANS_GUARD = "scripts/verify_gpu_rans_solver.sh"
GPU_PRIMITIVE_PROOF = "scripts/verify_gpu_algorithm.sh"
GPU_PRIMITIVE_NAME = "cell_flux_imbalance"

PRIMARY_FIELDS = ("U", "p", "phi")
TURBULENCE_FIELDS = ("nut", "k", "omega")
REQUIRED_FIELDS = (*PRIMARY_FIELDS, *TURBULENCE_FIELDS)

GPU_BACKEND_PROVIDER = "quadrants_cuda_rans_solver"
GPU_RUNTIME_UNAVAILABLE = "gpu_runtime_unavailable"
GPU_KERNELS_MISSING = "missing_gpu_solver_kernels"
GPU_STAGE_UNSUPPORTED = "unsupported_gpu_solver_stage"
GPU_NUMERICAL_MISMATCH = "gpu_numerical_mismatch"
GPU_STAGE_CAPABILITY_PREFIX = "solver_stage_contract:"
GPU_INPUT_SCHEMA_VERSION = 1
DEFAULT_LAMINAR_NU = 1.5e-5
DEFAULT_MOMENTUM_RELAXATION_ALPHA = 0.9
DEFAULT_MOMENTUM_PBICGSTAB_ITERATIONS = 50
DEFAULT_MOMENTUM_PBICGSTAB_RESIDUAL_TOLERANCE_SQUARED = 1.0e-16
DEFAULT_PRESSURE_CG_ITERATIONS = 300
DEFAULT_SIMPLE_CONSISTENT_RATU_FACTOR = 10.0
DEFAULT_OMEGA_WALL_BETA1 = 0.075

STAGE_OBSERVABILITY_GROUPS = (
    {
        "name": "momentum_assembly",
        "split_stages": ("assemble_momentum_terms", "assemble_UEqn"),
        "run_one_stages": ("assemble_UEqn",),
        "split_outputs": {
            "assemble_momentum_terms": ("terms",),
            "assemble_UEqn": ("UEqn",),
        },
    },
    {
        "name": "pressure_assembly",
        "split_stages": ("compute_pressure_inputs", "assemble_pEqn"),
        "run_one_stages": ("compute_pressure_inputs", "assemble_pEqn"),
        "split_outputs": {
            "compute_pressure_inputs": ("HbyA", "phiHbyA", "rAU"),
            "assemble_pEqn": ("pEqn",),
        },
    },
    {
        "name": "linear_solve_results",
        "split_stages": ("solve_UEqn", "solve_pEqn"),
        "run_one_stages": ("solve_UEqn", "solve_pEqn"),
        "split_outputs": {
            "solve_UEqn": ("performance", "field_after"),
            "solve_pEqn": ("performance", "p", "phi"),
        },
    },
    {
        "name": "final_correction",
        "split_stages": ("correct_velocity_pressure_flux",),
        "run_one_stages": ("update_phi_from_pEqn_flux", "correct_velocity_pressure_flux"),
        "split_outputs": {
            "correct_velocity_pressure_flux": ("U", "p", "phi"),
        },
    },
    {
        "name": "turbulence_updates",
        "split_stages": ("momentum_transport_predict", "momentum_transport_correct"),
        "run_one_stages": ("momentum_transport_predict", "momentum_transport_correct"),
        "split_outputs": {
            "momentum_transport_predict": ("case_path", "solver_name"),
            "momentum_transport_correct": ("U", "p", "phi", "nut", "k", "omega"),
        },
    },
)

GPU_STAGE_KERNEL_ENTRYPOINTS = {
    "momentum_assembly": ["gpu_rans_momentum_assembly", "gpu_rans_momentum_diffusion_coefficients", "gpu_rans_momentum_wall_diffusion_coefficients", "gpu_rans_momentum_convection_coefficients", "gpu_rans_momentum_equation_relaxation"],
    "pressure_assembly": ["gpu_rans_pressure_inputs", "gpu_rans_consistent_rAtU", "gpu_rans_momentum_hbyA_source", "gpu_rans_momentum_hbyA_face_accumulate", "gpu_rans_momentum_hbyA_finish", "gpu_rans_momentum_hbyA_fixed_value_boundary", "gpu_rans_surface_flux_from_cells", "gpu_rans_pressure_assembly", "gpu_rans_pressure_laplacian_coefficients", "gpu_rans_pressure_mixed_boundary_laplacian", "gpu_rans_pressure_source_from_flux", "gpu_rans_pressure_source_from_boundary_flux"],
    "linear_solve_results": ["gpu_ldu_matvec_vector_asymmetric_diag", "gpu_ldu_matvec_vector_asymmetric_face_accumulate", "gpu_bicgstab_initialize_vector", "gpu_bicgstab_dot_vector", "gpu_bicgstab_update_direction_vector", "gpu_bicgstab_precondition_vector", "gpu_bicgstab_update_intermediate_vector_preconditioned", "gpu_bicgstab_update_solution_residual_vector_preconditioned", "gpu_vector_residual_squared", "gpu_ldu_matvec_scalar_symmetric_diag", "gpu_ldu_matvec_scalar_symmetric_face_accumulate", "gpu_pcg_initialize_scalar", "gpu_cg_dot_scalar", "gpu_pcg_update_solution_residual_scalar", "gpu_pcg_update_direction_scalar"],
    "final_correction": ["gpu_rans_pressure_flux_correction", "gpu_rans_final_correction", "gpu_rans_pressure_velocity_correction"],
    "turbulence_updates": ["gpu_rans_turbulence_update", "gpu_rans_omega_wall_update"],
}