// This file is part of GOAST, a C++ library for variational methods in Geometry Processing
//
// Copyright (C) 2023 Behrend Heeren & Josua Sassen, University of Bonn <goast@ins.uni-bonn.de>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#include <nanobind/eigen/dense.h>
#include <nanobind/eigen/sparse.h>
#include <shell-energy/shell_energy.h>

namespace nb = nanobind;

using namespace nb::literals;

NB_MODULE( pyshell, m ) {
  // Energies
  m.def( "bending_energy", &shell::bending_energy, "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a );
  m.def( "membrane_energy", &shell::membrane_energy, "v_undef"_a, "v_def"_a, "f"_a, "mu"_a = 1., "lambda"_a = 1. );
  m.def( "shell_energy", &shell::shell_energy, "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a,
         "bending_weight"_a, "mu"_a = 1., "lambda"_a = 1. );

  // Gradients
  m.def( "bending_deformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
             const Eigen::MatrixXi &F, const Eigen::MatrixXi &uE,
             const Eigen::VectorXi &EMAP,
             const Eigen::MatrixXi &EF, const Eigen::MatrixXi &EI ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::bending_deformed_gradient( V_undef, V_def, F, uE, EMAP, EF, EI, out );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a
  );

  m.def( "bending_undeformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
             const Eigen::MatrixXi &F, const Eigen::MatrixXi &uE,
             const Eigen::VectorXi &EMAP,
             const Eigen::MatrixXi &EF, const Eigen::MatrixXi &EI ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::bending_undeformed_gradient( V_undef, V_def, F, uE, EMAP, EF, EI, out );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a
  );

  m.def( "membrane_deformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
             const Eigen::MatrixXi &F, double mu, double lambda ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::membrane_deformed_gradient( V_undef, V_def, F, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "membrane_undeformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
             const Eigen::MatrixXi &F, double mu, double lambda ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::membrane_undeformed_gradient( V_undef, V_def, F, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "shell_deformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, double bendingWeight, double mu,
             double lambda ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::shell_deformed_gradient( V_undef, V_def, F, uE, EMAP, EF, EI, bendingWeight, out, mu, lambda );
           return out;
         }, "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a,
         "bending_weight"_a, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "shell_undeformed_gradient",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, double bendingWeight, double mu,
             double lambda ) -> Eigen::MatrixXd {
           Eigen::MatrixXd out;
           shell::shell_undeformed_gradient( V_undef, V_def, F, uE, EMAP, EF, EI, bendingWeight, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a, "bending_weight"_a, "mu"_a = 1.,
         "lambda"_a = 1.
  );

  // Hessians

  m.def( "bending_deformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::bending_deformed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, out );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a
  );

  m.def( "bending_undeformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::bending_undeformed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, out );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a
  );

  m.def( "bending_mixed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, bool FirstDerivativeWRTDef ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::bending_mixed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, out, FirstDerivativeWRTDef );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a, "first_deriv_wrt_def"_a = true
  );

  m.def( "membrane_deformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::membrane_deformed_hessian( V_undef, V_def, F, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "membrane_undeformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::membrane_undeformed_hessian( V_undef, V_def, F, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "membrane_mixed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             bool FirstDerivativeWRTDef, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::membrane_mixed_hessian( V_undef, V_def, F, out, FirstDerivativeWRTDef,
                                          mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "first_deriv_wrt_def"_a = true, "mu"_a = 1., "lambda"_a = 1.
  );

  m.def( "shell_deformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, double bendingWeight, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::shell_deformed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, bendingWeight, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a, "bending_weight"_a, "mu"_a = 1.,
         "lambda"_a = 1.
  );

  m.def( "shell_undeformed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, double bendingWeight, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::shell_undeformed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, bendingWeight, out, mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a, "bending_weight"_a, "mu"_a = 1.,
         "lambda"_a = 1.
  );

  m.def( "shell_mixed_hessian",
         []( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def, const Eigen::MatrixXi &F,
             const Eigen::MatrixXi &uE, const Eigen::VectorXi &EMAP, const Eigen::MatrixXi &EF,
             const Eigen::MatrixXi &EI, double bendingWeight, bool FirstDerivativeWRTDef, double mu,
             double lambda ) -> Eigen::SparseMatrix<double> {
           Eigen::SparseMatrix<double> out;
           shell::shell_mixed_hessian( V_undef, V_def, F, uE, EMAP, EF, EI, bendingWeight, out, FirstDerivativeWRTDef,
                                       mu, lambda );
           return out;
         },
         "v_undef"_a, "v_def"_a, "f"_a, "E"_a, "EMAP"_a, "EF"_a, "EI"_a, "bending_weight"_a,
         "first_deriv_wrt_def"_a = true, "mu"_a = 1., "lambda"_a = 1.
  );
}