// 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_gradient( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
                                        const Eigen::MatrixXi &F, Eigen::MatrixXd &grad,
                                        double mu, double lambda ) {

  grad.resize( V_def.rows(), V_def.cols());
  grad.setZero();

  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 );

    //! trace part of gradient, e_tr = volUndef * _mu/2. *  traceDistTensor   
    Eigen::Vector3d factors;
    for ( int i : { 0, 1, 2 } )
      factors[i] = -0.25 * mu * undefEdges[( i + 2 ) % 3].dot( undefEdges[( i + 1 ) % 3] ) / volUndef;
    double factor = 2. * ( lambdaQuarter * volDef / volUndef - muHalfPlusLambdaQuarter * volUndef / volDef );

    for ( int i : { 0, 1, 2 } ) {
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], temp );
      for ( int j : { 0, 1, 2 } )
        grad( F( faceIdx, i ), j ) += factor * temp[j] + factors[( i + 1 ) % 3] * defEdges[( i + 1 ) % 3][j] -
                                      factors[( i + 2 ) % 3] * defEdges[( i + 2 ) % 3][j];
    }

  }
}

void shell::membrane_undeformed_gradient( const Eigen::MatrixXd &V_undef, const Eigen::MatrixXd &V_def,
                                          const Eigen::MatrixXi &F, Eigen::MatrixXd &grad,
                                          double mu, double lambda ) {

  grad.resize( V_def.rows(), V_def.cols());
  grad.setZero();

  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];

    double factor1 = ( 0.125 * mu * traceTerm + lambdaQuarter * volDefSqr ) / volUndefSqr;
    double factor2 = muHalfPlusLambdaQuarter * std::log( volDefSqr / volUndefSqr ) + ( mu + lambdaQuarter ) -
                     2 * muHalfPlusLambdaQuarter;
    double factorAreaGrad = factor1 + factor2;
    double factorTraceGrad = 0.125 * mu / volUndef;

    std::array<Eigen::Vector3d, 3> gradTrace;
    for ( int i : { 0, 1, 2 } )
      for ( int j : { 0, 1, 2 } )
        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] );

    // E = (0.125 * _mu * traceTerm + _lambdaQuarter * volDefSqr ) / volUndef;
    // grad E[j] = 0.125 * _mu * grad traceTerm[j] / volUndef  - (0.125 * _mu * traceTerm + _lambdaQuarter * volDefSqr )  grad volUndef[j] / volUndef^2
    for ( int i : { 0, 1, 2 } ) {
      getAreaGradient( nodes[( i + 1 ) % 3], nodes[( i + 2 ) % 3], nodes[i], temp );
      for ( int j : { 0, 1, 2 } )
        grad( F( faceIdx, i ), j ) += factorTraceGrad * gradTrace[i][j] - factorAreaGrad * temp[j];
    }

  }
}