/******************************************************************************
 * Copyright (c) 2024 Lawrence Livermore National Security, LLC and other
 * HYPRE Project Developers. See the top-level COPYRIGHT file for details.
 *
 * SPDX-License-Identifier: MIT
 ******************************************************************************/

#include "internal/scaling.h"
#include <math.h>
#include <mpi.h>
#include <string.h>
#include "_hypre_IJ_mv.h"
#include "_hypre_parcsr_mv.h"
#include "internal/error.h"
#include "internal/field.h"
#include "internal/gen_macros.h"
#include "internal/linsys.h"
#include "internal/utils.h"
#include "logging.h"

/*
 * gcovr: branches in #else arms for HYPRE < v3.0.0, in HYPRE_USING_GPU-only
 * sections, and on a few impractical error paths are excluded so branch
 * coverage reflects the active CPU / Hypre configuration (see markers below).
 */

#define Scaling_FIELDS(_prefix)                                       \
   ADD_FIELD_OFFSET_ENTRY(_prefix, enabled, hypredrv_FieldTypeIntSet) \
   ADD_FIELD_OFFSET_ENTRY(_prefix, type, hypredrv_FieldTypeIntSet)    \
   ADD_FIELD_OFFSET_ENTRY(_prefix, custom_values, hypredrv_FieldTypeDoubleArraySet)

/* Define num_fields macro */
#define Scaling_NUM_FIELDS \
   (sizeof(Scaling_field_offset_map) / sizeof(Scaling_field_offset_map[0]))

GENERATE_PREFIXED_COMPONENTS(Scaling) // LCOV_EXCL_LINE

/*-----------------------------------------------------------------------------
 * ScalingGetValidValues
 *-----------------------------------------------------------------------------*/

