/*  dynamo:- Event driven molecular dynamics simulator
    http://www.marcusbannerman.co.uk/dynamo
    Copyright (C) 2011  Marcus N Campbell Bannerman <m.bannerman@gmail.com>

    This program is free software: you can redistribute it and/or
    modify it under the terms of the GNU General Public License
    version 3 as published by the Free Software Foundation.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#pragma once

#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_proxy.hpp>
#include <exception>

namespace ublas = boost::numeric::ublas;

class Spline : private std::vector<std::pair<double, double> > {
 public:
  // The boundary conditions available
  enum BC_type { FIXED_1ST_DERIV_BC, FIXED_2ND_DERIV_BC, PARABOLIC_RUNOUT_BC };

  enum Spline_type { LINEAR, CUBIC };

  // Constructor takes the boundary conditions as arguments, this
  // sets the first derivative (gradient) at the lower and upper
  // end points
  Spline()
      : _valid(false),
        _BCLow(FIXED_2ND_DERIV_BC),
        _BCHigh(FIXED_2ND_DERIV_BC),
        _BCLowVal(0),
        _BCHighVal(0),
        _type(CUBIC) {}

  typedef std::vector<std::pair<double, double> > base;
  typedef base::const_iterator const_iterator;

  // Standard STL read-only container stuff
  const_iterator begin() const { return base::begin(); }
  const_iterator end() const { return base::end(); }
  void clear() {
    _valid = false;
    base::clear();
    _data.clear();
  }
  size_t size() const { return base::size(); }
  size_t max_size() const { return base::max_size(); }
  size_t capacity() const { return base::capacity(); }
  bool empty() const { return base::empty(); }

  // Add a point to the spline, and invalidate it so its
  // recalculated on the next access
  inline void addPoint(double x, double y) {
    _valid = false;
    base::push_back(std::pair<double, double>(x, y));
  }

  // Reset the boundary conditions
  inline void setLowBC(BC_type BC, double val = 0) {
    _BCLow = BC;
    _BCLowVal = val;
    _valid = false;
  }

  inline void setHighBC(BC_type BC, double val = 0) {
    _BCHigh = BC;
    _BCHighVal = val;
    _valid = false;
  }

  void setType(Spline_type type) {
    _type = type;
    _valid = false;
  }

  // Check if the spline has been calculated, then generate the
  // spline interpolated value
  double operator()(double xval) {
    if (!_valid) generate();

    // Special cases when we're outside the range of the spline points
    if (xval <= x(0)) return lowCalc(xval);
    if (xval >= x(size() - 1)) return highCalc(xval);

    // Check all intervals except the last one
    for (std::vector<SplineData>::const_iterator iPtr = _data.begin(); iPtr != _data.end() - 1; ++iPtr)
      if ((xval >= iPtr->x) && (xval <= (iPtr + 1)->x)) return splineCalc(iPtr, xval);

    return splineCalc(_data.end() - 1, xval);
  }

 private:
  ///////PRIVATE DATA MEMBERS
  struct SplineData {
    double x, a, b, c, d;
  };
  // vector of calculated spline data
  std::vector<SplineData> _data;
  // Second derivative at each point
  ublas::vector<double> _ddy;
  // Tracks whether the spline parameters have been calculated for
  // the current set of points
  bool _valid;
  // The boundary conditions
  BC_type _BCLow, _BCHigh;
  // The values of the boundary conditions
  double _BCLowVal, _BCHighVal;

  Spline_type _type;

  ///////PRIVATE FUNCTIONS
  // Function to calculate the value of a given spline at a point xval
  inline double splineCalc(std::vector<SplineData>::const_iterator i, double xval) {
    const double lx = xval - i->x;
    return ((i->a * lx + i->b) * lx + i->c) * lx + i->d;
  }

  inline double lowCalc(double xval) {
    const double lx = xval - x(0);

    if (_type == LINEAR) return lx * _BCHighVal + y(0);

    const double firstDeriv = (y(1) - y(0)) / h(0) - 2 * h(0) * (_data[0].b + 2 * _data[1].b) / 6;

    switch (_BCLow) {
      case FIXED_1ST_DERIV_BC:
        return lx * _BCLowVal + y(0);
      case FIXED_2ND_DERIV_BC:
        return lx * lx * _BCLowVal + firstDeriv * lx + y(0);
      case PARABOLIC_RUNOUT_BC:
        return lx * lx * _ddy[0] + lx * firstDeriv + y(0);
    }
    throw std::runtime_error("Unknown BC");
  }

  inline double highCalc(double xval) {
    const double lx = xval - x(size() - 1);

    if (_type == LINEAR) return lx * _BCHighVal + y(size() - 1);

    const double firstDeriv = 2 * h(size() - 2) * (_ddy[size() - 2] + 2 * _ddy[size() - 1]) / 6 +
                              (y(size() - 1) - y(size() - 2)) / h(size() - 2);

    switch (_BCHigh) {
      case FIXED_1ST_DERIV_BC:
        return lx * _BCHighVal + y(size() - 1);
      case FIXED_2ND_DERIV_BC:
        return lx * lx * _BCHighVal + firstDeriv * lx + y(size() - 1);
      case PARABOLIC_RUNOUT_BC:
        return lx * lx * _ddy[size() - 1] + lx * firstDeriv + y(size() - 1);
    }
    throw std::runtime_error("Unknown BC");
  }

  // These just provide access to the point data in a clean way
  inline double x(size_t i) const { return operator[](i).first; }
  inline double y(size_t i) const { return operator[](i).second; }
  inline double h(size_t i) const { return x(i + 1) - x(i); }

  // Invert a arbitrary matrix using the boost ublas library
  template <class T>
  bool InvertMatrix(ublas::matrix<T> A, ublas::matrix<T>& inverse) {
    using namespace ublas;

    // create a permutation matrix for the LU-factorization
    permutation_matrix<std::size_t> pm(A.size1());

    // perform LU-factorization
    int res = lu_factorize(A, pm);
    if (res != 0) return false;

    // create identity matrix of "inverse"
    inverse.assign(ublas::identity_matrix<T>(A.size1()));

    // backsubstitute to get the inverse
    lu_substitute(A, pm, inverse);

    return true;
  }

  // This function will recalculate the spline parameters and store
  // them in _data, ready for spline interpolation
  void generate() {
    if (size() < 2) throw std::runtime_error("Spline requires at least 2 points");

    // If any spline points are at the same x location, we have to
    // just slightly seperate them
    {
      bool testPassed(false);
      while (!testPassed) {
        testPassed = true;
        std::sort(base::begin(), base::end());

        for (base::iterator iPtr = base::begin(); iPtr != base::end() - 1; ++iPtr)
          if (iPtr->first == (iPtr + 1)->first) {
            if ((iPtr + 1)->first != 0)
              (iPtr + 1)->first += (iPtr + 1)->first * std::numeric_limits<double>::epsilon() * 10;
            else
              (iPtr + 1)->first = std::numeric_limits<double>::epsilon() * 10;
            testPassed = false;
            break;
          }
      }
    }

    const size_t e = size() - 1;

    switch (_type) {
      case LINEAR: {
        _data.resize(e);
        for (size_t i(0); i < e; ++i) {
          _data[i].x = x(i);
          _data[i].a = 0;
          _data[i].b = 0;
          _data[i].c = (y(i + 1) - y(i)) / (x(i + 1) - x(i));
          _data[i].d = y(i);
        }
        break;
      }
      case CUBIC: {
        ublas::matrix<double> A(size(), size());
        for (size_t yv(0); yv <= e; ++yv)
          for (size_t xv(0); xv <= e; ++xv) A(xv, yv) = 0;

        for (size_t i(1); i < e; ++i) {
          A(i - 1, i) = h(i - 1);
          A(i, i) = 2 * (h(i - 1) + h(i));
          A(i + 1, i) = h(i);
        }

        ublas::vector<double> C(size());
        for (size_t xv(0); xv <= e; ++xv) C(xv) = 0;

        for (size_t i(1); i < e; ++i) C(i) = 6 * ((y(i + 1) - y(i)) / h(i) - (y(i) - y(i - 1)) / h(i - 1));

        // Boundary conditions
        switch (_BCLow) {
          case FIXED_1ST_DERIV_BC:
            C(0) = 6 * ((y(1) - y(0)) / h(0) - _BCLowVal);
            A(0, 0) = 2 * h(0);
            A(1, 0) = h(0);
            break;
          case FIXED_2ND_DERIV_BC:
            C(0) = _BCLowVal;
            A(0, 0) = 1;
            break;
          case PARABOLIC_RUNOUT_BC:
            C(0) = 0;
            A(0, 0) = 1;
            A(1, 0) = -1;
            break;
        }

        switch (_BCHigh) {
          case FIXED_1ST_DERIV_BC:
            C(e) = 6 * (_BCHighVal - (y(e) - y(e - 1)) / h(e - 1));
            A(e, e) = 2 * h(e - 1);
            A(e - 1, e) = h(e - 1);
            break;
          case FIXED_2ND_DERIV_BC:
            C(e) = _BCHighVal;
            A(e, e) = 1;
            break;
          case PARABOLIC_RUNOUT_BC:
            C(e) = 0;
            A(e, e) = 1;
            A(e - 1, e) = -1;
            break;
        }

        ublas::matrix<double> AInv(size(), size());
        InvertMatrix(A, AInv);

        _ddy = ublas::prod(C, AInv);

        _data.resize(size() - 1);
        for (size_t i(0); i < e; ++i) {
          _data[i].x = x(i);
          _data[i].a = (_ddy(i + 1) - _ddy(i)) / (6 * h(i));
          _data[i].b = _ddy(i) / 2;
          _data[i].c = (y(i + 1) - y(i)) / h(i) - _ddy(i + 1) * h(i) / 6 - _ddy(i) * h(i) / 3;
          _data[i].d = y(i);
        }
      }
    }
    _valid = true;
  }
};
