// Copyright (c) 2020 Chris Richardson
// FEniCS Project
// SPDX-License-Identifier:    MIT

#include "basix.h"
#include "core/cell.h"
#include "core/finite-element.h"
#include "core/mappings.h"
#include "core/quadrature.h"
#include "core/span.hpp"
#include <algorithm>
#include <iterator>
#include <memory>
#include <vector>

using namespace basix;

std::vector<std::unique_ptr<FiniteElement>> _registry;

void check_handle(int handle)
{
  if (handle < 0 or handle >= (int)_registry.size())
    throw std::range_error("Bad index");
  if (!_registry[handle])
    throw std::runtime_error("Invalid element");
}
namespace

{
template <typename T>
void _map_push_forward(int handle, T* physical_data, const T* reference_data,
                       const double* J, const double* detJ, const double* K,
                       const int physical_dim, const int physical_value_size,
                       const int nresults, const int npoints)
{
  check_handle(handle);
  const int tdim = cell::topological_dimension(_registry[handle]->cell_type());
  const int vs = _registry[handle]->value_size();
  _registry[handle]->map_push_forward_m<T>(
      Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::RowMajor>>(reference_data,
                                                      npoints * nresults, vs),
      Eigen::Map<const Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::RowMajor>>(J, npoints,
                                                      physical_dim * tdim),
      tcb::span(detJ, npoints),
      Eigen::Map<const Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::RowMajor>>(K, npoints,
                                                      physical_dim * tdim),
      physical_data);
}

template <typename T>
void _map_pull_back(int handle, T* reference_data, const T* physical_data,
                    const double* J, const double* detJ, const double* K,
                    const int physical_dim, const int physical_value_size,
                    const int nresults, const int npoints)
{
  check_handle(handle);
  const int tdim = cell::topological_dimension(_registry[handle]->cell_type());
  _registry[handle]->map_pull_back_m<T>(
      Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::ColMajor>>(
          physical_data, npoints, physical_value_size * nresults),
      Eigen::Map<const Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::RowMajor>>(J, npoints,
                                                      physical_dim * tdim),
      tcb::span(detJ, npoints),
      Eigen::Map<const Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::RowMajor>>(K, npoints,
                                                      physical_dim * tdim),
      reference_data);
}
} // namespace

int basix::register_element(const char* family_name, const char* cell_type,
                            int degree)
{
  _registry.push_back(std::make_unique<FiniteElement>(
      create_element(family_name, cell_type, degree)));
  return _registry.size() - 1;
}

void basix::release_element(int handle)
{
  check_handle(handle);
  _registry[handle].reset();
  while (!_registry.empty() and !_registry.back())
    _registry.pop_back();
}

void basix::tabulate(int handle, double* basis_values, int nd, const double* x,
                     int npoints)
{
  check_handle(handle);

  // gdim and tdim are the same for all cells in basix
  const int gdim = cell::topological_dimension(_registry[handle]->cell_type());
  Eigen::Map<const Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic,
                                Eigen::RowMajor>>
      _x(x, npoints, gdim);

  _registry[handle]->tabulate(nd, _x, basis_values);

  // std::vector<Eigen::ArrayXXd> values = _registry[handle]->tabulate(nd, _x);

  // const int m = values[0].rows() * values[0].cols();
  // for (std::size_t i = 0; i < values.size(); ++i)
  //   std::copy(values[i].data(), values[i].data() + m, basis_values + i * m);
}

void basix::map_push_forward_real(int handle, double* physical_data,
                                  const double* reference_data, const double* J,
                                  const double* detJ, const double* K,
                                  const int physical_dim,
                                  const int physical_value_size,
                                  const int nresults, const int npoints)
{
  _map_push_forward<double>(handle, physical_data, reference_data, J, detJ, K,
                            physical_dim, physical_value_size, nresults,
                            npoints);
}