StrIntMapArray
hypredrv_ScalingGetValidValues(const char *key)
{
   /* GCOVR_EXCL_BR_START */
   if (!strcmp(key, "enabled"))
   {
      return STR_INT_MAP_ARRAY_CREATE_ON_OFF();
   }
   if (!strcmp(key, "type"))
   {
      static StrIntMap map[] = {
         {"rhs_l2", (int)SCALING_RHS_L2},
         {"dofmap_mag", (int)SCALING_DOFMAP_MAG},
         {"dofmap_custom", (int)SCALING_DOFMAP_CUSTOM},
         {"dofmap_row_custom", (int)SCALING_DOFMAP_ROW_CUSTOM},
         {"dofmap_col_custom", (int)SCALING_DOFMAP_COL_CUSTOM},
         {"dofmap_similarity_custom", (int)SCALING_DOFMAP_SIMILARITY_CUSTOM},
      };
      return STR_INT_MAP_ARRAY_CREATE(map);
   }
   return STR_INT_MAP_ARRAY_VOID();
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * ScalingSetDefaultArgs
 *-----------------------------------------------------------------------------*/

void
hypredrv_ScalingSetDefaultArgs(Scaling_args *args)
{
   args->enabled       = 0;
   args->type          = SCALING_RHS_L2;
   args->custom_values = NULL;
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingContextCreate
 *-----------------------------------------------------------------------------*/

void
hypredrv_ScalingContextCreate(MPI_Comm comm, Scaling_context **ctx_ptr)
{
   /* GCOVR_EXCL_BR_START */
   Scaling_context *ctx        = (Scaling_context *)malloc(sizeof(Scaling_context));
   ctx->enabled                = 0;
   ctx->type                   = SCALING_RHS_L2;
   ctx->is_applied             = 0;
   ctx->matrices_are_scaled    = 0;
   ctx->rhs_is_scaled          = 0;
   ctx->x_is_scaled            = 0;
   ctx->scalar_factor          = 1.0;
   ctx->scaling_vector         = NULL;
   ctx->inverse_scaling_vector = NULL;
   ctx->scaling_ijvec          = NULL;
   *ctx_ptr                    = ctx;
   HYPREDRV_LOG_COMMF(3, comm, NULL, 0, "scaling context created");
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * ScalingContextFreeVector
 *
 * Free the scaling vector held by the context, handling both ownership models:
 *   - dofmap_custom: scaling_vector is owned by scaling_ijvec (IJVector wrapper)
 *   - dofmap_mag:    scaling_vector is directly owned (no IJVector wrapper)
 *-----------------------------------------------------------------------------*/

static void
ScalingContextFreeVector(Scaling_context *ctx)
{
#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   /* GCOVR_EXCL_BR_START */
   if (ctx->inverse_scaling_vector)
   {
      HYPRE_SAFE_CALL(HYPRE_ParVectorDestroy(ctx->inverse_scaling_vector));
      ctx->inverse_scaling_vector = NULL;
   }
   if (ctx->scaling_ijvec)
   {
      HYPRE_SAFE_CALL(HYPRE_IJVectorDestroy(ctx->scaling_ijvec));
      ctx->scaling_ijvec  = NULL;
      ctx->scaling_vector = NULL;
   }
   else if (ctx->scaling_vector)
   {
      HYPRE_SAFE_CALL(HYPRE_ParVectorDestroy(ctx->scaling_vector));
      ctx->scaling_vector = NULL;
   }
   /* GCOVR_EXCL_BR_STOP */
#endif
}

static int
ScalingContextCacheInverse(Scaling_context *ctx)
{
#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   if (!ctx || !ctx->scaling_vector)
   {
      return 0;
   }
   if (ctx->inverse_scaling_vector)
   {
      HYPRE_ParVectorDestroy(ctx->inverse_scaling_vector);
      ctx->inverse_scaling_vector = NULL;
   }
   hypre_ParVector *inverse = NULL;
   hypre_ParVectorPointwiseInverse((hypre_ParVector *)ctx->scaling_vector, &inverse);
   ctx->inverse_scaling_vector = (HYPRE_ParVector)inverse;
   return inverse != NULL;
#else
   (void)ctx;
   return 0;
#endif
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingContextDestroy
 *-----------------------------------------------------------------------------*/

void
hypredrv_ScalingContextDestroy(MPI_Comm comm, Scaling_context **ctx_ptr)
{
   /* GCOVR_EXCL_BR_START */
   if (!ctx_ptr || !*ctx_ptr)
   {
      return;
   }

   Scaling_context *ctx = *ctx_ptr;

   ScalingContextFreeVector(ctx);

   free(ctx);
   *ctx_ptr = NULL;
   HYPREDRV_LOG_COMMF(3, comm, NULL, 0, "scaling context destroyed");
   /* GCOVR_EXCL_BR_STOP */
}

static MPI_Comm
ScalingCommFromMatrix(HYPRE_IJMatrix mat)
{
   /* GCOVR_EXCL_BR_START */
   if (!mat)
   {
      return MPI_COMM_NULL;
   }

   void *obj = NULL;
   HYPRE_IJMatrixGetObject(mat, &obj);
   if (!obj)
   {
      return MPI_COMM_NULL;
   }

   return hypre_ParCSRMatrixComm((hypre_ParCSRMatrix *)obj);
   /* GCOVR_EXCL_BR_STOP */
}

static MPI_Comm
ScalingCommFromVector(HYPRE_IJVector vec)
{
   /* GCOVR_EXCL_BR_START */
   if (!vec)
   {
      return MPI_COMM_NULL;
   }

   return hypre_IJVectorComm((hypre_IJVector *)vec);
   /* GCOVR_EXCL_BR_STOP */
}

static MPI_Comm
ScalingSystemCommResolve(HYPRE_IJMatrix mat_A, HYPRE_IJMatrix mat_M, HYPRE_IJVector vec_b,
                         HYPRE_IJVector vec_x)
{
   /* GCOVR_EXCL_BR_START */
   MPI_Comm comm = ScalingCommFromMatrix(mat_A);
   if (comm != MPI_COMM_NULL)
   {
      return comm;
   }

   comm = ScalingCommFromMatrix(mat_M);
   if (comm != MPI_COMM_NULL)
   {
      return comm;
   }

   comm = ScalingCommFromVector(vec_b);
   if (comm != MPI_COMM_NULL)
   {
      return comm;
   }

   return ScalingCommFromVector(vec_x);
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingCompute (rhs_l2 strategy)
 *-----------------------------------------------------------------------------*/

static void
ScalingComputeRHSL2(MPI_Comm comm, Scaling_context *ctx, HYPRE_IJVector vec_b)
{
   /* GCOVR_EXCL_BR_START */
   (void)comm;
   double b_norm = 0.0;

   hypredrv_LinearSystemComputeVectorNorm(vec_b, "L2", &b_norm);

   if (b_norm > 0.0)
   {
      ctx->scalar_factor = 1.0 / sqrt(b_norm);
   }
   else
   {
      ctx->scalar_factor = 1.0;
   }
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingCompute (dofmap_mag strategy)
 *-----------------------------------------------------------------------------*/

static void
ScalingComputeDofmapMag(MPI_Comm comm, Scaling_args *args, Scaling_context *ctx,
                        HYPRE_IJMatrix mat_A, IntArray *dofmap)
{
#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   /* GCOVR_EXCL_BR_START */
   (void)args;
   HYPRE_MemoryLocation memloc_tags = HYPRE_MEMORY_HOST;
#if defined(HYPRE_USING_GPU)
   HYPRE_MemoryLocation orig_mat_memloc = HYPRE_MEMORY_HOST;
#endif
   void              *obj_A  = NULL;
   HYPRE_ParCSRMatrix par_A  = NULL;
   HYPRE_BigInt       ilower = 0, iupper = 0;
   HYPRE_Int          num_local_rows = 0;
   HYPRE_Int         *tags           = NULL;
   HYPRE_Int          num_tags       = 0;
   HYPRE_Int          max_tag        = -1;
   int                myid           = 0;

   MPI_Comm_rank(comm, &myid);

   if (!dofmap || !dofmap->data)
   {
      hypredrv_ErrorCodeSet(ERROR_MISSING_DOFMAP);
      hypredrv_ErrorMsgAdd("dofmap scaling requires a dofmap to be set");
      return;
   }

   HYPRE_IJMatrixGetObject(mat_A, &obj_A);
   par_A = (HYPRE_ParCSRMatrix)obj_A;
#if defined(HYPRE_USING_GPU)
   /* GCOVR_EXCL_BR_START */
   orig_mat_memloc = hypre_ParCSRMatrixMemoryLocation((hypre_ParCSRMatrix *)par_A);
   /* GCOVR_EXCL_BR_STOP */
#endif

   /* Get local range from ParCSRMatrix directly instead of IJMatrix to avoid potential
    * issues */
   HYPRE_BigInt row_start = 0, row_end = 0, col_start = 0, col_end = 0;
   hypre_ParCSRMatrixGetLocalRange((hypre_ParCSRMatrix *)par_A, &row_start, &row_end,
                                   &col_start, &col_end);
   ilower         = row_start;
   iupper         = row_end;
   num_local_rows = (HYPRE_Int)(iupper - ilower + 1);

   if (num_local_rows < 0 || dofmap->size != (size_t)num_local_rows)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("dofmap size (%d) does not match local matrix rows (%d)",
                           dofmap->size, num_local_rows);
      return;
   }

   /* Allocate and fill tags array from dofmap */
   tags = (HYPRE_Int *)malloc((size_t)num_local_rows * sizeof(HYPRE_Int));
   for (HYPRE_Int i = 0; i < num_local_rows; i++)
   {
      tags[i] = (HYPRE_Int)dofmap->data[i];
      if (tags[i] > max_tag)
      {
         max_tag = tags[i];
      }
   }

   /* Compute global max tag */
   HYPRE_Int global_max_tag = 0;
   MPI_Allreduce(&max_tag, &global_max_tag, 1, MPI_INT, MPI_MAX, comm);
   num_tags = global_max_tag + 1;

   /* Destroy previous scaling vector if it exists (from previous system) */
   ScalingContextFreeVector(ctx);

   /* Work around hypre's device TaggedFnorm kernel, which can read tags out of bounds
    * for dofmap_mag on GPU builds. Compute the tagged scaling on host, then migrate the
    * resulting scaling vector back to the matrix memory location. */
#if defined(HYPRE_USING_GPU)
   /* GCOVR_EXCL_BR_START */
   if (hypre_GetExecPolicy1(orig_mat_memloc) == HYPRE_EXEC_DEVICE)
   {
      hypre_ParCSRMatrixMigrate((hypre_ParCSRMatrix *)par_A, HYPRE_MEMORY_HOST);
   }
   /* GCOVR_EXCL_BR_STOP */
#endif

   /* Compute scaling into a fresh ParVector */
   HYPRE_SAFE_CALL(HYPRE_ParCSRMatrixComputeScalingTagged(par_A, 1, memloc_tags, num_tags,
                                                          tags, &ctx->scaling_vector));

#if defined(HYPRE_USING_GPU)
   /* GCOVR_EXCL_BR_START */
   if (hypre_GetExecPolicy1(orig_mat_memloc) == HYPRE_EXEC_DEVICE)
   {
      hypre_ParVectorMigrate((hypre_ParVector *)ctx->scaling_vector, orig_mat_memloc);
      hypre_ParCSRMatrixMigrate((hypre_ParCSRMatrix *)par_A, orig_mat_memloc);
   }
   /* GCOVR_EXCL_BR_STOP */
#endif

   /* Free memory */
   free(tags);
   /* GCOVR_EXCL_BR_STOP */
#else
   /* GCOVR_EXCL_BR_START */
   (void)comm;
   (void)args;
   (void)ctx;
   (void)mat_A;
   (void)dofmap;
   /* Scaling disabled on older hypre versions */
   /* GCOVR_EXCL_BR_STOP */
#endif
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingCompute (dofmap_custom strategy)
 *-----------------------------------------------------------------------------*/

static void
ScalingComputeDofmapCustom(MPI_Comm comm, Scaling_args *args, Scaling_context *ctx,
                           HYPRE_IJMatrix mat_A, IntArray *dofmap)
{
#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   /* GCOVR_EXCL_BR_START */
   void              *obj_A  = NULL;
   HYPRE_ParCSRMatrix par_A  = NULL;
   HYPRE_BigInt       ilower = 0, iupper = 0;
   HYPRE_Int          num_local_rows = 0;
   int                myid           = 0;

   MPI_Comm_rank(comm, &myid);

   if (!dofmap || !dofmap->data)
   {
      hypredrv_ErrorCodeSet(ERROR_MISSING_DOFMAP);
      hypredrv_ErrorMsgAdd("custom dofmap scaling requires a dofmap to be set");
      return;
   }

   if (!args->custom_values || args->custom_values->size == 0)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("custom dofmap scaling requires custom_values to be set");
      return;
   }

   for (size_t i = 0; i < args->custom_values->size; i++)
   {
      if (args->custom_values->data[i] == 0.0)
      {
         hypredrv_ErrorCodeSet(ERROR_INVALID_VAL);
         hypredrv_ErrorMsgAdd(
            "custom dofmap scaling requires nonzero custom_values (entry %zu is zero)",
            i);
         return;
      }
   }

   HYPRE_IJMatrixGetObject(mat_A, &obj_A);
   par_A = (HYPRE_ParCSRMatrix)obj_A;

   /* Get local range from ParCSRMatrix directly */
   HYPRE_BigInt row_start = 0, row_end = 0, col_start = 0, col_end = 0;
   hypre_ParCSRMatrixGetLocalRange((hypre_ParCSRMatrix *)par_A, &row_start, &row_end,
                                   &col_start, &col_end);
   ilower         = row_start;
   iupper         = row_end;
   num_local_rows = (HYPRE_Int)(iupper - ilower + 1);

   if (num_local_rows < 0 || dofmap->size != (size_t)num_local_rows)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("dofmap size (%d) does not match local matrix rows (%d)",
                           dofmap->size, num_local_rows);
      return;
   }

   /* Find the maximum tag value in the dofmap */
   HYPRE_Int max_tag = -1;
   for (HYPRE_Int i = 0; i < num_local_rows; i++)
   {
      if (dofmap->data[i] > max_tag)
      {
         max_tag = dofmap->data[i];
      }
   }

   /* Compute global max tag */
   HYPRE_Int global_max_tag = 0;
   MPI_Allreduce(&max_tag, &global_max_tag, 1, MPI_INT, MPI_MAX, comm);
   HYPRE_Int num_tags = global_max_tag + 1;

   /* Verify that the number of custom values matches the number of unique tags */
   if ((size_t)num_tags != args->custom_values->size)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd(
         "dofmap_custom: number of custom_values (%zu) does not match number of "
         "unique dofmap tags (%d)",
         args->custom_values->size, num_tags);
      return;
   }

   /* Destroy previous scaling vector if it exists (from previous system) */
   ScalingContextFreeVector(ctx);

   /* Create IJVector wrapper for scaling vector */
   HYPRE_MemoryLocation memory_location =
      hypre_ParCSRMatrixMemoryLocation((hypre_ParCSRMatrix *)par_A);
   HYPRE_Complex *h_values =
      hypre_TAlloc(HYPRE_Complex, num_local_rows, HYPRE_MEMORY_HOST);
   const HYPRE_Complex *values = h_values;
#ifdef HYPRE_USING_GPU
   /* GCOVR_EXCL_BR_START */
   HYPRE_Complex *d_values = NULL;
   /* GCOVR_EXCL_BR_STOP */
#endif

   if (memory_location == HYPRE_MEMORY_UNDEFINED)
   {
      memory_location = HYPRE_MEMORY_HOST;
   }

   HYPRE_SAFE_CALL(HYPRE_IJVectorCreate(comm, ilower, iupper, &ctx->scaling_ijvec));
   HYPRE_SAFE_CALL(HYPRE_IJVectorSetObjectType(ctx->scaling_ijvec, HYPRE_PARCSR));
   HYPRE_SAFE_CALL(HYPRE_IJVectorInitialize_v2(ctx->scaling_ijvec, memory_location));

#ifdef HYPRE_USING_GPU
   if (memory_location == HYPRE_MEMORY_DEVICE)
   {
      values = d_values =
         hypre_TAlloc(HYPRE_Complex, num_local_rows, HYPRE_MEMORY_DEVICE);
   }
#endif

   /* Fill scaling vector: for each row i, use custom_values[dofmap[i]] */
   for (HYPRE_Int i = 0; i < num_local_rows; i++)
   {
      HYPRE_Int tag = dofmap->data[i];
      if (tag < 0 || tag >= (HYPRE_Int)args->custom_values->size)
      {
         hypre_TFree(h_values, HYPRE_MEMORY_HOST);
#ifdef HYPRE_USING_GPU
         /* GCOVR_EXCL_BR_START */
         if (d_values)
         {
            hypre_TFree(d_values, HYPRE_MEMORY_DEVICE);
         }
         /* GCOVR_EXCL_BR_STOP */
#endif
         HYPRE_SAFE_CALL(HYPRE_IJVectorDestroy(ctx->scaling_ijvec));
         ctx->scaling_ijvec = NULL;
         hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
         hypredrv_ErrorMsgAdd(
            "dofmap_custom: invalid tag %d at local row %d (expected 0-%zu)", tag, i,
            args->custom_values->size - 1);
         return;
      }
      h_values[i] = (HYPRE_Complex)args->custom_values->data[tag];
   }

#ifdef HYPRE_USING_GPU
   /* GCOVR_EXCL_BR_START */
   if (values != h_values)
   {
      hypre_TMemcpy(d_values, h_values, HYPRE_Complex, num_local_rows,
                    HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
   }
   /* GCOVR_EXCL_BR_STOP */
#endif

   HYPRE_SAFE_CALL(
      HYPRE_IJVectorSetValues(ctx->scaling_ijvec, num_local_rows, NULL, values));

   /* Assemble the vector */
   HYPRE_SAFE_CALL(HYPRE_IJVectorAssemble(ctx->scaling_ijvec));
   hypre_TFree(h_values, HYPRE_MEMORY_HOST);
#ifdef HYPRE_USING_GPU
   /* GCOVR_EXCL_BR_START */
   if (d_values)
   {
      hypre_TFree(d_values, HYPRE_MEMORY_DEVICE);
   }
   /* GCOVR_EXCL_BR_STOP */
#endif

   /* Cache the assembled ParVector owned by the IJVector */
   void *obj_scaling = NULL;
   HYPRE_IJVectorGetObject(ctx->scaling_ijvec, &obj_scaling);
   ctx->scaling_vector = (HYPRE_ParVector)obj_scaling;
   /* GCOVR_EXCL_BR_STOP */
#else
   /* GCOVR_EXCL_BR_START */
   (void)comm;
   (void)args;
   (void)ctx;
   (void)mat_A;
   (void)dofmap;
   /* Scaling disabled on older hypre versions */
   /* GCOVR_EXCL_BR_STOP */
#endif
}

/*-----------------------------------------------------------------------------
 * hypredrv_ScalingCompute
 *-----------------------------------------------------------------------------*/

void
hypredrv_ScalingCompute(MPI_Comm comm, Scaling_args *args, Scaling_context *ctx,
                        HYPRE_IJMatrix mat_A, HYPRE_IJVector vec_b, IntArray *dofmap)
{
   /* GCOVR_EXCL_BR_START */
   int log_rank = -1;
   if (hypredrv_LogEnabled(2))
   {
      log_rank = hypredrv_LogRankFromComm(comm);
   }
   if (!args || !ctx)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("hypredrv_ScalingCompute: args or ctx is NULL");
      HYPREDRV_LOGF(2, log_rank, NULL, 0,
                    "scaling compute failed: args or context is NULL");
      return;
   }

   if (!args->enabled)
   {
      ctx->enabled = 0;
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling compute skipped (disabled)");
      return;
   }

   ctx->enabled = 1;
   ctx->type    = args->type;
   HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling compute begin (type=%d)",
                 (int)args->type);

#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   switch (args->type)
   {
      case SCALING_RHS_L2:
         ScalingComputeRHSL2(comm, ctx, vec_b);
         break;

      case SCALING_DOFMAP_MAG:
         ScalingComputeDofmapMag(comm, args, ctx, mat_A, dofmap);
         break;

      case SCALING_DOFMAP_CUSTOM:
      case SCALING_DOFMAP_ROW_CUSTOM:
      case SCALING_DOFMAP_COL_CUSTOM:
      case SCALING_DOFMAP_SIMILARITY_CUSTOM:
         ScalingComputeDofmapCustom(comm, args, ctx, mat_A, dofmap);
         break;

      default:
         hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
         hypredrv_ErrorMsgAdd("hypredrv_ScalingCompute: unknown scaling type");
         HYPREDRV_LOGF(2, log_rank, NULL, 0,
                       "scaling compute failed: unknown scaling type=%d",
                       (int)args->type);
         break;
   }
   if (args->type != SCALING_RHS_L2)
   {
      int inverse_ok        = !hypredrv_ErrorCodeGet() && ScalingContextCacheInverse(ctx);
      int global_inverse_ok = 0;
      MPI_Allreduce(&inverse_ok, &global_inverse_ok, 1, MPI_INT, MPI_MIN, comm);
      if (!global_inverse_ok)
      {
         if (ctx->inverse_scaling_vector)
         {
            HYPRE_ParVectorDestroy(ctx->inverse_scaling_vector);
            ctx->inverse_scaling_vector = NULL;
         }
         hypredrv_ErrorCodeSet(ERROR_ALLOCATION);
         hypredrv_ErrorMsgAdd("Failed to cache the inverse scaling vector");
      }
   }
#else
   (void)mat_A;
   (void)vec_b;
   (void)dofmap;
   HYPREDRV_LOGF(
      2, log_rank, NULL, 0,
      "scaling requested but disabled (requires Hypre >= v3.0.0; linked release=%d)",
      (int)HYPREDRV_HYPRE_RELEASE_NUMBER);
   ctx->enabled = 0;
#endif
   HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling compute end (enabled=%d)", ctx->enabled);
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * Vector scaling helpers
 *-----------------------------------------------------------------------------*/

static int
ScalingTransformVectorRHSL2(const Scaling_context *ctx, HYPRE_IJVector vec,
                            scaling_vector_kind_t kind, int apply)
{
   /* GCOVR_EXCL_BR_START */
   MPI_Comm log_comm = ScalingCommFromVector(vec);
   if (!vec)
   {
      return 0;
   }

   HYPRE_Complex s = ctx->scalar_factor;
   if (s == 0.0)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("ScalingTransformVectorRHSL2: invalid scaling factor");
      HYPREDRV_LOG_COMMF(2, log_comm, NULL, 0,
                         "scaling vector transform failed: invalid scalar factor");
      return 0;
   }

   HYPRE_Complex factor = 1.0;
   if (kind == SCALING_VECTOR_RHS)
   {
      factor = apply ? s : (1.0 / s);
   }
   else
   {
      factor = apply ? (1.0 / s) : s;
   }

   void            *obj_vec = NULL;
   hypre_ParVector *par_vec = NULL;
   HYPRE_IJVectorGetObject(vec, &obj_vec);
   if (!obj_vec)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("ScalingTransformVectorRHSL2: vector object is NULL");
      HYPREDRV_LOG_COMMF(2, log_comm, NULL, 0,
                         "scaling vector transform failed: IJ vector object is NULL");
      return 0;
   }
   par_vec = (hypre_ParVector *)obj_vec;
   hypre_ParVectorScale(factor, par_vec);
   return 1;
   /* GCOVR_EXCL_BR_STOP */
}

static int
ScalingTransformVectorDofmap(const Scaling_context *ctx, HYPRE_IJVector vec,
                             scaling_vector_kind_t kind, int apply)
{
   /* GCOVR_EXCL_BR_START */
   MPI_Comm log_comm = ScalingCommFromVector(vec);
   if (!vec)
   {
      return 0;
   }

   if (!ctx->scaling_vector)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("ScalingTransformVectorDofmap: scaling vector not computed");
      HYPREDRV_LOG_COMMF(2, log_comm, NULL, 0,
                         "scaling vector transform failed: scaling vector not computed");
      return 0;
   }

   void            *obj_vec     = NULL;
   hypre_ParVector *par_vec     = NULL;
   hypre_ParVector *par_scaling = (hypre_ParVector *)ctx->scaling_vector;

   HYPRE_IJVectorGetObject(vec, &obj_vec);
   if (!obj_vec)
   {
      hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
      hypredrv_ErrorMsgAdd("ScalingTransformVectorDofmap: vector object is NULL");
      HYPREDRV_LOG_COMMF(2, log_comm, NULL, 0,
                         "scaling vector transform failed: IJ vector object is NULL");
      return 0;
   }
   par_vec = (hypre_ParVector *)obj_vec;

#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   if ((kind == SCALING_VECTOR_RHS && apply) ||
       (kind == SCALING_VECTOR_UNKNOWN && !apply))
   {
      hypre_ParVectorPointwiseProduct(par_scaling, par_vec, &par_vec);
   }
   else
   {
#if HYPRE_CHECK_MIN_VERSION(30100, 18)
      hypre_ParVectorPointwiseDivision(par_scaling, par_vec, &par_vec);
#else
      /* hypre_ParVectorPointwiseDivision(x, y, z) computes z=y/x on host but z=x/y on
       * device (GPU bug). Use inverse+product to get z=y/x=vec/scaling consistently. */
      hypre_ParVector *inv_scaling = (hypre_ParVector *)ctx->inverse_scaling_vector;
      if (!inv_scaling)
      {
         hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
         hypredrv_ErrorMsgAdd(
            "ScalingTransformVectorDofmap: inverse scaling vector not computed");
         return 0;
      }
      hypre_ParVectorPointwiseProduct(inv_scaling, par_vec, &par_vec);
#endif
   }
