/*
 * This file belongs to the Galois project, a C++ library for exploiting parallelism.
 * The code is being released under the terms of the 3-Clause BSD License (a
 * copy is located in LICENSE.txt at the top-level directory).
 *
 * Copyright (C) 2018, The University of Texas at Austin. All rights reserved.
 * UNIVERSITY EXPRESSLY DISCLAIMS ANY AND ALL WARRANTIES CONCERNING THIS
 * SOFTWARE AND DOCUMENTATION, INCLUDING ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR ANY PARTICULAR PURPOSE, NON-INFRINGEMENT AND WARRANTIES OF
 * PERFORMANCE, AND ANY WARRANTY THAT MIGHT OTHERWISE ARISE FROM COURSE OF
 * DEALING OR USAGE OF TRADE.  NO WARRANTY IS EITHER EXPRESS OR IMPLIED WITH
 * RESPECT TO THE USE OF THE SOFTWARE OR DOCUMENTATION. Under no circumstances
 * shall University be liable for incidental, special, indirect, direct or
 * consequential damages or loss of profits, interruption of business, or
 * related expenses which may arise from use of Software or Documentation,
 * including but not limited to those resulting from defects in Software and/or
 * Documentation, or loss or inaccuracy of data of any kind.
 */

/*  This file is part of libDAI - http://www.libdai.org/
 *
 *  Copyright (c) 2006-2011, The libDAI authors. All rights reserved.
 *
 *  Use of this source code is governed by a BSD-style license that can be found
 * in the LICENSE file.
 */

#include <dai/alldai.h>

namespace dai {

using namespace std;

void DecMAP::setProperties(const PropertySet& opts) {
  DAI_ASSERT(opts.hasKey("ianame"));
  DAI_ASSERT(opts.hasKey("iaopts"));

  props.ianame = opts.getStringAs<string>("ianame");
  props.iaopts = opts.getStringAs<PropertySet>("iaopts");
  if (opts.hasKey("verbose"))
    props.verbose = opts.getStringAs<size_t>("verbose");
  else
    props.verbose = 0;
  if (opts.hasKey("reinit"))
    props.reinit = opts.getStringAs<bool>("reinit");
  else
    props.reinit = true;
}

PropertySet DecMAP::getProperties() const {
  PropertySet opts;
  opts.set("verbose", props.verbose);
  opts.set("reinit", props.reinit);
  opts.set("ianame", props.ianame);
  opts.set("iaopts", props.iaopts);
  return opts;
}

string DecMAP::printProperties() const {
  stringstream s(stringstream::out);
  s << "[";
  s << "verbose=" << props.verbose << ",";
  s << "reinit=" << props.reinit << ",";
  s << "ianame=" << props.ianame << ",";
  s << "iaopts=" << props.iaopts << "]";
  return s.str();
}

DecMAP::DecMAP(const FactorGraph& fg, const PropertySet& opts)
    : DAIAlgFG(fg), _state(), _logp(), _maxdiff(), _iters(), props() {
  setProperties(opts);

  _state = vector<size_t>(nrVars(), 0);
  _logp  = -INFINITY;
}

Factor DecMAP::belief(const VarSet& vs) const {
  if (vs.size() == 0)
    return Factor();
  else {
    map<Var, size_t> state;
    for (VarSet::const_iterator v = vs.begin(); v != vs.end(); v++)
      state[*v] = _state[findVar(*v)];
    return createFactorDelta(vs, calcLinearState(vs, state));
  }
}

Factor DecMAP::beliefV(size_t i) const {
  return createFactorDelta(var(i), _state[i]);
}

vector<Factor> DecMAP::beliefs() const {
  vector<Factor> result;
  for (size_t i = 0; i < nrVars(); ++i)
    result.push_back(beliefV(i));
  for (size_t I = 0; I < nrFactors(); ++I)
    result.push_back(beliefF(I));
  return result;
}

Real DecMAP::run() {
  if (props.verbose >= 1)
    cerr << "Starting " << identify() << "...";
  if (props.verbose >= 2)
    cerr << endl;

  // the variables which have not been clamped yet
  SmallSet<size_t> freeVars;
  for (size_t i = 0; i < nrVars(); i++)
    freeVars |= i;

  // prepare the inference algorithm object
  InfAlg* clamped = newInfAlg(props.ianame, fg(), props.iaopts);

  // decimate until no free variables remain
  while (freeVars.size()) {
    Real md = clamped->run();
    if (md > _maxdiff)
      _maxdiff = md;
    _iters += clamped->Iterations();

    // store the variables that need initialization
    VarSet varsToInit;
    SmallSet<size_t> varsToClamp;

    // schedule clamping for the free variables with zero entropy
    for (SmallSet<size_t>::const_iterator it = freeVars.begin();
         it != freeVars.end();) {
      if (clamped->beliefV(*it).entropy() == 0.0) {
        // this variable should be clamped
        varsToInit |= var(*it);
        varsToClamp |= *it;
        _state[*it] = clamped->beliefV(*it).p().argmax().first;
        freeVars.erase(*it);
      } else
        it++;
    }

    // find the free factor with lowest entropy
    size_t bestI = 0;
    Real bestEnt = INFINITY;
    for (size_t I = 0; I < nrFactors(); I++) {
      // check if the factor is still free
      if (freeVars.intersects(bipGraph().nb2Set(I))) {
        Real EntI = clamped->beliefF(I).entropy();
        if (EntI < bestEnt) {
          bestI   = I;
          bestEnt = EntI;
        }
      }
    }

    // schedule clamping for the factor with lowest entropy
    vector<size_t> Istate(1, 0);
    Istate[0]                  = clamped->beliefF(bestI).p().argmax().first;
    map<Var, size_t> Istatemap = calcState(factor(bestI).vars(), Istate[0]);
    diaforeach(size_t i, bipGraph().nb2Set(bestI) & freeVars) {
      varsToInit |= var(i);
      varsToClamp |= i;
      _state[i] = Istatemap[var(i)];
      freeVars.erase(i);
    }

    // clamp all variables scheduled for clamping
    diaforeach(size_t i, varsToClamp) clamped->clamp(i, _state[i], false);

    // initialize clamped for the next run
    if (props.reinit)
      clamped->init();
    else
      clamped->init(varsToInit);
  }

  // calculate MAP state
  map<Var, size_t> state;
  for (size_t i = 0; i < nrVars(); i++)
    state[var(i)] = _state[i];
  _logp = 0.0;
  for (size_t I = 0; I < nrFactors(); I++)
    _logp += dai::log(factor(I)[calcLinearState(factor(I).vars(), state)]);

  // clean up
  delete clamped;

  return _maxdiff;
}

} // end of namespace dai
