// $Id$
//
//  Copyright (C) 2004-2006 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 "DistViolationContrib.h"
#include <ForceField/ForceField.h>
#include <RDGeneral/Invariant.h>

namespace DistGeom {

DistViolationContrib::DistViolationContrib(ForceFields::ForceField *owner,
                                           unsigned int idx1, unsigned int idx2,
                                           double ub, double lb,
                                           double weight) {
  PRECONDITION(owner, "bad owner");
  URANGE_CHECK(idx1, owner->positions().size() - 1);
  URANGE_CHECK(idx2, owner->positions().size() - 1);

  dp_forceField = owner;
  d_end1Idx = idx1;
  d_end2Idx = idx2;
  d_ub = ub;
  d_lb = lb;
  d_weight = weight;
}

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

  double d =
      this->dp_forceField->distance(this->d_end1Idx, this->d_end2Idx, pos);
  double val = 0.0;
  if (d > d_ub) {
    val = ((d * d) / (d_ub * d_ub)) - 1.0;
  } else if (d < d_lb) {
    val = ((2 * d_lb * d_lb) / (d_lb * d_lb + d * d)) - 1.0;
  }
  double res;
  if (val > 0.0) {
    res = d_weight * val * val;
  } else {
    res = 0;
  }
  return res;
}

void DistViolationContrib::getGrad(double *pos, double *grad) const {
  PRECONDITION(dp_forceField, "no owner");
  PRECONDITION(pos, "bad vector");
  PRECONDITION(grad, "bad vector");
  unsigned int dim = this->dp_forceField->dimension();
  double d =
      this->dp_forceField->distance(this->d_end1Idx, this->d_end2Idx, pos);
  double preFactor = 0.0;
  if (d > d_ub) {
    double u2 = d_ub * d_ub;
    preFactor = 4. * (((d * d) / u2) - 1.0) * (d / u2);
  } else if (d < d_lb) {
    double d2 = d * d;
    double l2 = d_lb * d_lb;
    double l2d2 = d2 + l2;
    preFactor = 8. * l2 * d * (1. - 2 * l2 / l2d2) / (l2d2 * l2d2);
    // preFactor = -8.*((l2-d2)/pow(l2+d2,3))*d*l2;
  } else {
    return;
  }

  double *end1Coords = &(pos[dim * this->d_end1Idx]);
  double *end2Coords = &(pos[dim * this->d_end2Idx]);

  for (unsigned int i = 0; i < dim; i++) {
    double dGrad;
    if (d > 0.0) {
      dGrad = d_weight * preFactor * (end1Coords[i] - end2Coords[i]) / d;
    } else {
      // FIX: this likely isn't right
      dGrad = d_weight * preFactor * (end1Coords[i] - end2Coords[i]);
    }
    grad[dim * this->d_end1Idx + i] += dGrad;
    grad[dim * this->d_end2Idx + i] -= dGrad;
  }
}
}