#else
   (void)kind;
   (void)apply;
   hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
   hypredrv_ErrorMsgAdd("ScalingTransformVectorDofmap: requires Hypre >= v3.0.0");
   HYPREDRV_LOG_COMMF(2, log_comm, NULL, 0,
                      "scaling vector transform failed: requires Hypre >= v3.0.0");
   return 0;
#endif
   return 1;
   /* GCOVR_EXCL_BR_STOP */
}

static int
ScalingTransformVector(const Scaling_context *ctx, HYPRE_IJVector vec,
                       scaling_vector_kind_t kind, int apply)
{
   /* GCOVR_EXCL_BR_START */
   if (!ctx || !ctx->enabled || !vec)
   {
      return 0;
   }

#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   switch (ctx->type)
   {
      case SCALING_RHS_L2:
         return ScalingTransformVectorRHSL2(ctx, vec, kind, apply);

      case SCALING_DOFMAP_MAG:
      case SCALING_DOFMAP_CUSTOM:
         return ScalingTransformVectorDofmap(ctx, vec, kind, apply);

      case SCALING_DOFMAP_ROW_CUSTOM:
         if (kind == SCALING_VECTOR_RHS)
         {
            return ScalingTransformVectorDofmap(ctx, vec, kind, apply);
         }
         return 0;

      case SCALING_DOFMAP_COL_CUSTOM:
         if (kind == SCALING_VECTOR_UNKNOWN)
         {
            return ScalingTransformVectorDofmap(ctx, vec, kind, apply);
         }
         return 0;

      case SCALING_DOFMAP_SIMILARITY_CUSTOM:
         if (kind == SCALING_VECTOR_RHS)
         {
            /* c = S^-1 b */
            return ScalingTransformVectorDofmap(ctx, vec, kind, !apply);
         }
         else
         {
            /* y = S^-1 x on apply; x = S y on undo */
            return ScalingTransformVectorDofmap(ctx, vec, kind, apply);
         }

      default:
         hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
         hypredrv_ErrorMsgAdd(apply
                                 ? "hypredrv_ScalingApplyToVector: unknown scaling type"
                                 : "hypredrv_ScalingUndoOnVector: unknown scaling type");
         if (apply)
         {
            HYPREDRV_LOG_COMMF(2, ScalingCommFromVector(vec), NULL, 0,
                               "scaling apply-to-vector failed: unknown scaling type=%d",
                               (int)ctx->type);
         }
         else
         {
            HYPREDRV_LOG_COMMF(2, ScalingCommFromVector(vec), NULL, 0,
                               "scaling undo-on-vector failed: unknown scaling type=%d",
                               (int)ctx->type);
         }
         return 0;
   }
#else
   (void)kind;
   if (apply)
   {
      HYPREDRV_LOG_COMMF(2, ScalingCommFromVector(vec), NULL, 0,
                         "scaling apply-to-vector failed: requires Hypre >= v3.0.0");
   }
   else
   {
      HYPREDRV_LOG_COMMF(2, ScalingCommFromVector(vec), NULL, 0,
                         "scaling undo-on-vector failed: requires Hypre >= v3.0.0");
   }
#endif
   return 0;
   /* GCOVR_EXCL_BR_STOP */
}

