// This file is part of GOAST, a C++ library for variational methods in Geometry Processing
//
// Copyright (C) 2021 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 <array>

#include "shell-energy/bending_energy.h"
#include "shell-energy/auxiliary.h"

void shell::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> &Hess ) {
  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 16 * 9 * uE.rows());

  auto numV = V_def.rows();

  auto localToGlobal = [&]( int k, int l, const Eigen::Matrix3d &localMatrix ) {
    for ( int i = 0; i < 3; i++ )
      for ( int j = 0; j < 3; j++ )
        tripletList.emplace_back( i * numV + k, j * numV + l,
                                  3. * localMatrix( i, j )); // Factor 3 was omitted in original code

    if ( k != l )
      for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
          tripletList.emplace_back( i * numV + l, j * numV + k,
                                    3. * localMatrix( j, i )); // Factor 3 was omitted in original code
  };

  for ( int edgeIdx = 0; edgeIdx < uE.rows(); ++edgeIdx ) {
    // Check for boundary edge
    if ( EF( edgeIdx, 0 ) == -1 || EF( edgeIdx, 1 ) == -1 )
      continue;

    int pi = uE( edgeIdx, 0 );
    int pj = uE( edgeIdx, 1 );
    int pk = F( EF( edgeIdx, 0 ), EI( edgeIdx, 0 ));
    int pl = F( EF( edgeIdx, 1 ), EI( edgeIdx, 1 ));

    // set up vertices and edges
    Eigen::Vector3d Pi, Pj, Pk, Pl, temp;

    // first get undefomed quantities
    Pi = V_undef.row( pi );
    Pj = V_undef.row( pj );
    Pk = V_undef.row( pk );
    Pl = V_undef.row( pl );
    double delTheta = getDihedralAngle( Pi, Pj, Pk, Pl );


    // compute Ak + Al
    temp = ( Pk - Pj ).cross( Pi - Pk );
    double vol = temp.norm() / 2.;
    temp = ( Pl - Pi ).cross( Pj - Pl );
    vol += temp.norm() / 2.;
    double elengthSqr = ( Pj - Pi ).dot( Pj - Pi );

    // now get the deformed values
    Pi = V_def.row( pi );
    Pj = V_def.row( pj );
    Pk = V_def.row( pk );
    Pl = V_def.row( pl );

    // compute weighted differnce of dihedral angles
    delTheta -= getDihedralAngle( Pi, Pj, Pk, Pl );
    delTheta *= -2. * elengthSqr / vol;
    double factor = 2. * elengthSqr / vol;

    // compute first derivatives of dihedral angle
    Eigen::Vector3d thetak, thetal, thetai, thetaj;
    getThetaGradK( Pi, Pj, Pk, thetak );
    getThetaGradK( Pj, Pi, Pl, thetal );
    getThetaGradI( Pi, Pj, Pk, Pl, thetai );
    getThetaGradJ( Pi, Pj, Pk, Pl, thetaj );

    // now compute second derivatives of dihedral angle
    Eigen::Matrix3d tensorProduct, H, aux;

    //kk
    getHessThetaKK( Pi, Pj, Pk, aux );
    tensorProduct = thetak * thetak.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pk, pk, H );

    //ik & ki (Hki = Hik)
    getHessThetaIK( Pi, Pj, Pk, aux );
    tensorProduct = thetai * thetak.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pi, pk, H );

    //jk & kj (Hkj = Hjk)
    getHessThetaJK( Pi, Pj, Pk, aux );
    tensorProduct = thetaj * thetak.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pj, pk, H );

    //ll
    getHessThetaKK( Pj, Pi, Pl, aux );
    tensorProduct = thetal * thetal.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pl, pl, H );

    //il & li (Hli = Hil)
    getHessThetaJK( Pj, Pi, Pl, aux );
    tensorProduct = thetai * thetal.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pi, pl, H );

    //jl & lj (Hlj = Hjl)
    getHessThetaIK( Pj, Pi, Pl, aux );
    tensorProduct = thetaj * thetal.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pj, pl, H );

    //kl/lk: Hkl = 0 and Hlk = 0
    tensorProduct = thetak * thetal.transpose();
    tensorProduct *= factor;
    localToGlobal( pk, pl, tensorProduct );

    //ii
    getHessThetaII( Pi, Pj, Pk, Pl, aux );
    tensorProduct = thetai * thetai.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pi, pi, H );

    //jj
    getHessThetaII( Pj, Pi, Pl, Pk, aux );
    tensorProduct = thetaj * thetaj.transpose();
    H = factor * tensorProduct + delTheta * aux;
    localToGlobal( pj, pj, H );

    //ij & ji (Hij = Hji)
    getHessThetaJI( Pi, Pj, Pk, Pl, H );
    H *= delTheta;
    tensorProduct = thetai * thetaj.transpose();
    H += factor * tensorProduct;
    localToGlobal( pi, pj, H );
  }

  // fill matrix from triplets
  Hess.setFromTriplets( tripletList.begin(), tripletList.end());
}

