#include <mesh/MeshSmoother.hpp>

#include <algebra/TinyMatrix.hpp>
#include <algebra/TinyVector.hpp>
#include <language/utils/InterpolateItemValue.hpp>
#include <mesh/Connectivity.hpp>
#include <mesh/ItemValueUtils.hpp>
#include <mesh/Mesh.hpp>
#include <mesh/MeshFlatNodeBoundary.hpp>
#include <mesh/MeshLineNodeBoundary.hpp>
#include <mesh/MeshNodeBoundary.hpp>
#include <scheme/AxisBoundaryConditionDescriptor.hpp>
#include <scheme/FixedBoundaryConditionDescriptor.hpp>
#include <scheme/SymmetryBoundaryConditionDescriptor.hpp>
#include <utils/RandomEngine.hpp>

#include <variant>

template <size_t Dimension>
class MeshSmootherHandler::MeshSmoother
{
 private:
  using Rd               = TinyVector<Dimension>;
  using Rdxd             = TinyMatrix<Dimension>;
  using ConnectivityType = Connectivity<Dimension>;
  using MeshType         = Mesh<ConnectivityType>;

  const MeshType& m_given_mesh;

  class AxisBoundaryCondition;
  class FixedBoundaryCondition;
  class SymmetryBoundaryCondition;

  using BoundaryCondition = std::variant<AxisBoundaryCondition, FixedBoundaryCondition, SymmetryBoundaryCondition>;

  using BoundaryConditionList = std::vector<BoundaryCondition>;
  BoundaryConditionList m_boundary_condition_list;

  BoundaryConditionList
  _getBCList(const MeshType& mesh,
             const std::vector<std::shared_ptr<const IBoundaryConditionDescriptor>>& bc_descriptor_list)
  {
    BoundaryConditionList bc_list;

    for (const auto& bc_descriptor : bc_descriptor_list) {
      switch (bc_descriptor->type()) {
      case IBoundaryConditionDescriptor::Type::axis: {
        if constexpr (Dimension == 1) {
          bc_list.emplace_back(FixedBoundaryCondition{getMeshNodeBoundary(mesh, bc_descriptor->boundaryDescriptor())});
        } else {
          bc_list.emplace_back(
            AxisBoundaryCondition{getMeshLineNodeBoundary(mesh, bc_descriptor->boundaryDescriptor())});
        }
        break;
      }
      case IBoundaryConditionDescriptor::Type::symmetry: {
        bc_list.emplace_back(
          SymmetryBoundaryCondition{getMeshFlatNodeBoundary(mesh, bc_descriptor->boundaryDescriptor())});
        break;
      }
      case IBoundaryConditionDescriptor::Type::fixed: {
        bc_list.emplace_back(FixedBoundaryCondition{getMeshNodeBoundary(mesh, bc_descriptor->boundaryDescriptor())});
        break;
      }
      default: {
        std::ostringstream error_msg;
        error_msg << *bc_descriptor << " is an invalid boundary condition for mesh smoother";
        throw NormalError(error_msg.str());
      }
      }
    }

    return bc_list;
  }

  void
  _applyBC(NodeValue<Rd>& shift) const
  {
    for (auto&& boundary_condition : m_boundary_condition_list) {
      std::visit(
        [&](auto&& bc) {
          using BCType = std::decay_t<decltype(bc)>;
          if constexpr (std::is_same_v<BCType, SymmetryBoundaryCondition>) {
            const Rd& n = bc.outgoingNormal();

            const Rdxd I   = identity;
            const Rdxd nxn = tensorProduct(n, n);
            const Rdxd P   = I - nxn;

            const Array<const NodeId>& node_list = bc.nodeList();
            parallel_for(
              node_list.size(), PUGS_LAMBDA(const size_t i_node) {
                const NodeId node_id = node_list[i_node];

                shift[node_id] = P * shift[node_id];
              });

          } else if constexpr (std::is_same_v<BCType, AxisBoundaryCondition>) {
            if constexpr (Dimension > 1) {
              const Rd& t = bc.direction();

              const Rdxd txt = tensorProduct(t, t);

              const Array<const NodeId>& node_list = bc.nodeList();
              parallel_for(
                node_list.size(), PUGS_LAMBDA(const size_t i_node) {
                  const NodeId node_id = node_list[i_node];

                  shift[node_id] = txt * shift[node_id];
                });
            } else {
              throw UnexpectedError("AxisBoundaryCondition make no sense in dimension 1");
            }

          } else if constexpr (std::is_same_v<BCType, FixedBoundaryCondition>) {
            const Array<const NodeId>& node_list = bc.nodeList();
            parallel_for(
              node_list.size(), PUGS_LAMBDA(const size_t i_node) {
                const NodeId node_id = node_list[i_node];
                shift[node_id]       = zero;
              });

          } else {
            throw UnexpectedError("invalid boundary condition type");
          }
        },
        boundary_condition);
    }
  }

