// $Id$
//
//  Copyright (C) 2004-2013 Rational Discovery LLC
//
//   @@ All Rights Reserved @@
//  This file is part of the RDKit.
//  The contents are covered by the terms of the BSD license
//  which is included in the file license.txt, found at the root
//  of the RDKit source tree.
//
#include "AngleBend.h"
#include "BondStretch.h"
#include "Params.h"
#include <math.h>
#include <ForceField/ForceField.h>
#include <RDGeneral/Invariant.h>
#include <RDGeneral/utils.h>

namespace ForceFields {
  namespace UFF {

    namespace Utils {
      double calcAngleForceConstant(double theta0,
                                    double bondOrder12,double bondOrder23,
                                    const AtomicParams *at1Params,
                                    const AtomicParams *at2Params,
                                    const AtomicParams *at3Params){
        double cosTheta0=cos(theta0);
        double r12 = calcBondRestLength(bondOrder12,at1Params,at2Params);
        double r23 = calcBondRestLength(bondOrder23,at2Params,at3Params);
        double r13 = sqrt(r12*r12 + r23*r23 - 2.*r12*r23*cosTheta0);
        double beta = 2.*Params::G/(r12*r23);

        double preFactor = beta*at1Params->Z1*at3Params->Z1 / int_pow<5>(r13);
        double rTerm = r12*r23;
        double innerBit = 3.*rTerm*(1.-cosTheta0*cosTheta0) - r13*r13*cosTheta0;
        double res=preFactor*rTerm*innerBit;
        return res;
      }
      
      void calcAngleBendGrad(RDGeom::Point3D *r, double *dist,
        double **g, double &dE_dTheta, double &cosTheta, double &sinTheta)
      {
        // -------
        // dTheta/dx is trickier:
        double dCos_dS[6] = {
          1.0 / dist[0] * (r[1].x - cosTheta * r[0].x),
          1.0 / dist[0] * (r[1].y - cosTheta * r[0].y),
          1.0 / dist[0] * (r[1].z - cosTheta * r[0].z),
          1.0 / dist[1] * (r[0].x - cosTheta * r[1].x),
          1.0 / dist[1] * (r[0].y - cosTheta * r[1].y),
          1.0 / dist[1] * (r[0].z - cosTheta * r[1].z)
        };

        g[0][0] += dE_dTheta * dCos_dS[0] / (-sinTheta);
        g[0][1] += dE_dTheta * dCos_dS[1] / (-sinTheta);
        g[0][2] += dE_dTheta * dCos_dS[2] / (-sinTheta);

        g[1][0] += dE_dTheta * (-dCos_dS[0] - dCos_dS[3]) / (-sinTheta);
        g[1][1] += dE_dTheta * (-dCos_dS[1] - dCos_dS[4]) / (-sinTheta);
        g[1][2] += dE_dTheta * (-dCos_dS[2] - dCos_dS[5]) / (-sinTheta);
      
        g[2][0] += dE_dTheta * dCos_dS[3] / (-sinTheta);
        g[2][1] += dE_dTheta * dCos_dS[4] / (-sinTheta);
        g[2][2] += dE_dTheta * dCos_dS[5] / (-sinTheta);
      }
    } // end of namespace Utils
  
    AngleBendContrib::AngleBendContrib(ForceField *owner,
               unsigned int idx1,unsigned int idx2,unsigned int idx3,
               double bondOrder12,double bondOrder23,
               const AtomicParams *at1Params,
               const AtomicParams *at2Params,
               const AtomicParams *at3Params,
               unsigned int order){
      PRECONDITION(owner,"bad owner");
      PRECONDITION(at1Params,"bad params pointer");
      PRECONDITION(at2Params,"bad params pointer");
      PRECONDITION(at3Params,"bad params pointer");
      PRECONDITION((idx1!=idx2&&idx2!=idx3&&idx1!=idx3),"degenerate points");
      RANGE_CHECK(0,idx1,owner->positions().size()-1);
      RANGE_CHECK(0,idx2,owner->positions().size()-1);
      RANGE_CHECK(0,idx3,owner->positions().size()-1);
      // the following is a hack to get decent geometries
      // with 3- and 4-membered rings incorporating sp2 atoms
      double theta0 = at2Params->theta0;
      if (order >= 30) {
        switch (order) {
          case 30:
            theta0 = 150.0 / 180.0 * M_PI;
            break;
          case 35:
            theta0 = 60.0 / 180.0 * M_PI;
            break;
          case 40:
            theta0 = 135.0 / 180.0 * M_PI;
            break;
          case 45:
            theta0 = 90.0 / 180.0 * M_PI;
            break;
        }
        order = 0;
      }
      // end of the hack
      dp_forceField = owner;
      d_at1Idx = idx1;
      d_at2Idx = idx2;
      d_at3Idx = idx3;
      d_order = order;
      d_forceConstant = Utils::calcAngleForceConstant
                  (theta0, bondOrder12,bondOrder23,
                  at1Params,at2Params,at3Params);
      if(order==0){
  double sinTheta0=sin(theta0);
  double cosTheta0=cos(theta0);
  d_C2 = 1./(4.*std::max(sinTheta0*sinTheta0,1e-8));
  d_C1 = -4.*d_C2*cosTheta0;
  d_C0 = d_C2*(2.*cosTheta0*cosTheta0 + 1.);
      }
    }