void shell::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> &Hess ) {
  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 16 * 9 * uE.rows());

  auto numV = V_def.rows();

  auto localToGlobal = [&]( int k, int l, const Eigen::Matrix3d &localMatrix ) {
    for ( int i = 0; i < 3; i++ )
      for ( int j = 0; j < 3; j++ )
        tripletList.emplace_back( i * numV + k, j * numV + l,
                                  3. * localMatrix( i, j )); // Factor 3 was omitted in original code

    if ( k != l )
      for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
          tripletList.emplace_back( i * numV + l, j * numV + k,
                                    3. * localMatrix( j, i )); // Factor 3 was omitted in original code
  };

  for ( int edgeIdx = 0; edgeIdx < uE.rows(); ++edgeIdx ) {
    // Check for boundary edge
    if ( EF( edgeIdx, 0 ) == -1 || EF( edgeIdx, 1 ) == -1 )
      continue;

    int pi = uE( edgeIdx, 0 );
    int pj = uE( edgeIdx, 1 );
    int pk = F( EF( edgeIdx, 0 ), EI( edgeIdx, 0 ));
    int pl = F( EF( edgeIdx, 1 ), EI( edgeIdx, 1 ));

    // set up vertices and edges
    Eigen::Vector3d Pi, Pj, Pk, Pl, temp;

    // first get deformed quantities
    Pi = V_def.row( pi );
    Pj = V_def.row( pj );
    Pk = V_def.row( pk );
    Pl = V_def.row( pl );
    double delTheta = getDihedralAngle( Pi, Pj, Pk, Pl );

    // now get the undeformed values
    Pi = V_undef.row( pi );
    Pj = V_undef.row( pj );
    Pk = V_undef.row( pk );
    Pl = V_undef.row( pl );

    // compute Ak + Al, |e|^2 and diff. o dihedral angles
    temp = ( Pk - Pj ).cross( Pi - Pk );
    double vol = temp.norm() / 2.;
    temp = ( Pl - Pi ).cross( Pj - Pl );
    vol += temp.norm() / 2.;
    double elengthSqr = ( Pj - Pi ).dot( Pj - Pi );
    // note sign here!
    delTheta -= getDihedralAngle( Pi, Pj, Pk, Pl );
    delTheta *= -1.;

    // compute first derivatives of dihedral angle
    Eigen::Vector3d thetak, thetal, thetai, thetaj;
    getThetaGradK( Pi, Pj, Pk, thetak );
    getThetaGradK( Pj, Pi, Pl, thetal );
    getThetaGradI( Pi, Pj, Pk, Pl, thetai );
    getThetaGradJ( Pi, Pj, Pk, Pl, thetaj );

    // compute first derivatives of area
    Eigen::Vector3d areak, areal, areai, areaj;
    getAreaGradK( Pi, Pj, Pk, areak );
    getAreaGradK( Pj, Pi, Pl, areal );
    getAreaGradK( Pj, Pk, Pi, areai );
    getAreaGradK( Pl, Pj, Pi, temp );
    areai += temp;
    getAreaGradK( Pk, Pi, Pj, areaj );
    getAreaGradK( Pi, Pl, Pj, temp );
    areaj += temp;

    // now compute second derivatives
    Eigen::Matrix3d H, auxMat;
    Eigen::Vector3d auxVec, e( Pj - Pi );

    //*k
    auxVec = elengthSqr * thetak - delTheta * elengthSqr / vol * areak;
    temp = 2. * thetak - delTheta / vol * areak;

    //kk      
    H = thetak * auxVec.transpose();
    getHessThetaKK( Pi, Pj, Pk, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areak * auxVec.transpose();

    getHessAreaKK( Pi, Pj, Pk, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H *= 2. / vol;
    localToGlobal( pk, pk, H );

    //lk      
    H = thetal * auxVec.transpose();
    H -= delTheta / vol * areal * auxVec.transpose();

    H *= 2. / vol;
    localToGlobal( pl, pk, H );

    //ik
    H = thetai * auxVec.transpose();
    getHessThetaIK( Pi, Pj, Pk, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areai * auxVec.transpose();

    getHessAreaIK( Pi, Pj, Pk, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;


    H -= delTheta * e * temp.transpose();

    H *= 2. / vol;
    localToGlobal( pi, pk, H );

    //jk
    H = thetaj * auxVec.transpose();
    getHessThetaJK( Pi, Pj, Pk, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areaj * auxVec.transpose();
    getHessAreaIK( Pj, Pi, Pk, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H += delTheta * e * temp.transpose();
    H *= 2. / vol;
    localToGlobal( pj, pk, H );

    //*l
    auxVec = elengthSqr * thetal - delTheta * elengthSqr / vol * areal;
    temp = 2. * thetal - delTheta / vol * areal;

    //ll      
    H = thetal * auxVec.transpose();
    getHessThetaKK( Pj, Pi, Pl, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areal * auxVec.transpose();
    getHessAreaKK( Pj, Pi, Pl, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H *= 2. / vol;
    localToGlobal( pl, pl, H );

    //il
    H = thetai * auxVec.transpose();
    getHessThetaJK( Pj, Pi, Pl, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areai * auxVec.transpose();
    getHessAreaIK( Pi, Pj, Pl, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H -= delTheta * e * temp.transpose();
    H *= 2. / vol;
    localToGlobal( pi, pl, H );

    //jl
    H = thetaj * auxVec.transpose();
    getHessThetaIK( Pj, Pi, Pl, auxMat );
    H += delTheta * elengthSqr * auxMat;
    H -= delTheta / vol * areaj * auxVec.transpose();
    getHessAreaIK( Pj, Pi, Pl, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H += delTheta * e * temp.transpose();
    H *= 2. / vol;
    localToGlobal( pj, pl, H );

    //*j
    auxVec = elengthSqr * thetaj - delTheta * elengthSqr / vol * areaj;
    temp = 2. * thetaj - delTheta / vol * areaj;
    auxVec += 2. * delTheta * e; // caution with factor 2!!!!!!!

    //jj     
    H = thetaj * auxVec.transpose();
    getHessThetaII( Pj, Pi, Pl, Pk, auxMat );
    H += delTheta * elengthSqr * auxMat;
    auxVec -= delTheta * e;
    H -= delTheta / vol * areaj * auxVec.transpose();
    getHessAreaKK( Pk, Pi, Pj, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    getHessAreaKK( Pi, Pl, Pj, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H += delTheta * e * temp.transpose();
    H.diagonal().array() += delTheta * delTheta;
    H *= 2. / vol;
    localToGlobal( pj, pj, H );

    //ij     
    auxVec += delTheta * e;
    H = thetai * auxVec.transpose();
    getHessThetaJI( Pi, Pj, Pk, Pl, auxMat );
    H += delTheta * elengthSqr * auxMat;
    auxVec -= delTheta * e;
    H -= delTheta / vol * areai * auxVec.transpose();
    getHessAreaIK( Pi, Pk, Pj, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    getHessAreaIK( Pi, Pl, Pj, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H -= delTheta * e * temp.transpose();
    H.diagonal().array() -= delTheta * delTheta;
    H *= 2. / vol;
    localToGlobal( pi, pj, H );

    //*i
    auxVec = elengthSqr * thetai - delTheta * elengthSqr / vol * areai;
    temp = 2. * thetai - delTheta / vol * areai;

    auxVec -= 2. * delTheta * e; // caution with factor 2!!!!!!!

    //ii     
    H = thetai * auxVec.transpose();
    getHessThetaII( Pi, Pj, Pk, Pl, auxMat );
    H += delTheta * elengthSqr * auxMat;
    auxVec += delTheta * e;
    H -= delTheta / vol * areai * auxVec.transpose();
    getHessAreaKK( Pl, Pj, Pi, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    getHessAreaKK( Pj, Pk, Pi, auxMat );
    H -= 0.5 * delTheta * delTheta * elengthSqr / vol * auxMat;
    H -= delTheta * e * temp.transpose();
    H.diagonal().array() += delTheta * delTheta;
    H *= 2. / vol;
    localToGlobal( pi, pi, H );
  }

  // fill matrix from triplets
  Hess.setFromTriplets( tripletList.begin(), tripletList.end());
}


void shell::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,
                                   Eigen::SparseMatrix<double> &Hess, bool FirstDerivativeWRTDef ) {
  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 16 * 9 * uE.rows());

  auto numV = V_def.rows();

  auto localToGlobal = [&]( int k, int l, const Eigen::Matrix3d &localMatrix ) {
    if ( FirstDerivativeWRTDef ) {
      for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
          tripletList.emplace_back( i * numV + k, j * numV + l,
                                    3 * localMatrix( i, j )); // Factor 3 was omitted in original code
    }
    else {
      for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
          tripletList.emplace_back( i * numV + l, j * numV + k,
                                    3 * localMatrix( j, i )); // Factor 3 was omitted in original code
    }
  };

  for ( int edgeIdx = 0; edgeIdx < uE.rows(); ++edgeIdx ) {
    // Check for boundary edge
    if ( EF( edgeIdx, 0 ) == -1 || EF( edgeIdx, 1 ) == -1 )
      continue;

    int pi = uE( edgeIdx, 0 );
    int pj = uE( edgeIdx, 1 );
    int pk = F( EF( edgeIdx, 0 ), EI( edgeIdx, 0 ));
    int pl = F( EF( edgeIdx, 1 ), EI( edgeIdx, 1 ));

    std::array<int, 4> idx{ pi, pj, pk, pl };

    // set up vertices and edges
    Eigen::Vector3d Pi, Pj, Pk, Pl, gradArea, gradTheta, temp;
    std::array<Eigen::Vector3d, 4> undefGrads, defGrads;

    // first get deformed quantities
    Pi = V_def.row( pi );
    Pj = V_def.row( pj );
    Pk = V_def.row( pk );
    Pl = V_def.row( pl );
    double delTheta = getDihedralAngle( Pi, Pj, Pk, Pl );

    getThetaGradI( Pi, Pj, Pk, Pl, defGrads[0] );
    getThetaGradJ( Pi, Pj, Pk, Pl, defGrads[1] );
    getThetaGradK( Pi, Pj, Pk, defGrads[2] );
    getThetaGradK( Pj, Pi, Pl, defGrads[3] );

    // now get the undeformed values
    Pi = V_undef.row( pi );
    Pj = V_undef.row( pj );
    Pk = V_undef.row( pk );
    Pl = V_undef.row( pl );

    // compute Ak + Al, |e|^2 and diff. o dihedral angles
    temp = ( Pk - Pj ).cross( Pi - Pk );
    double vol = temp.norm() / 2.;
    temp = ( Pl - Pi ).cross( Pj - Pl );
    vol += temp.norm() / 2.;
    double elengthSqr = ( Pj - Pi ).dot( Pj - Pi );
    // note sign here!
    delTheta -= getDihedralAngle( Pi, Pj, Pk, Pl );

    // factors
    double factorGradTheta = -2. * elengthSqr / vol;
    double factorGradArea = -2. * delTheta * elengthSqr / ( vol * vol );
    double factorGradEdgeLengthSqr = 4. * delTheta / vol;

    // d_i
    getThetaGradI( Pi, Pj, Pk, Pl, gradTheta );
    getAreaGradK( Pj, Pk, Pi, gradArea );
    undefGrads[0] = factorGradTheta * gradTheta + factorGradArea * gradArea;
    getAreaGradK( Pl, Pj, Pi, gradArea );
    undefGrads[0] += factorGradArea * gradArea;
    undefGrads[0] += factorGradEdgeLengthSqr * ( Pi - Pj );

    // d_j
    getThetaGradJ( Pi, Pj, Pk, Pl, gradTheta );
    getAreaGradK( Pk, Pi, Pj, gradArea );
    undefGrads[1] = factorGradTheta * gradTheta + factorGradArea * gradArea;
    getAreaGradK( Pi, Pl, Pj, gradArea );
    undefGrads[1] += factorGradArea * gradArea, factorGradArea;
    undefGrads[1] += factorGradEdgeLengthSqr * ( Pj - Pi );

    // d_k
    getThetaGradK( Pi, Pj, Pk, gradTheta );
    getAreaGradK( Pi, Pj, Pk, gradArea );
    undefGrads[2] = factorGradTheta * gradTheta + factorGradArea * gradArea;

    // d_l
    getThetaGradK( Pj, Pi, Pl, gradTheta );
    getAreaGradK( Pj, Pi, Pl, gradArea );
    undefGrads[3] = factorGradTheta * gradTheta + factorGradArea * gradArea;

    // local to global
    for ( int m = 0; m < 4; m++ ) {
      for ( int n = 0; n < 4; n++ ) {
        Eigen::Matrix3d matrix = defGrads[m] * undefGrads[n].transpose();
        localToGlobal( idx[m], idx[n], matrix );
      }
    }
  }

  // fill matrix from triplets
  Hess.setFromTriplets( tripletList.begin(), tripletList.end());
}