// 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/membrane_energy.h"
#include "shell-energy/auxiliary.h"


void shell::membrane_deformed_hessian( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
                                       const Eigen::MatrixXi &F, Eigen::SparseMatrix<double> &Hess,
                                       double mu, double lambda ) {

  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 9 * 9 * F.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, localMatrix( i, j ));

    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, localMatrix( j, i ));
  };

  double lambdaQuarter = lambda / 4.;
  double muHalfPlusLambdaQuarter = mu / 2. + lambdaQuarter;

  for ( int faceIdx = 0; faceIdx < F.rows(); ++faceIdx ) {

    std::array<Eigen::Vector3d, 3> nodes, undefEdges, defEdges;
    Eigen::Vector3d temp;

    //! get undeformed quantities
    for ( int j  : { 0, 1, 2 } )
      nodes[j] = V_undef.row( F( faceIdx, j ));
    for ( int j  : { 0, 1, 2 } )
      undefEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];
    // compute volume
    temp = ( nodes[2] - nodes[1] ).cross( nodes[0] - nodes[2] );
    double volUndefSqr = temp.squaredNorm() / 4.;
    double volUndef = std::sqrt( volUndefSqr );

    //! get deformed quantities
    for ( int j  : { 0, 1, 2 } )
      nodes[j] = V_def.row( F( faceIdx, j ));
    for ( int j  : { 0, 1, 2 } )
      defEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];
    // compute volume
    temp = ( nodes[2] - nodes[1] ).cross( nodes[0] - nodes[2] );
    double volDefSqr = temp.squaredNorm() / 4.;
    double volDef = std::sqrt( volDefSqr );

    Eigen::Vector3d traceFactors;
    for ( int i = 0; i < 3; i++ )
      traceFactors[i] = -0.25 * mu * undefEdges[( i + 2 ) % 3].dot( undefEdges[( i + 1 ) % 3] ) / volUndef;

    double mixedFactor = 0.5 * lambda / volUndef + 2. * muHalfPlusLambdaQuarter * volUndef / volDefSqr;
    double areaFactor = 0.5 * lambda * volDef / volUndef - 2. * muHalfPlusLambdaQuarter * volUndef / volDef;

    // precompute area gradients
    std::array<Eigen::Vector3d, 3> gradArea;
    for ( int i = 0; i < 3; i++ )
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], gradArea[i] );

    // compute local matrices
    Eigen::Matrix3d tensorProduct, H, auxMat;

    // i==j
    for ( int i = 0; i < 3; i++ ) {
      getHessAreaKK( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], auxMat );
      tensorProduct = gradArea[i] * gradArea[i].transpose();
      H = areaFactor * auxMat + mixedFactor * tensorProduct;
      H.diagonal().array() += traceFactors[( i + 1 ) % 3] + traceFactors[( i + 2 ) % 3];
      localToGlobal( F( faceIdx, i ), F( faceIdx, i ), H );
    }

    // i!=j
    for ( int i = 0; i < 3; i++ ) {
      getHessAreaIK( nodes[i], nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], auxMat );
      tensorProduct = gradArea[i] * gradArea[( i + 2 ) % 3].transpose();
      H = areaFactor * auxMat + mixedFactor * tensorProduct;
      H.diagonal().array() -= traceFactors[( i + 1 ) % 3];
      localToGlobal( F( faceIdx, i ), F( faceIdx, ( i + 2 ) % 3 ), H );
    }

  }

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


