// Copyright (c) 2010-2023, Lawrence Livermore National Security, LLC. Produced
// at the Lawrence Livermore National Laboratory. All Rights reserved. See files
// LICENSE and NOTICE for details. LLNL-CODE-806117.
//
// This file is part of the MFEM library. For more information and source code
// availability visit https://mfem.org.
//
// MFEM is free software; you can redistribute it and/or modify it under the
// terms of the BSD-3 license. We welcome feedback and contributions, see file
// CONTRIBUTING.md for details.

#include "lor_ams.hpp"
#include "../../general/forall.hpp"
#include "../../fem/pbilinearform.hpp"

namespace mfem
{

#ifdef MFEM_USE_MPI

void BatchedLOR_AMS::Form2DEdgeToVertex(Array<int> &edge2vert)
{
   const FiniteElementCollection *fec = edge_fes.FEColl();
   if (dynamic_cast<const ND_FECollection*>(fec))
   {
      Form2DEdgeToVertex_ND(edge2vert);
   }
   else if (dynamic_cast<const RT_FECollection*>(fec))
   {
      Form2DEdgeToVertex_RT(edge2vert);
   }
   else
   {
      MFEM_ABORT("Bad finite element type.")
   }
}

void BatchedLOR_AMS::Form2DEdgeToVertex_ND(Array<int> &edge2vert)
{
   const int o = order;
   const int op1 = o + 1;
   const int nedge = static_cast<int>(dim*o*pow(op1, dim-1));

   edge2vert.SetSize(2*nedge);
   auto e2v = Reshape(edge2vert.HostWrite(), 2, nedge);

   for (int c=0; c<dim; ++c)
   {
      const int nx = (c == 0) ? o : op1;
      for (int i2=0; i2<op1; ++i2)
      {
         for (int i1=0; i1<o; ++i1)
         {
            const int ix = (c == 0) ? i1 : i2;
            const int iy = (c == 0) ? i2 : i1;

            const int iedge = ix + iy*nx + c*o*op1;

            const int ix1 = (c == 0) ? ix + 1 : ix;
            const int iy1 = (c == 1) ? iy + 1 : iy;

            const int iv0 = ix + iy*op1;
            const int iv1 = ix1 + iy1*op1;

            e2v(0, iedge) = iv0;
            e2v(1, iedge) = iv1;
         }
      }
   }
}

void BatchedLOR_AMS::Form2DEdgeToVertex_RT(Array<int> &edge2vert)
{
   const int o = order;
   const int op1 = o + 1;
   const int nedge = static_cast<int>(dim*o*pow(op1, dim-1));

   edge2vert.SetSize(2*nedge);
   auto e2v = Reshape(edge2vert.HostWrite(), 2, nedge);

   for (int c=0; c<dim; ++c)
   {
      const int nx = (c == 0) ? op1 : o;
      for (int i=0; i<o*op1; ++i)
      {
         const int ix = i%nx;
         const int iy = i/nx;

         const int iedge = ix + iy*nx + c*o*op1;

         const int ix1 = (c == 0) ? ix : ix + 1;
         const int iy1 = (c == 1) ? iy : iy + 1;

         const int iv0 = ix + iy*op1;
         const int iv1 = ix1 + iy1*op1;

         // Rotated gradient in 2D (-dy, dx), so flip the sign for the first
         // component (c == 0).
         e2v(0, iedge) = (c == 1) ? iv0 : iv1;
         e2v(1, iedge) = (c == 1) ? iv1 : iv0;
      }
   }
}

void BatchedLOR_AMS::Form3DEdgeToVertex(Array<int> &edge2vert)
{
   const int o = order;
   const int op1 = o + 1;
   const int nedge = static_cast<int>(dim*o*pow(op1, dim-1));

   edge2vert.SetSize(2*nedge);
   auto e2v = Reshape(edge2vert.HostWrite(), 2, nedge);

   for (int c=0; c<dim; ++c)
   {
      const int nx = (c == 0) ? o : op1;
      const int ny = (c == 1) ? o : op1;
      for (int i=0; i<o*op1*op1; ++i)
      {
         const int ix = i%nx;
         const int iy = (i/nx)%ny;
         const int iz = i/nx/ny;

         const int iedge = ix + iy*nx + iz*nx*ny + c*o*op1*op1;

         const int ix1 = (c == 0) ? ix + 1 : ix;
         const int iy1 = (c == 1) ? iy + 1 : iy;
         const int iz1 = (c == 2) ? iz + 1 : iz;

         const int iv0 = ix + iy*op1 + iz*op1*op1;
         const int iv1 = ix1 + iy1*op1 + iz1*op1*op1;

         e2v(0, iedge) = iv0;
         e2v(1, iedge) = iv1;
      }
   }
}

void BatchedLOR_AMS::FormGradientMatrix()
{
   // The gradient matrix maps from LOR vertices to LOR edges. Given an edge
   // (defined by its two vertices) e_i = (v_j1, v_j2), the matrix has nonzeros
   // A(i, j1) = -1 and A(i, j2) = 1, so there are always exactly two nonzeros
   // per row.
   const int nedge_dof = edge_fes.GetNDofs();
   const int nvert_dof = vert_fes.GetNDofs();

   SparseMatrix G_local;
   G_local.OverrideSize(nedge_dof, nvert_dof);

   G_local.GetMemoryI().New(nedge_dof+1, Device::GetDeviceMemoryType());
   // Each row always has two nonzeros
   const int nnz = 2*nedge_dof;
   auto I = G_local.WriteI();
   mfem::forall(nedge_dof+1, [=] MFEM_HOST_DEVICE (int i) { I[i] = 2*i; });

   // edge2vertex is a mapping of size (2, nedge_per_el), such that with a macro
   // element, edge i (in lexicographic ordering) has vertices (also in
   // lexicographic ordering) given by the entries (0, i) and (1, i) of the
   // matrix.
   Array<int> edge2vertex;
   if (dim == 2) { Form2DEdgeToVertex(edge2vertex); }
   else { Form3DEdgeToVertex(edge2vertex); }

   ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
   const auto *R_v = dynamic_cast<const ElementRestriction*>(
                        vert_fes.GetElementRestriction(ordering));
   const auto *R_e = dynamic_cast<const ElementRestriction*>(
                        edge_fes.GetElementRestriction(ordering));
   MFEM_VERIFY(R_v != NULL && R_e != NULL, "");

   const int nel_ho = edge_fes.GetNE();
   const int nedge_per_el = static_cast<int>(dim*order*pow(order + 1, dim - 1));
   const int nvert_per_el = static_cast<int>(pow(order + 1, dim));

   const auto offsets_e = R_e->Offsets().Read();
   const auto indices_e = R_e->Indices().Read();
   const auto gather_v = Reshape(R_v->GatherMap().Read(), nvert_per_el, nel_ho);

   const auto e2v = Reshape(edge2vertex.Read(), 2, nedge_per_el);

   // Fill J and data
   G_local.GetMemoryJ().New(nnz, Device::GetDeviceMemoryType());
   G_local.GetMemoryData().New(nnz, Device::GetDeviceMemoryType());

   auto J = G_local.WriteJ();
   auto V = G_local.WriteData();

   // Loop over Nedelec L-DOFs
   mfem::forall(nedge_dof, [=] MFEM_HOST_DEVICE (int i)
   {
      const int sj = indices_e[offsets_e[i]]; // signed
      const int j = (sj >= 0) ? sj : -1 - sj;
      const int sgn = (sj >= 0) ? 1 : -1;
      const int j_loc = j%nedge_per_el;
      const int j_el = j/nedge_per_el;

      const int jv0_loc = e2v(0, j_loc);
      const int jv1_loc = e2v(1, j_loc);

      J[i*2 + 0] = gather_v(jv0_loc, j_el);
      J[i*2 + 1] = gather_v(jv1_loc, j_el);

      V[i*2 + 0] = -sgn;
      V[i*2 + 1] = sgn;
   });

   // Create a block diagonal parallel matrix
   OperatorHandle G_diag(Operator::Hypre_ParCSR);
   G_diag.MakeRectangularBlockDiag(vert_fes.GetComm(),
                                   edge_fes.GlobalVSize(),
                                   vert_fes.GlobalVSize(),
                                   edge_fes.GetDofOffsets(),
                                   vert_fes.GetDofOffsets(),
                                   &G_local);

   // Assemble the parallel gradient matrix, must be deleted by the caller
   if (IsIdentityProlongation(vert_fes.GetProlongationMatrix()))
   {
      G = G_diag.As<HypreParMatrix>();
      G_diag.SetOperatorOwner(false);
      HypreStealOwnership(*G, G_local);
   }
   else
   {
      OperatorHandle Rt(Transpose(*edge_fes.GetRestrictionMatrix()));
      OperatorHandle Rt_diag(Operator::Hypre_ParCSR);
      Rt_diag.MakeRectangularBlockDiag(edge_fes.GetComm(),
                                       edge_fes.GlobalVSize(),
                                       edge_fes.GlobalTrueVSize(),
                                       edge_fes.GetDofOffsets(),
                                       edge_fes.GetTrueDofOffsets(),
                                       Rt.As<SparseMatrix>());
      G = RAP(Rt_diag.As<HypreParMatrix>(),
              G_diag.As<HypreParMatrix>(),
              vert_fes.Dof_TrueDof_Matrix());
   }
   G->CopyRowStarts();
   G->CopyColStarts();
}

template <typename T>
static inline const T *HypreRead(const Memory<T> &mem)
{
   return mem.Read(GetHypreMemoryClass(), mem.Capacity());
}

template <typename T>
static inline T *HypreWrite(Memory<T> &mem)
{
   return mem.Write(GetHypreMemoryClass(), mem.Capacity());
}

void BatchedLOR_AMS::FormCoordinateVectors(const Vector &X_vert)
{
   // Create true-DOF vectors x, y, and z that contain the coordinates of the
   // vertices of the LOR mesh. The vertex coordinates are already computed in
   // E-vector format and passed in in X_vert.
   //
   // In this function, we need to convert X_vert (which has the shape (dim,
   // ndof_per_el, nel_ho)) to T-DOF format.
   //
   // We place the results in the vector xyz_tvec, which has shape (ntdofs, dim)
   // and then make the hypre vectors x, y, and z point to subvectors.
   //
   // In 2D, z is NULL.

   // Create the H1 vertex space and get the element restriction
   ElementDofOrdering ordering = ElementDofOrdering::LEXICOGRAPHIC;
   const Operator *op = vert_fes.GetElementRestriction(ordering);
   const auto *el_restr = dynamic_cast<const ElementRestriction*>(op);
   MFEM_VERIFY(el_restr != NULL, "");
   const SparseMatrix *R = vert_fes.GetRestrictionMatrix();

   const int nel_ho = vert_fes.GetNE();
   const int ndp1 = order + 1;
   const int ndof_per_el = static_cast<int>(pow(ndp1, dim));
   const int sdim = dim;
   const int ntdofs = R->Height();

   const MemoryClass mc = GetHypreMemoryClass();
   bool dev = (mc == MemoryClass::DEVICE);

   xyz_tvec = new Vector(ntdofs*dim);

   auto xyz_tv = Reshape(HypreWrite(xyz_tvec->GetMemory()), ntdofs, dim);
   const auto xyz_e =
      Reshape(HypreRead(X_vert.GetMemory()), dim, ndof_per_el, nel_ho);
   const auto d_offsets = HypreRead(el_restr->Offsets().GetMemory());
   const auto d_indices = HypreRead(el_restr->Indices().GetMemory());
   const auto ltdof_ldof = HypreRead(R->GetMemoryJ());

   // Go from E-vector format directly to T-vector format
   MFEM_HYPRE_FORALL(i, ntdofs,
   {
      const int j = d_offsets[ltdof_ldof[i]];
      for (int c = 0; c < sdim; ++c)
      {
         const int idx_j = d_indices[j];
         xyz_tv(i,c) = xyz_e(c, idx_j%ndof_per_el, idx_j/ndof_per_el);
      }
   });

   // Make x, y, z HypreParVectors point to T-vector data
   HYPRE_BigInt glob_size = vert_fes.GlobalTrueVSize();
   HYPRE_BigInt *cols = vert_fes.GetTrueDofOffsets();

   double *d_x_ptr = xyz_tv + 0*ntdofs;
   x = new HypreParVector(vert_fes.GetComm(), glob_size, d_x_ptr, cols, dev);
   double *d_y_ptr = xyz_tv + 1*ntdofs;
   y = new HypreParVector(vert_fes.GetComm(), glob_size, d_y_ptr, cols, dev);
   if (dim == 3)
   {
      double *d_z_ptr = xyz_tv + 2*ntdofs;
      z = new HypreParVector(vert_fes.GetComm(), glob_size, d_z_ptr, cols, dev);
   }
   else
   {
      z = NULL;
   }
}

HypreParMatrix *BatchedLOR_AMS::StealGradientMatrix()
{
   return StealPointer(G);
}

Vector *BatchedLOR_AMS::StealCoordinateVector()
{
   return StealPointer(xyz_tvec);
}

HypreParVector *BatchedLOR_AMS::StealXCoordinate()
{
   return StealPointer(x);
}

HypreParVector *BatchedLOR_AMS::StealYCoordinate()
{
   return StealPointer(y);
}

HypreParVector *BatchedLOR_AMS::StealZCoordinate()
{
   return StealPointer(z);
}

BatchedLOR_AMS::~BatchedLOR_AMS()
{
   delete x;
   delete y;
   delete z;
   delete xyz_tvec;
   delete G;
}

BatchedLOR_AMS::BatchedLOR_AMS(ParFiniteElementSpace &pfes_ho_,
                               const Vector &X_vert)
   : edge_fes(pfes_ho_),
     dim(edge_fes.GetParMesh()->Dimension()),
     order(edge_fes.GetMaxElementOrder()),
     vert_fec(order, dim),
     vert_fes(edge_fes.GetParMesh(), &vert_fec)
{
   FormCoordinateVectors(X_vert);
   FormGradientMatrix();
}

LORSolver<HypreAMS>::LORSolver(
   ParBilinearForm &a_ho, const Array<int> &ess_tdof_list, int ref_type)
{
   if (BatchedLORAssembly::FormIsSupported(a_ho))
   {
      ParFiniteElementSpace &pfes = *a_ho.ParFESpace();
      BatchedLORAssembly batched_lor(pfes);
      batched_lor.Assemble(a_ho, ess_tdof_list, A);
      BatchedLOR_AMS lor_ams(pfes, batched_lor.GetLORVertexCoordinates());
      xyz = lor_ams.StealCoordinateVector();
      solver = new HypreAMS(*A.As<HypreParMatrix>(),
                            lor_ams.StealGradientMatrix(),
                            lor_ams.StealXCoordinate(),
                            lor_ams.StealYCoordinate(),
                            lor_ams.StealZCoordinate());
   }
   else
   {
      ParLORDiscretization lor(a_ho, ess_tdof_list, ref_type);
      // Assume ownership of the system matrix so that `lor` can be safely
      // deleted
      A.Reset(lor.GetAssembledSystem().Ptr());
      lor.GetAssembledSystem().SetOperatorOwner(false);
      solver = new HypreAMS(lor.GetAssembledMatrix(), &lor.GetParFESpace());
   }
   width = solver->Width();
   height = solver->Height();
}

void LORSolver<HypreAMS>::SetOperator(const Operator &op)
{
   solver->SetOperator(op);
}

void LORSolver<HypreAMS>::Mult(const Vector &x, Vector &y) const
{
   solver->Mult(x, y);
}

HypreAMS &LORSolver<HypreAMS>::GetSolver() { return *solver; }

const HypreAMS &LORSolver<HypreAMS>::GetSolver() const { return *solver; }

LORSolver<HypreAMS>::~LORSolver()
{
   delete solver;
   delete xyz;
}

#endif

} // namespace mfem