void
hypredrv_ScalingApplyToVector(const Scaling_context *ctx, HYPRE_IJVector vec,
                              scaling_vector_kind_t kind)
{
   ScalingTransformVector(ctx, vec, kind, 1);
}

void
hypredrv_ScalingUndoOnVector(const Scaling_context *ctx, HYPRE_IJVector vec,
                             scaling_vector_kind_t kind)
{
   ScalingTransformVector(ctx, vec, kind, 0);
}

/*-----------------------------------------------------------------------------
 * System scaling helpers
 *-----------------------------------------------------------------------------*/

#if HYPRE_CHECK_MIN_VERSION(30000, 0)
static void
ScalingDiagScaleMatrices(HYPRE_ParCSRMatrix par_A, HYPRE_ParCSRMatrix par_M,
                         hypre_ParVector *left, hypre_ParVector *right)
{
   hypre_ParCSRMatrixDiagScale(par_A, left, right);
   if (par_M != par_A)
   {
      hypre_ParCSRMatrixDiagScale(par_M, left, right);
   }
}

static void
ScalingDofmapMatrixFactors(const Scaling_context *ctx, int apply,
                           hypre_ParVector **left_ptr, hypre_ParVector **right_ptr)
{
   hypre_ParVector *scaling = (hypre_ParVector *)ctx->scaling_vector;
   hypre_ParVector *inverse = (hypre_ParVector *)ctx->inverse_scaling_vector;

   switch (ctx->type)
   {
      case SCALING_DOFMAP_ROW_CUSTOM:
         *left_ptr  = apply ? scaling : inverse;
         *right_ptr = NULL;
         break;
      case SCALING_DOFMAP_COL_CUSTOM:
         *left_ptr  = NULL;
         *right_ptr = apply ? scaling : inverse;
         break;
      case SCALING_DOFMAP_SIMILARITY_CUSTOM:
         *left_ptr  = apply ? inverse : scaling;
         *right_ptr = apply ? scaling : inverse;
         break;
      default:
         *left_ptr = *right_ptr = apply ? scaling : inverse;
         break;
   }
}
#endif