    double AngleBendContrib::getEnergy(double *pos) const {
      PRECONDITION(dp_forceField,"no owner");
      PRECONDITION(pos,"bad vector");

      double dist1=dp_forceField->distance(d_at1Idx,d_at2Idx,pos);
      double dist2=dp_forceField->distance(d_at2Idx,d_at3Idx,pos);

      RDGeom::Point3D p1(pos[3*d_at1Idx],
       pos[3*d_at1Idx+1],
       pos[3*d_at1Idx+2]);
      RDGeom::Point3D p2(pos[3*d_at2Idx],
       pos[3*d_at2Idx+1],
       pos[3*d_at2Idx+2]);
      RDGeom::Point3D p3(pos[3*d_at3Idx],
       pos[3*d_at3Idx+1],
       pos[3*d_at3Idx+2]);
      RDGeom::Point3D p12=p1-p2;
      RDGeom::Point3D p32=p3-p2;
      double cosTheta = p12.dotProduct(p32)/(dist1*dist2);
      clipToOne(cosTheta);
      // we need sin^2(theta) to get cos(2*theta), so compute that:
      double sinThetaSq = 1.-cosTheta*cosTheta;
    
      double angleTerm = getEnergyTerm(cosTheta,sinThetaSq);
      double res = d_forceConstant*angleTerm;

      return res;
    }

    void AngleBendContrib::getGrad(double *pos,double *grad) const {
      PRECONDITION(dp_forceField,"no owner");
      PRECONDITION(pos,"bad vector");
      PRECONDITION(grad,"bad vector");

      double dist[2] = {
        dp_forceField->distance(d_at1Idx, d_at2Idx, pos),
        dp_forceField->distance(d_at2Idx, d_at3Idx, pos)
      };

      RDGeom::Point3D p1(pos[3 * d_at1Idx],
        pos[3 * d_at1Idx + 1], pos[3 * d_at1Idx + 2]);
      RDGeom::Point3D p2(pos[3 * d_at2Idx],
        pos[3 * d_at2Idx + 1], pos[3 * d_at2Idx + 2]);
      RDGeom::Point3D p3(pos[3 * d_at3Idx],
        pos[3 * d_at3Idx + 1], pos[3 * d_at3Idx + 2]);
      double *g[3] = {
        &(grad[3 * d_at1Idx]),
        &(grad[3 * d_at2Idx]),
        &(grad[3 * d_at3Idx])
      };
      RDGeom::Point3D r[2] = {
        (p1 - p2) / dist[0],
        (p3 - p2) / dist[1]
      };
      double cosTheta = r[0].dotProduct(r[1]);
      clipToOne(cosTheta);
      double sinThetaSq = 1.0 - cosTheta * cosTheta;
      double sinTheta = std::max(((sinThetaSq > 0.0) ? sqrt(sinThetaSq) : 0.0), 1.0e-8);

      //std::cerr << "GRAD: " << cosTheta << " (" << acos(cosTheta)<< "), ";
      //std::cerr << sinTheta << " (" << asin(sinTheta)<< ")" << std::endl;
    
      // use the chain rule:
      // dE/dx = dE/dTheta * dTheta/dx

      // dE/dTheta is independent of cartesians:
      double dE_dTheta=getThetaDeriv(cosTheta,sinTheta);
    
      Utils::calcAngleBendGrad(r, dist, g, dE_dTheta, cosTheta, sinTheta);
    }


    double AngleBendContrib::getEnergyTerm(double cosTheta,double sinThetaSq) const {
      PRECONDITION(d_order==0||d_order==1||d_order==2||d_order==3||d_order==4,"bad order");
      // cos(2x) = cos^2(x) - sin^2(x);
      double cos2Theta = cosTheta*cosTheta - sinThetaSq;

      double res=0.0;
      if(d_order==0){
  res=d_C0 + d_C1*cosTheta + d_C2*cos2Theta;
      } else {
  switch(d_order){
  case 1:
    res=-cosTheta;
    break;
  case 2:
    res=cos2Theta;
    break;
  case 3:
    // cos(3x) = cos^3(x) - 3*cos(x)*sin^2(x)
    res = cosTheta*(cosTheta*cosTheta-3.*sinThetaSq);
    break;
  case 4:
    // cos(4x) = cos^4(x) - 6*cos^2(x)*sin^2(x)+sin^4(x)
    res = int_pow<4>(cosTheta) - 6.*cosTheta*cosTheta*sinThetaSq + sinThetaSq*sinThetaSq;
    break;
  }
  res = 1.-res;
  res /= (double)(d_order*d_order);
      }
      return res;
    }


  
    double AngleBendContrib::getThetaDeriv(double cosTheta,double sinTheta) const {
      PRECONDITION(d_order==0||d_order==1||d_order==2||d_order==3||d_order==4,"bad order");

      double dE_dTheta=0.0;
      double sin2Theta = 2.*sinTheta*cosTheta;

      if(d_order==0){
  dE_dTheta =  -1.*d_forceConstant*(d_C1*sinTheta +
                 2.*d_C2*sin2Theta);
      } else {
  // E = k/n^2 [1-cos(n theta)]
  // dE = - k/n^2 * d cos(n theta)

  // these all use:
  // d cos(ax) = -a sin(ax)
      
  switch(d_order){
  case 1:
    dE_dTheta = -sinTheta;
    break;
  case 2:
    // sin(2*x) = 2*cos(x)*sin(x)
    dE_dTheta = sin2Theta;
    break;
  case 3:
    // sin(3*x) = 3*sin(x) - 4*sin^3(x)
    dE_dTheta = sinTheta*(3.-4.*sinTheta*sinTheta);
    break;
  case 4:
    // sin(4*x) = cos(x)*(4*sin(x) - 8*sin^3(x))
    dE_dTheta = cosTheta*sinTheta*(4.-8.*sinTheta*sinTheta);
    break;
  }
  dE_dTheta *= d_forceConstant/(double)(d_order);
      }
      return dE_dTheta;
    }
  
  }
}
