// SPDX-FileCopyrightText: 2021-2024 The Ikarus Developers mueller@ibb.uni-stuttgart.de
// SPDX-License-Identifier: LGPL-3.0-or-later

#pragma once

#include "registerelement.hh"

#include <dune/fufem/boundarypatch.hh>
#include <dune/functions/functionspacebases/lagrangebasis.hh>
#include <dune/functions/functionspacebases/powerbasis.hh>
#include <dune/grid/yaspgrid.hh>
#include <dune/python/common/typeregistry.hh>
#include <dune/python/functions/globalbasis.hh>
#include <dune/python/pybind11/eigen.h>
#include <dune/python/pybind11/functional.h>
#include <dune/python/pybind11/pybind11.h>
#include <dune/python/pybind11/stl.h>

#include <ikarus/finiteelements/ferequirements.hh>
#include <ikarus/finiteelements/mechanics/nonlinearelastic.hh>
#include <ikarus/utils/basis.hh>

namespace Ikarus::Python {

template <class NonLinearElastic, class... options>
void registerNonLinearElastic(pybind11::handle scope, pybind11::class_<NonLinearElastic, options...> cls) {
  using pybind11::operator""_a;

  using GlobalBasis    = typename NonLinearElastic::BasisHandler;
  using FlatBasis      = typename NonLinearElastic::FlatBasis;
  using GridView       = typename GlobalBasis::GridView;
  using Element        = typename NonLinearElastic::Element;
  using Traits         = typename NonLinearElastic::Traits;
  using FERequirements = typename NonLinearElastic::FERequirementType;
  using Material       = typename NonLinearElastic::Material;
  using FERequirements = typename NonLinearElastic::FERequirementType;

  cls.def(pybind11::init([](const GlobalBasis& basis, const Element& element, const Material& mat) {
            return new NonLinearElastic(basis, element, mat);
          }),
          pybind11::keep_alive<1, 2>(), pybind11::keep_alive<1, 3>());

  using LoadFunction = std::function<Eigen::Vector<double, Traits::worlddim>(
      Dune::FieldVector<double, Traits::worlddim>, const double&)>;

  cls.def(pybind11::init(
              [](const GlobalBasis& basis, const Element& element, const Material& mat, const LoadFunction volumeLoad) {
                return new NonLinearElastic(basis, element, mat, volumeLoad);
              }),
          pybind11::keep_alive<1, 2>(), pybind11::keep_alive<1, 3>());

  cls.def(pybind11::init([](const GlobalBasis& basis, const Element& element, const Material& mat,
                            const LoadFunction volumeLoad, const BoundaryPatch<GridView>& bp,
                            const LoadFunction neumannBoundaryLoad) {
            return new NonLinearElastic(basis, element, mat, volumeLoad, &bp, neumannBoundaryLoad);
          }),
          pybind11::keep_alive<1, 2>(), pybind11::keep_alive<1, 3>(), pybind11::keep_alive<1, 7>());

  pybind11::module scopedf = pybind11::module::import("dune.functions");

  typedef Dune::Python::LocalViewWrapper<FlatBasis> LocalViewWrapper;
  auto includes = Dune::Python::IncludeFiles{"dune/python/functions/globalbasis.hh"};
  auto lv       = Dune::Python::insertClass<LocalViewWrapper>(
                scopedf, "LocalViewWrapper",
                Dune::Python::GenerateTypeName("Dune::Python::LocalViewWrapperWrapper", Dune::MetaType<FlatBasis>()),
                includes)
                .first;
  lv.def("bind", &LocalViewWrapper::bind);
  lv.def("unbind", &LocalViewWrapper::unbind);
  lv.def("index", [](const LocalViewWrapper& localView, int index) { return localView.index(index); });
  lv.def("__len__", [](LocalViewWrapper& self) -> int { return self.size(); });

  Dune::Python::Functions::registerTree<typename LocalViewWrapper::Tree>(lv);
  lv.def("tree", [](const LocalViewWrapper& view) { return view.tree(); });

  auto basisName = Dune::className<FlatBasis>();
  auto entry     = Dune::Python::insertClass<FlatBasis>(scopedf, "DefaultGlobalBasis", pybind11::buffer_protocol(),
                                                    Dune::Python::GenerateTypeName(basisName),
                                                    Dune::Python::IncludeFiles{"dune/python/functions/globalbasis.hh"});
  if (entry.second)
    Dune::Python::registerGlobalBasis(scopedf, entry.first);

  cls.def(
      "localView",
      [](NonLinearElastic& self) -> LocalViewWrapper {
        auto lvWrapped = LocalViewWrapper(self.localView().globalBasis());
        // this can be simplified when https://gitlab.dune-project.org/staging/dune-functions/-/merge_requests/418
        // becomes available
        pybind11::object obj = pybind11::cast(self.localView().element());
        lvWrapped.bind(obj);
        return lvWrapped;
      },
      pybind11::keep_alive<0, 1>());
  cls.def("calculateScalar",
          [](NonLinearElastic& self, const FERequirements& req) { return self.calculateScalar(req); });
  cls.def("calculateVector", [](NonLinearElastic& self, const FERequirements& req, Eigen::Ref<Eigen::VectorXd> vec) {
    return self.calculateVector(req, vec);
  });
  cls.def(
      "calculateMatrix",
      [](NonLinearElastic& self, const FERequirements& req, Eigen::Ref<Eigen::MatrixXd> mat) {
        return self.calculateMatrix(req, mat);
      },
      pybind11::arg("FERequirements"), pybind11::arg("elementMatrix").noconvert());

  auto resultTypes = Dune::makeTupleVector(makeRT<ResultType::PK2Stress>());
  registerCalculateAt(scope, cls, resultTypes);
}

} // namespace Ikarus::Python