/* Cleanup must not mistake an error from the failed setup/solve for an error
 * raised by the inverse transforms themselves. Keep queued messages intact,
 * temporarily remove the saved bits, then merge them back after restoration. */
static void
ScalingRestoreErrorCode(uint32_t saved_error)
{
   for (uint32_t i = 0; i < 32; i++)
   {
      uint32_t bit = UINT32_C(1) << i;
      if (saved_error & bit)
      {
         hypredrv_ErrorCodeSet((hypredrv_error_t)bit);
      }
   }
}

static void
ScalingUpdateAppliedState(Scaling_context *ctx)
{
   ctx->is_applied = ctx->matrices_are_scaled || ctx->rhs_is_scaled || ctx->x_is_scaled;
}

static void
ScalingTransformSystem(Scaling_context *ctx, HYPRE_IJMatrix mat_A, HYPRE_IJMatrix mat_M,
                       HYPRE_IJVector vec_b, HYPRE_IJVector vec_x, int apply)
{
   /* GCOVR_EXCL_BR_START */
   MPI_Comm log_comm = ScalingSystemCommResolve(mat_A, mat_M, vec_b, vec_x);
   int      log_rank = -1;
   if (hypredrv_LogEnabled(2))
   {
      log_rank = hypredrv_LogRankFromComm(log_comm);
   }
   if (!ctx || !ctx->enabled || (!apply && !ctx->is_applied))
   {
      return;
   }

   if (apply && ctx->is_applied)
   {
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling apply skipped (already applied)");
      return;
   }

   uint32_t saved_error = ERROR_NONE;
   if (!apply)
   {
      saved_error = hypredrv_ErrorCodeGet();
      hypredrv_ErrorCodeReset(saved_error);
   }

   if (apply)
   {
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling apply begin (type=%d)",
                    (int)ctx->type);
   }
   else
   {
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling undo begin (type=%d)", (int)ctx->type);
   }

