/**
 *  \file PMFTable.h
 *  \brief A table for storing potentials
 *
 *  Copyright 2007-2021 IMP Inventors. All rights reserved.
 *
 */

#ifndef IMPSCORE_FUNCTOR_INTERNAL_PMF_TABLE_H
#define IMPSCORE_FUNCTOR_INTERNAL_PMF_TABLE_H

#include <IMP/score_functor/score_functor_config.h>
#include "RawOpenCubicSpline.h"
#include <IMP/algebra/GridD.h>
#include <IMP/algebra/grid_storages.h>
#include <IMP/exception.h>
#include <IMP/Object.h>
#include <IMP/Array.h>
#include <IMP/file.h>
#include <cmath>
#include <vector>

IMPSCOREFUNCTOR_BEGIN_INTERNAL_NAMESPACE

// slightly evil using the grid storage, but...
template <bool SPARSE>
struct StorageSelector {
  typedef algebra::DenseGridStorageD<2, RawOpenCubicSpline> Type;
};

template <>
struct StorageSelector<true> {
  typedef algebra::SparseGridStorageD<2, RawOpenCubicSpline,
                                      algebra::UnboundedGridRangeD<2> > Type;
};

template <bool BIPARTITE, bool INTERPOLATE, bool SPARSE = false>
struct PMFTable : public Object {
  unsigned int split_;
  double inverse_bin_width_;
  double bin_width_;
  double max_;
  double offset_;
  typedef typename StorageSelector<SPARSE>::Type Storage;
  Storage data_;
  void order(unsigned int &i, unsigned int &j) const {
    if (i > j) {
      std::swap(i, j);
    }
    IMP_USAGE_CHECK(!BIPARTITE || j >= split_,
                    "One of the particles should be of each type: " << i << " "
                                                                    << j);
    j -= split_;
  }
  const RawOpenCubicSpline &get(int i, int j) const {
    Array<2, int> is;
    is[0] = i;
    is[1] = j;
    typename Storage::ExtendedIndex ei(is.begin(), is.end());
    return data_[data_.get_index(ei)];
  }
  template <class Key>
  void initialize(TextInput tin) {
    std::istream &in = tin;
    double bin;
    std::string line;
    std::getline(in, line);
    std::istringstream iss(line);
    iss >> bin;
    if (!iss) {
      IMP_THROW("Error reading bin size from line " << line, IOException);
    }
    unsigned int np, nl;
    iss >> np;
    if (BIPARTITE) {
      iss >> nl;
    } else {
      nl = np;
    }
    if (!iss) {
      IMP_THROW("Error number of types from line " << line, IOException);
    }
    iss >> offset_;
    if (!iss) {
      offset_ = 0;
    }
    IMP_LOG_TERSE("Reading " << np << " by " << nl << " from file "
                             << tin.get_name() << std::endl);
    /*if (np != ni
        || nl != nj) {
      IMP_THROW("Expected number of protein and ligand types not found. "
                << "Expected " << ni
                << " " << nj
                << " but got " << np << " " << nl,
                IOException);
                }*/

    bin_width_ = bin;
    inverse_bin_width_ = 1.0 / bin;
    Ints dims(2);
    dims[0] = np;
    dims[1] = nl;
    data_ = Storage(dims);
    int bins_read = -1;
    unsigned int read_entries = 0;
    while (true) {
      std::getline(in, line);
      if (!in) {
        break;
      }
      std::istringstream ins(line);
      std::string pname, lname;
      ins >> pname >> lname;
      unsigned int i = Key(pname).get_index();
      unsigned int j = Key(lname).get_index();
      IMP_INTERNAL_CHECK(Key(pname) == Key(i),
                         "Expected and found protein types don't match: "
                             << "expected \"" << Key(i).get_string() << " got "
                             << pname << " at " << i << " " << j);
      IMP_INTERNAL_CHECK(Key(lname) == Key(j),
                         "Expected and found ligand types don't match: "
                             << "expected \"" << Key(j).get_string() << " got "
                             << pname << " at " << i << " " << j);
      //      std::cout << line << std::endl;
      //      if(j>1) exit(EXIT_FAILURE);
      int cur_bins_read = 0;
      Floats data;
      while (true) {
        double potentialvalue;
        ins >> potentialvalue;
        if (ins) {
          data.push_back(potentialvalue);
          ++cur_bins_read;
        } else {
          break;
        }
      }
      order(i, j);
      Array<2, int> is;
      is[0] = i;
      is[1] = j;
      typename Storage::ExtendedIndex ei(is.begin(), is.end());
      if (!data_.get_has_index(ei)) {
        data_.add_voxel(ei, score_functor::internal::RawOpenCubicSpline(
                                data, bin_width_, inverse_bin_width_));
      } else {
        data_[data_.get_index(ei)] =
            score_functor::internal::RawOpenCubicSpline(data, bin_width_,
                                                        inverse_bin_width_);
      }
      if (bins_read != -1 && cur_bins_read != bins_read) {
        IMP_THROW("Read wrong number of bins from line: "
                      << line << "\nExpected " << bins_read << " got "
                      << cur_bins_read,
                  IOException);
      }
      bins_read = cur_bins_read;
      ins.clear();
      ++read_entries;
    }
    max_ = bin_width_ * bins_read;
    if (BIPARTITE) {
      if (read_entries != np * nl) {
        IMP_THROW("Read " << read_entries << " from table, but expected all of"
                          << np << "x" << nl << "=" << np * nl,
                  IOException);
      }
    } else {
      if (read_entries != np * (np + 1) / 2) {
        IMP_THROW("Read " << read_entries << " from table, but expected all of"
                          << np << "x" << np + 1 << "/2=" << np * (np + 1) / 2,
                  IOException);
      }
    }
    IMP_LOG_TERSE("PMF table entries have " << bins_read << " bins with width "
                                            << bin_width_ << std::endl);
  }

 public:
  template <class Key>
  PMFTable(TextInput name, unsigned int split, Key)
      : Object("PMFTable-" + name.get_name()), split_(split) {
    initialize<Key>(name);
  }

  double get_score(unsigned int i, unsigned int j, double dist) const {
    if (dist >= max_ || dist <= offset_) return 0;
    order(i, j);
    if (INTERPOLATE) {
      return get(i, j).evaluate(dist - .5 * bin_width_ - offset_, bin_width_,
                                inverse_bin_width_);
    } else {
      return get(i, j).get_bin(dist - offset_, bin_width_, inverse_bin_width_);
    }
  }
  double get_max() const { return max_; }
  DerivativePair get_score_with_derivative(unsigned int i, unsigned int j,
                                           double dist) const {
    if (dist >= max_ - .5 * bin_width_ || dist <= offset_) {
      return DerivativePair(0, 0);
    }
    order(i, j);
    if (dist <= .5 * bin_width_)
      return DerivativePair(get_score(i, j, dist), 0);
    // shift by .5 for the splines so as to be between the centers of the cells
    return get(i, j).evaluate_with_derivative(dist - .5 * bin_width_ - offset_,
                                              bin_width_, inverse_bin_width_);
  }
};

IMPSCOREFUNCTOR_END_INTERNAL_NAMESPACE

#endif /* IMPSCORE_FUNCTOR_INTERNAL_PMF_TABLE_H */