void shell::membrane_undeformed_hessian( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
                                         const Eigen::MatrixXi &F, Eigen::SparseMatrix<double> &Hess,
                                         double mu, double lambda ) {

  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 9 * 9 * F.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, localMatrix( i, j ));

    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, localMatrix( j, i ));
  };


  double lambdaQuarter = lambda / 4.;
  double muHalfPlusLambdaQuarter = mu / 2. + lambdaQuarter;

  for ( int faceIdx = 0; faceIdx < F.rows(); ++faceIdx ) {
    std::array<Eigen::Vector3d, 3> nodes, undefEdges, fixedEdges;
    Eigen::Vector3d temp;

    //! get fixed edgess
    for ( int j : { 0, 1, 2 } )
      nodes[j] = V_def.row( F( faceIdx, j ));
    for ( int j : { 0, 1, 2 } )
      fixedEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];

    // compute volume
    temp = ( nodes[1] - nodes[0] ).cross( nodes[2] - nodes[1] );
    double volDefSqr = temp.squaredNorm() / 4.;
    double volDef = std::sqrt( volDefSqr );

    Eigen::Vector3d defLengthSqr;
    for ( int i : { 0, 1, 2 } )
      defLengthSqr[i] = fixedEdges[i].squaredNorm();

    //! get undeformed quantities
    for ( int j = 0; j < 3; j++ )
      nodes[j] = V_undef.row( F( faceIdx, j ));
    for ( int j = 0; j < 3; j++ )
      undefEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];

    // compute volume
    temp = ( nodes[1] - nodes[0] ).cross( nodes[2] - nodes[1] );
    double volUndefSqr = temp.squaredNorm() / 4.;
    double volUndef = std::sqrt( volUndefSqr );

    double traceTerm = 0.;
    for ( int i : { 0, 1, 2 } )
      traceTerm -= undefEdges[( i + 1 ) % 3].dot( undefEdges[( i + 2 ) % 3] ) * defLengthSqr[i];

    std::array<Eigen::Vector3d, 3> gradTrace;
    for ( int i = 0; i < 3; i++ )
      for ( int j = 0; j < 3; j++ )
        gradTrace[i][j] = defLengthSqr[i] * ( undefEdges[( i + 1 ) % 3][j] - undefEdges[( i + 2 ) % 3][j] ) +
                          undefEdges[i][j] * ( defLengthSqr[( i + 1 ) % 3] - defLengthSqr[( i + 2 ) % 3] );

    // precompute area gradients
    std::array<Eigen::Vector3d, 3> gradArea;
    for ( int i = 0; i < 3; i++ )
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], gradArea[i] );

    double areaFactor = 0.125 * mu * traceTerm + lambdaQuarter * volDefSqr;
    double negHessAreaFactor =
            muHalfPlusLambdaQuarter * ( std::log( volDefSqr / volUndefSqr ) - 2. ) + ( mu + lambdaQuarter ) +
            areaFactor / volUndefSqr;
    double mixedAreaFactor = 2 * ( areaFactor / volUndefSqr + muHalfPlusLambdaQuarter ) / volUndef;
    double mixedFactor = -0.125 * mu / volUndefSqr;
    double hessTraceFactor = 0.125 * mu / volUndef;


    // compute local matrices
    Eigen::Matrix3d tensorProduct, H, auxMat;

    // i==j
    for ( int i = 0; i < 3; i++ ) {
      getHessAreaKK( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], H );
      H *= -1. * negHessAreaFactor;

      H += mixedFactor * gradArea[i] * gradTrace[i].transpose();
      H += mixedFactor * gradTrace[i] * gradArea[i].transpose();
      H += mixedAreaFactor * gradArea[i] * gradArea[i].transpose();
      H.diagonal().array() += hessTraceFactor * 2 * defLengthSqr[i];

      localToGlobal( F( faceIdx, i ), F( faceIdx, i ), H );
    }

    // i!=j
    for ( int i = 0; i < 3; i++ ) {
      getHessAreaIK( nodes[i], nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], H );
      H *= -1. * negHessAreaFactor;

      H += mixedFactor * gradArea[i] * gradTrace[( i + 2 ) % 3].transpose();
      H += mixedFactor * gradTrace[i] * gradArea[( i + 2 ) % 3].transpose();
      H += mixedAreaFactor * gradArea[i] * gradArea[( i + 2 ) % 3].transpose();
      H.diagonal().array() +=
              hessTraceFactor * ( defLengthSqr[( i + 1 ) % 3] - defLengthSqr[i] - defLengthSqr[( i + 2 ) % 3] );

      localToGlobal( F( faceIdx, i ), F( faceIdx, ( i + 2 ) % 3 ), H );
    }

  }

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