void basix::map_pull_back_real(int handle, double* reference_data,
                               const double* physical_data, const double* J,
                               const double* detJ, const double* K,
                               const int physical_dim,
                               const int physical_value_size,
                               const int nresults, const int npoints)
{
  _map_pull_back<double>(handle, reference_data, physical_data, J, detJ, K,
                         physical_dim, physical_value_size, nresults, npoints);
}

void basix::map_push_forward_complex(int handle,
                                     std::complex<double>* physical_data,
                                     const std::complex<double>* reference_data,
                                     const double* J, const double* detJ,
                                     const double* K, const int physical_dim,
                                     const int physical_value_size,
                                     const int nresults, const int npoints)
{
  _map_push_forward<std::complex<double>>(
      handle, physical_data, reference_data, J, detJ, K, physical_dim,
      physical_value_size, nresults, npoints);
}

void basix::map_pull_back_complex(int handle,
                                  std::complex<double>* reference_data,
                                  const std::complex<double>* physical_data,
                                  const double* J, const double* detJ,
                                  const double* K, const int physical_dim,
                                  const int physical_value_size,
                                  const int nresults, const int npoints)
{
  _map_pull_back<std::complex<double>>(handle, reference_data, physical_data, J,
                                       detJ, K, physical_dim,
                                       physical_value_size, nresults, npoints);
}

const char* basix::cell_type(int handle)
{
  check_handle(handle);
  return cell::type_to_str(_registry[handle]->cell_type()).c_str();
}

int basix::degree(int handle)
{
  check_handle(handle);
  return _registry[handle]->degree();
}

int basix::dim(int handle)
{
  check_handle(handle);
  return _registry[handle]->dim();
}

int basix::value_rank(int handle)
{
  check_handle(handle);
  return _registry[handle]->value_shape().size();
}

void basix::value_shape(int handle, int* dimensions)
{
  check_handle(handle);
  std::vector<int> dims = _registry[handle]->value_shape();
  std::copy(dims.begin(), dims.end(), dimensions);
}

int basix::interpolation_num_points(int handle)
{
  check_handle(handle);
  return _registry[handle]->num_points();
}

void basix::interpolation_points(int handle, double* points)
{
  check_handle(handle);
  Eigen::Map<
      Eigen::Array<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
      points, interpolation_num_points(handle),
      cell_geometry_dimension(cell_type(handle)))
      = _registry[handle]->points();
}

void basix::interpolation_matrix(int handle, double* matrix)
{
  check_handle(handle);
  Eigen::Map<
      Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
      matrix, dim(handle),
      interpolation_num_points(handle) * _registry[handle]->value_size())
      = _registry[handle]->interpolation_matrix();
}

void basix::entity_dofs(int handle, int dim, int* num_dofs)
{
  check_handle(handle);
  std::vector<std::vector<int>> dof_counts = _registry[handle]->entity_dofs();
  std::copy(dof_counts[dim].begin(), dof_counts[dim].end(), num_dofs);
}

const char* basix::family_name(int handle)
{
  check_handle(handle);
  return element::type_to_str(_registry[handle]->family()).c_str();
}

const char* basix::mapping_name(int handle)
{
  check_handle(handle);
  return mapping::type_to_str(_registry[handle]->mapping_type()).c_str();
}

int basix::cell_geometry_num_points(const char* cell_type)
{
  cell::type ct = cell::str_to_type(cell_type);
  return cell::geometry(ct).shape[0];
}

int basix::cell_geometry_dimension(const char* cell_type)
{
  cell::type ct = cell::str_to_type(cell_type);
  return cell::geometry(ct).shape[1];
}

void basix::cell_geometry(const char* cell_type, double* points)
{
  cell::type ct = cell::str_to_type(cell_type);
  ndarray<double, 2> pts = cell::geometry(ct);
  std::copy(pts.data(), pts.data() + pts.size(), points);
}

std::vector<std::vector<std::vector<int>>>
basix::topology(const char* cell_type)
{
  cell::type ct = cell::str_to_type(cell_type);
  return cell::topology(ct);
}