#if HYPRE_CHECK_MIN_VERSION(30000, 0)
   void              *obj_A = NULL, *obj_M = NULL;
   HYPRE_ParCSRMatrix par_A = NULL, par_M = NULL;

   switch (ctx->type)
   {
      case SCALING_RHS_L2:
         if (ctx->scalar_factor == 0.0)
         {
            hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
            hypredrv_ErrorMsgAdd(apply ? "ScalingApplyRHSL2: invalid scaling factor"
                                       : "ScalingUndoRHSL2: invalid scaling factor");
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          apply ? "scaling apply failed: invalid scalar factor"
                                : "scaling undo failed: invalid scalar factor");
            goto done;
         }
         break;

      case SCALING_DOFMAP_MAG:
      case SCALING_DOFMAP_CUSTOM:
      case SCALING_DOFMAP_ROW_CUSTOM:
      case SCALING_DOFMAP_COL_CUSTOM:
      case SCALING_DOFMAP_SIMILARITY_CUSTOM:
         if (!ctx->scaling_vector || !ctx->inverse_scaling_vector)
         {
            hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
            hypredrv_ErrorMsgAdd(apply
                                    ? "ScalingApplyDofmap: scaling vectors not computed"
                                    : "ScalingUndoDofmap: scaling vectors not computed");
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          apply ? "scaling apply failed: scaling vectors not computed"
                                : "scaling undo failed: scaling vectors not computed");
            goto done;
         }
         break;

      default:
         hypredrv_ErrorCodeSet(ERROR_UNKNOWN);
         hypredrv_ErrorMsgAdd(apply
                                 ? "hypredrv_ScalingApplyToSystem: unknown scaling type"
                                 : "hypredrv_ScalingUndoOnSystem: unknown scaling type");
         if (apply)
         {
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          "scaling apply failed: unknown scaling type=%d",
                          (int)ctx->type);
         }
         else
         {
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          "scaling undo failed: unknown scaling type=%d", (int)ctx->type);
         }
         goto done;
   }

   HYPRE_IJMatrixGetObject(mat_A, &obj_A);
   par_A = (HYPRE_ParCSRMatrix)obj_A;

   if (mat_M != mat_A)
   {
      HYPRE_IJMatrixGetObject(mat_M, &obj_M);
      par_M = (HYPRE_ParCSRMatrix)obj_M;
   }
   else
   {
      par_M = par_A;
   }

   if (apply)
   {
      if (ctx->type == SCALING_RHS_L2)
      {
         HYPRE_Complex s2 = ctx->scalar_factor * ctx->scalar_factor;
         hypre_ParCSRMatrixScale(par_A, s2);
         if (par_M != par_A)
         {
            hypre_ParCSRMatrixScale(par_M, s2);
         }
      }
      else
      {
         hypre_ParVector *left = NULL, *right = NULL;
         ScalingDofmapMatrixFactors(ctx, 1, &left, &right);
         ScalingDiagScaleMatrices(par_A, par_M, left, right);
      }
      ctx->matrices_are_scaled = 1;
      ScalingUpdateAppliedState(ctx);

      if (ScalingTransformVector(ctx, vec_b, SCALING_VECTOR_RHS, 1))
      {
         ctx->rhs_is_scaled = 1;
         ScalingUpdateAppliedState(ctx);
      }
      if (hypredrv_ErrorCodeGet())
      {
         HYPREDRV_LOGF(2, log_rank, NULL, 0,
                       "scaling apply failed: RHS vector transform error");
         goto done;
      }
      if (ScalingTransformVector(ctx, vec_x, SCALING_VECTOR_UNKNOWN, 1))
      {
         ctx->x_is_scaled = 1;
         ScalingUpdateAppliedState(ctx);
      }
      if (hypredrv_ErrorCodeGet())
      {
         HYPREDRV_LOGF(2, log_rank, NULL, 0,
                       "scaling apply failed: unknown vector transform error");
         goto done;
      }
   }
   else
   {
      if (ctx->x_is_scaled)
      {
         if (ScalingTransformVector(ctx, vec_x, SCALING_VECTOR_UNKNOWN, 0))
         {
            ctx->x_is_scaled = 0;
            ScalingUpdateAppliedState(ctx);
         }
         if (hypredrv_ErrorCodeGet())
         {
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          "scaling undo failed: unknown vector transform error");
            goto done;
         }
      }
      if (ctx->rhs_is_scaled)
      {
         if (ScalingTransformVector(ctx, vec_b, SCALING_VECTOR_RHS, 0))
         {
            ctx->rhs_is_scaled = 0;
            ScalingUpdateAppliedState(ctx);
         }
         if (hypredrv_ErrorCodeGet())
         {
            HYPREDRV_LOGF(2, log_rank, NULL, 0,
                          "scaling undo failed: RHS vector transform error");
            goto done;
         }
      }

      if (ctx->matrices_are_scaled && ctx->type == SCALING_RHS_L2)
      {
         HYPRE_Complex s2     = ctx->scalar_factor * ctx->scalar_factor;
         HYPRE_Complex inv_s2 = 1.0 / s2;
         hypre_ParCSRMatrixScale(par_A, inv_s2);
         if (par_M != par_A)
         {
            hypre_ParCSRMatrixScale(par_M, inv_s2);
         }
         ctx->matrices_are_scaled = 0;
         ScalingUpdateAppliedState(ctx);
      }
      else if (ctx->matrices_are_scaled)
      {
         hypre_ParVector *left = NULL, *right = NULL;
         ScalingDofmapMatrixFactors(ctx, 0, &left, &right);
         ScalingDiagScaleMatrices(par_A, par_M, left, right);
         ctx->matrices_are_scaled = 0;
         ScalingUpdateAppliedState(ctx);
      }
   }