void shell::membrane_mixed_hessian( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
                                    const Eigen::MatrixXi &F, Eigen::SparseMatrix<double> &Hess,
                                    bool FirstDerivativeWRTDef, double mu, double lambda ) {

  Hess.resize( V_def.rows() * V_def.cols(), V_def.rows() * V_def.cols());
  Hess.setZero();

  std::vector<Eigen::Triplet<double>> tripletList;
  tripletList.reserve( 9 * 9 * F.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, localMatrix( i, j ));
    }
    else {
      for ( int i = 0; i < 3; i++ )
        for ( int j = 0; j < 3; j++ )
          tripletList.emplace_back( i * numV + l, j * numV + k, localMatrix( j, i ));
    }
  };

  double lambdaQuarter = lambda / 4.;
  double muHalfPlusLambdaQuarter = mu / 2. + lambdaQuarter;

  for ( int faceIdx = 0; faceIdx < F.rows(); ++faceIdx ) {
    std::array<Eigen::Vector3d, 3> nodes, undefEdges, defEdges;
    Eigen::Vector3d temp;

    //! get undeformed quantities
    for ( int j = 0; j < 3; j++ )
      nodes[j] = V_undef.row( F( faceIdx, j ));
    for ( int j = 0; j < 3; j++ )
      undefEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];

    // compute volume
    temp = undefEdges[0].cross( undefEdges[1] );
    double volUndefSqr = temp.squaredNorm() / 4.;
    double volUndef = std::sqrt( volUndefSqr );

    // precompute undeformed area gradients
    std::array<Eigen::Vector3d, 3> gradUndefArea;
    for ( int i = 0; i < 3; i++ )
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], gradUndefArea[i] );

    //
    Eigen::Vector3d factors;
    for ( int i = 0; i < 3; i++ )
      factors[i] = undefEdges[( i + 1 ) % 3].dot( undefEdges[( i + 2 ) % 3] );


    //! get fixed edgess
    for ( int j : { 0, 1, 2 } )
      nodes[j] = V_def.row( F( faceIdx, j ));
    for ( int j : { 0, 1, 2 } )
      defEdges[j] = nodes[( j + 2 ) % 3] - nodes[( j + 1 ) % 3];

    // compute volume
    temp = ( nodes[1] - nodes[0] ).cross( nodes[2] - nodes[1] );
    double volDefSqr = temp.squaredNorm() / 4.;
    double volDef = std::sqrt( volDefSqr );

    // precomputed deformed trace gradients
    std::vector<Eigen::Vector3d> gradDefTrace( 3 );
    for ( int i = 0; i < 3; i++ )
      gradDefTrace[i] = -2 * factors[( i + 1 ) % 3] * defEdges[( i + 1 ) % 3] +
                        2 * factors[( i + 2 ) % 3] * defEdges[( i + 2 ) % 3];

    // precompute deformed area gradients
    std::vector<Eigen::Vector3d> gradDefArea( 3 );
    for ( int i = 0; i < 3; i++ )
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], gradDefArea[i] );

    // compute local matrices
    Eigen::Matrix3d tensorProduct, H, auxMat;
    double mixedTraceHessFactor = 0.125 * mu / volUndef;
    double MixedAreaFactor =
            -2. * ( lambdaQuarter * volDef / volUndef + muHalfPlusLambdaQuarter * volUndef / volDef ) / volUndef;
    double MixedFactor = -0.125 * mu / volUndefSqr;

    // i!=j
    for ( int i = 0; i < 3; i++ ) {
      for ( int j = 0; j < 3; j++ ) {
        // Hess trace term
        if ( i == j ) {
          H = undefEdges[i] * defEdges[i].transpose();
        }
        else {
          int k = ( 2 * i + 2 * j ) % 3;
          H = ( undefEdges[j] - undefEdges[k] ) * defEdges[i].transpose();
          H += undefEdges[i] * defEdges[k].transpose();
        }
        H *= -2 * mixedTraceHessFactor;

        // mixed area term
        H += MixedAreaFactor * gradUndefArea[i] * gradDefArea[j].transpose();

        // mixed term
        H += MixedFactor * gradUndefArea[i] * gradDefTrace[j].transpose();

        localToGlobal( F( faceIdx, i ), F( faceIdx, j ), H );
      }
    }

  }

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