// Copyright (c) 2010-2024, 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 "pml.hpp"

namespace mfem
{

CartesianPML::CartesianPML(Mesh *mesh_, const Array2D<real_t> &length_)
   : mesh(mesh_), length(length_)
{
   dim = mesh->Dimension();
   SetBoundaries();
}

void CartesianPML::SetBoundaries()
{
   comp_dom_bdr.SetSize(dim, 2);
   dom_bdr.SetSize(dim, 2);
   // initialize
   for (int i = 0; i < dim; i++)
   {
      dom_bdr(i, 0) = infinity();
      dom_bdr(i, 1) = -infinity();
   }

   for (int i = 0; i < mesh->GetNBE(); i++)
   {
      Array<int> bdr_vertices;
      mesh->GetBdrElementVertices(i, bdr_vertices);
      for (int j = 0; j < bdr_vertices.Size(); j++)
      {
         for (int k = 0; k < dim; k++)
         {
            dom_bdr(k, 0) = std::min(dom_bdr(k, 0), mesh->GetVertex(bdr_vertices[j])[k]);
            dom_bdr(k, 1) = std::max(dom_bdr(k, 1), mesh->GetVertex(bdr_vertices[j])[k]);
         }
      }
   }

#ifdef MFEM_USE_MPI
   ParMesh * pmesh = dynamic_cast<ParMesh *>(mesh);
   if (pmesh)
   {
      for (int d=0; d<dim; d++)
      {
         MPI_Allreduce(MPI_IN_PLACE, &dom_bdr(d,0), 1,
                       MPITypeMap<real_t>::mpi_type, MPI_MIN,pmesh->GetComm());
         MPI_Allreduce(MPI_IN_PLACE, &dom_bdr(d,1), 1,
                       MPITypeMap<real_t>::mpi_type, MPI_MAX, pmesh->GetComm());
      }
   }
#endif

   for (int i = 0; i < dim; i++)
   {
      comp_dom_bdr(i, 0) = dom_bdr(i, 0) + length(i, 0);
      comp_dom_bdr(i, 1) = dom_bdr(i, 1) - length(i, 1);
   }
}

void CartesianPML::SetAttributes(Mesh *mesh_, Array<int> * attrNonPML,
                                 Array<int> * attrPML)
{
   int nrelem = mesh_->GetNE();
   elems.SetSize(nrelem);

   for (int i = 0; i < nrelem; ++i)
   {
      elems[i] = 1;
      bool in_pml = false;
      Element *el = mesh_->GetElement(i);
      Array<int> vertices;
      // Initialize Attribute
      el->SetAttribute(1);
      el->GetVertices(vertices);
      int nrvert = vertices.Size();
      // Check if any vertex is in the pml
      for (int iv = 0; iv < nrvert; ++iv)
      {
         int vert_idx = vertices[iv];
         real_t *coords = mesh_->GetVertex(vert_idx);
         for (int comp = 0; comp < dim; ++comp)
         {
            if (coords[comp] > comp_dom_bdr(comp, 1) ||
                coords[comp] < comp_dom_bdr(comp, 0))
            {
               in_pml = true;
               break;
            }
         }
      }
      if (in_pml)
      {
         elems[i] = 0;
         el->SetAttribute(2);
      }
   }
   mesh_->SetAttributes();

   if (mesh_->attributes.Size())
   {
      if (attrNonPML)
      {
         attrNonPML->SetSize(mesh_->attributes.Max());
         *attrNonPML = 0; (*attrNonPML)[0] = 1;

      }
      if (attrPML)
      {
         attrPML->SetSize(mesh_->attributes.Max());
         *attrPML = 0;
         if (mesh_->attributes.Max()>1)
         {
            (*attrPML)[1]=1;
         }
      }
   }

}

void CartesianPML::StretchFunction(const Vector &x,
                                   std::vector<std::complex<real_t>> &dxs)
{
   std::complex<real_t> zi = std::complex<real_t>(0., 1.);

   real_t n = 2.0;
   real_t c = 5.0;
   real_t coeff;
   real_t k = omega * sqrt(epsilon * mu);
   // Stretch in each direction independently
   for (int i = 0; i < dim; ++i)
   {
      dxs[i] = 1.0;
      if (x(i) >= comp_dom_bdr(i, 1))
      {
         coeff = n * c / k / pow(length(i, 1), n);
         dxs[i] = real_t(1.0) + zi * real_t(coeff * std::abs(pow(x(i) - comp_dom_bdr(i,
                                                                                     1), n - 1.0)));
      }
      if (x(i) <= comp_dom_bdr(i, 0))
      {
         coeff = n * c / k / pow(length(i, 0), n);
         dxs[i] = real_t(1.0) + zi * real_t(coeff * std::abs(pow(x(i) - comp_dom_bdr(i,
                                                                                     0), n - 1.0)));
      }
   }
}

// acoustics UW PML coefficients functions
// |J|
real_t detJ_r_function(const Vector & x, CartesianPML * pml)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det(1.0,0.0);
   pml->StretchFunction(x, dxs);
   for (int i=0; i<dim; ++i) { det *= dxs[i]; }
   return det.real();
}