  NodeValue<Rd>
  _getDisplacement() const
  {
    const ConnectivityType& connectivity = m_given_mesh.connectivity();
    NodeValue<const Rd> given_xr         = m_given_mesh.xr();

    auto node_to_cell_matrix        = connectivity.nodeToCellMatrix();
    auto cell_to_node_matrix        = connectivity.cellToNodeMatrix();
    auto node_number_in_their_cells = connectivity.nodeLocalNumbersInTheirCells();

    NodeValue<double> max_delta_xr{connectivity};
    parallel_for(
      connectivity.numberOfNodes(), PUGS_LAMBDA(const NodeId node_id) {
        const Rd& x0 = given_xr[node_id];

        const auto& node_cell_list = node_to_cell_matrix[node_id];
        double min_distance_2      = std::numeric_limits<double>::max();

        for (size_t i_cell = 0; i_cell < node_cell_list.size(); ++i_cell) {
          const size_t i_cell_node = node_number_in_their_cells(node_id, i_cell);

          const CellId cell_id       = node_cell_list[i_cell];
          const auto& cell_node_list = cell_to_node_matrix[cell_id];

          for (size_t i_node = 0; i_node < cell_node_list.size(); ++i_node) {
            if (i_node != i_cell_node) {
              const NodeId cell_node_id = cell_node_list[i_node];
              const Rd delta            = x0 - given_xr[cell_node_id];
              min_distance_2            = std::min(min_distance_2, dot(delta, delta));
            }
          }
        }
        double max_delta = std::sqrt(min_distance_2);

        max_delta_xr[node_id] = max_delta;
      });

    NodeValue<Rd> shift_r{connectivity};

    parallel_for(
      m_given_mesh.numberOfNodes(), PUGS_LAMBDA(const NodeId node_id) {
        const auto& node_cell_list = node_to_cell_matrix[node_id];
        Rd mean_position(zero);
        size_t number_of_neighbours = 0;

        for (size_t i_cell = 0; i_cell < node_cell_list.size(); ++i_cell) {
          const size_t i_cell_node = node_number_in_their_cells(node_id, i_cell);

          const CellId cell_id       = node_cell_list[i_cell];
          const auto& cell_node_list = cell_to_node_matrix[cell_id];
          for (size_t i_node = 0; i_node < cell_node_list.size(); ++i_node) {
            if (i_node != i_cell_node) {
              const NodeId cell_node_id = cell_node_list[i_node];
              mean_position += given_xr[cell_node_id];
              number_of_neighbours++;
              // std::cout << node_id << " position " << given_xr[node_id] << " position voisin " <<
              // given_xr[cell_node_id]
              //           << " mean_position " << mean_position << "\n";
            }
          }
        }
        mean_position    = 1. / number_of_neighbours * mean_position;
        shift_r[node_id] = mean_position - given_xr[node_id];
        // std::cout << " mean position " << mean_position << given_xr[node_id] << "\n";

        // double nshift    = sqrt(dot(shift_r[node_id], shift_r[node_id]));
        // if (nshift > max_delta_xr[node_id]) {
        //   shift_r[node_id] = max_delta_xr[node_id] / nshift * shift_r[node_id];
        // }
      });

    this->_applyBC(shift_r);

    synchronize(shift_r);

    return shift_r;
  }

 public:
  std::shared_ptr<const IMesh>
  getSmoothedMesh() const
  {
    NodeValue<const Rd> given_xr = m_given_mesh.xr();

    NodeValue<Rd> xr = this->_getDisplacement();

    parallel_for(
      m_given_mesh.numberOfNodes(), PUGS_LAMBDA(const NodeId node_id) { xr[node_id] += given_xr[node_id]; });

    return std::make_shared<MeshType>(m_given_mesh.shared_connectivity(), xr);
  }

  std::shared_ptr<const IMesh>
  getSmoothedMesh(const FunctionSymbolId& function_symbol_id) const
  {
    NodeValue<const Rd> given_xr = m_given_mesh.xr();

    NodeValue<const bool> is_displaced =
      InterpolateItemValue<bool(const Rd)>::interpolate(function_symbol_id, given_xr);

    NodeValue<Rd> xr = this->_getDisplacement();

    parallel_for(
      m_given_mesh.numberOfNodes(),
      PUGS_LAMBDA(const NodeId node_id) { xr[node_id] = is_displaced[node_id] * xr[node_id] + given_xr[node_id]; });

    return std::make_shared<MeshType>(m_given_mesh.shared_connectivity(), xr);
  }