#else
   (void)mat_A;
   (void)mat_M;
   (void)vec_b;
   (void)vec_x;
   if (apply)
   {
      HYPREDRV_LOGF(2, log_rank, NULL, 0,
                    "scaling apply failed: requires Hypre >= v3.0.0");
   }
   else
   {
      HYPREDRV_LOGF(2, log_rank, NULL, 0,
                    "scaling undo failed: requires Hypre >= v3.0.0");
   }
#endif

done:
   ScalingUpdateAppliedState(ctx);
   if (!apply)
   {
      ScalingRestoreErrorCode(saved_error);
   }
   if (apply)
   {
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling apply end (is_applied=%d)",
                    ctx->is_applied);
   }
   else
   {
      HYPREDRV_LOGF(3, log_rank, NULL, 0, "scaling undo end (is_applied=%d)",
                    ctx->is_applied);
   }
   /* GCOVR_EXCL_BR_STOP */
}

/*-----------------------------------------------------------------------------
 * System scaling public wrappers
 *-----------------------------------------------------------------------------*/

void
hypredrv_ScalingApplyToSystem(Scaling_context *ctx, HYPRE_IJMatrix mat_A,
                              HYPRE_IJMatrix mat_M, HYPRE_IJVector vec_b,
                              HYPRE_IJVector vec_x)
{
   ScalingTransformSystem(ctx, mat_A, mat_M, vec_b, vec_x, 1);
}

void
hypredrv_ScalingUndoOnSystem(Scaling_context *ctx, HYPRE_IJMatrix mat_A,
                             HYPRE_IJMatrix mat_M, HYPRE_IJVector vec_b,
                             HYPRE_IJVector vec_x)
{
   ScalingTransformSystem(ctx, mat_A, mat_M, vec_b, vec_x, 0);
}