real_t detJ_i_function(const Vector & x, CartesianPML * pml)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det(1.0,0.0);
   pml->StretchFunction(x, dxs);
   for (int i=0; i<dim; ++i) { det *= dxs[i]; }
   return det.imag();
}

real_t abs_detJ_2_function(const Vector & x, CartesianPML * pml)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det(1.0,0.0);
   pml->StretchFunction(x, dxs);
   for (int i=0; i<dim; ++i) { det *= dxs[i]; }
   return det.imag()*det.imag() + det.real()*det.real();
}

// J^T J / |J|
void Jt_J_detJinv_r_function(const Vector & x, CartesianPML * pml,
                             DenseMatrix & M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det(1.0,0.0);
   pml->StretchFunction(x, dxs);
   for (int i = 0; i<dim; ++i) { det *= dxs[i]; }

   M=0.0;
   for (int i = 0; i<dim; ++i)
   {
      M(i,i) = (pow(dxs[i], real_t(2))/det).real();
   }
}

void Jt_J_detJinv_i_function(const Vector & x, CartesianPML * pml,
                             DenseMatrix & M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det = 1.0;
   pml->StretchFunction(x, dxs);
   for (int i = 0; i<dim; ++i) { det *= dxs[i]; }

   M=0.0;
   for (int i = 0; i<dim; ++i)
   {
      M(i,i) = (pow(dxs[i], real_t(2))/det).imag();
   }
}

void abs_Jt_J_detJinv_2_function(const Vector & x, CartesianPML * pml,
                                 DenseMatrix & M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det = 1.0;
   pml->StretchFunction(x, dxs);
   for (int i = 0; i<dim; ++i) { det *= dxs[i]; }

   M=0.0;
   for (int i = 0; i<dim; ++i)
   {
      std::complex<real_t> a = pow(dxs[i], real_t(2))/det;
      M(i,i) = a.imag() * a.imag() + a.real() * a.real();
   }
}


// Maxwell PML coefficients
void detJ_Jt_J_inv_r_function(const Vector &x, CartesianPML * pml,
                              DenseMatrix &M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det(1.0, 0.0);
   pml->StretchFunction(x, dxs);

   for (int i = 0; i < dim; ++i) { det *= dxs[i]; }

   M = 0.0;
   for (int i = 0; i < dim; ++i)
   {
      M(i, i) = (det / pow(dxs[i], real_t(2))).real();
   }
}

void detJ_Jt_J_inv_i_function(const Vector &x, CartesianPML * pml,
                              DenseMatrix &M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det = 1.0;
   pml->StretchFunction(x, dxs);

   for (int i = 0; i < dim; ++i) { det *= dxs[i]; }

   M = 0.0;
   for (int i = 0; i < dim; ++i)
   {
      M(i, i) = (det / pow(dxs[i], real_t(2))).imag();
   }
}

void abs_detJ_Jt_J_inv_2_function(const Vector &x, CartesianPML * pml,
                                  DenseMatrix &M)
{
   int dim = pml->dim;
   std::vector<std::complex<real_t>> dxs(dim);
   std::complex<real_t> det = 1.0;
   pml->StretchFunction(x, dxs);

   for (int i = 0; i < dim; ++i) { det *= dxs[i]; }

   M = 0.0;
   for (int i = 0; i < dim; ++i)
   {
      std::complex<real_t> a = det / pow(dxs[i], real_t(2));
      M(i, i) = a.real()*a.real() + a.imag()*a.imag();
   }
}

} // namespace mfem