  MeshSmoother(const MeshSmoother&) = delete;
  MeshSmoother(MeshSmoother&&)      = delete;

  MeshSmoother(const MeshType& given_mesh,
               const std::vector<std::shared_ptr<const IBoundaryConditionDescriptor>>& bc_descriptor_list)
    : m_given_mesh(given_mesh), m_boundary_condition_list(this->_getBCList(given_mesh, bc_descriptor_list))
  {}

  ~MeshSmoother() = default;
};

template <size_t Dimension>
class MeshSmootherHandler::MeshSmoother<Dimension>::AxisBoundaryCondition
{
 public:
  using Rd = TinyVector<Dimension, double>;

 private:
  const MeshLineNodeBoundary<Dimension> m_mesh_line_node_boundary;

 public:
  const Rd&
  direction() const
  {
    return m_mesh_line_node_boundary.direction();
  }

  const Array<const NodeId>&
  nodeList() const
  {
    return m_mesh_line_node_boundary.nodeList();
  }

  AxisBoundaryCondition(MeshLineNodeBoundary<Dimension>&& mesh_line_node_boundary)
    : m_mesh_line_node_boundary(mesh_line_node_boundary)
  {
    ;
  }

  ~AxisBoundaryCondition() = default;
};

template <>
class MeshSmootherHandler::MeshSmoother<1>::AxisBoundaryCondition
{
 public:
  AxisBoundaryCondition()  = default;
  ~AxisBoundaryCondition() = default;
};

template <size_t Dimension>
class MeshSmootherHandler::MeshSmoother<Dimension>::FixedBoundaryCondition
{
 private:
  const MeshNodeBoundary<Dimension> m_mesh_node_boundary;

 public:
  const Array<const NodeId>&
  nodeList() const
  {
    return m_mesh_node_boundary.nodeList();
  }

  FixedBoundaryCondition(MeshNodeBoundary<Dimension>&& mesh_node_boundary) : m_mesh_node_boundary{mesh_node_boundary} {}

  ~FixedBoundaryCondition() = default;
};

template <size_t Dimension>
class MeshSmootherHandler::MeshSmoother<Dimension>::SymmetryBoundaryCondition
{
 public:
  using Rd = TinyVector<Dimension, double>;

 private:
  const MeshFlatNodeBoundary<Dimension> m_mesh_flat_node_boundary;

 public:
  const Rd&
  outgoingNormal() const
  {
    return m_mesh_flat_node_boundary.outgoingNormal();
  }

  const Array<const NodeId>&
  nodeList() const
  {
    return m_mesh_flat_node_boundary.nodeList();
  }

  SymmetryBoundaryCondition(MeshFlatNodeBoundary<Dimension>&& mesh_flat_node_boundary)
    : m_mesh_flat_node_boundary(mesh_flat_node_boundary)
  {
    ;
  }

  ~SymmetryBoundaryCondition() = default;
};

std::shared_ptr<const IMesh>
MeshSmootherHandler::getSmoothedMesh(
  const IMesh& mesh,
  const std::vector<std::shared_ptr<const IBoundaryConditionDescriptor>>& bc_descriptor_list) const
{
  switch (mesh.dimension()) {
  case 1: {
    constexpr size_t Dimension = 1;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh();
  }
  case 2: {
    constexpr size_t Dimension = 2;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh();
  }
  case 3: {
    constexpr size_t Dimension = 3;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh();
  }
  default: {
    throw UnexpectedError("invalid mesh dimension");
  }
  }
}

std::shared_ptr<const IMesh>
MeshSmootherHandler::getSmoothedMesh(
  const IMesh& mesh,
  const std::vector<std::shared_ptr<const IBoundaryConditionDescriptor>>& bc_descriptor_list,
  const FunctionSymbolId& function_symbol_id) const
{
  switch (mesh.dimension()) {
  case 1: {
    constexpr size_t Dimension = 1;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh(function_symbol_id);
  }
  case 2: {
    constexpr size_t Dimension = 2;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh(function_symbol_id);
  }
  case 3: {
    constexpr size_t Dimension = 3;
    using MeshType             = Mesh<Connectivity<Dimension>>;
    MeshSmoother smoother(dynamic_cast<const MeshType&>(mesh), bc_descriptor_list);
    return smoother.getSmoothedMesh(function_symbol_id);
  }
  default: {
    throw UnexpectedError("invalid mesh dimension");
  }
  }
}