/****************************************************************************
 *   This file is part of the aGrUM/pyAgrum library.                        *
 *                                                                          *
 *   Copyright (c) 2005-2026 by                                             *
 *       - Pierre-Henri WUILLEMIN(_at_LIP6)                                 *
 *       - Christophe GONZALES(_at_AMU)                                     *
 *                                                                          *
 *   The aGrUM/pyAgrum library is free software; you can redistribute it    *
 *   and/or modify it under the terms of either :                           *
 *                                                                          *
 *    - the GNU Lesser General Public License as published by               *
 *      the Free Software Foundation, either version 3 of the License,      *
 *      or (at your option) any later version,                              *
 *    - the MIT license (MIT),                                              *
 *    - or both in dual license, as here.                                   *
 *                                                                          *
 *   (see https://agrum.gitlab.io/articles/dual-licenses-lgplv3mit.html)    *
 *                                                                          *
 *   This aGrUM/pyAgrum library is distributed in the hope that it will be  *
 *   useful, but WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,          *
 *   INCLUDING BUT NOT LIMITED TO THE WARRANTIES MERCHANTABILITY or FITNESS *
 *   FOR A PARTICULAR PURPOSE  AND NONINFRINGEMENT. IN NO EVENT SHALL THE   *
 *   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER *
 *   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,        *
 *   ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR  *
 *   OTHER DEALINGS IN THE SOFTWARE.                                        *
 *                                                                          *
 *   See LICENCES for more details.                                         *
 *                                                                          *
 *   SPDX-FileCopyrightText: Copyright 2005-2026                            *
 *       - Pierre-Henri WUILLEMIN(_at_LIP6)                                 *
 *       - Christophe GONZALES(_at_AMU)                                     *
 *   SPDX-License-Identifier: LGPL-3.0-or-later OR MIT                      *
 *                                                                          *
 *   Contact  : info_at_agrum_dot_org                                       *
 *   homepage : http://agrum.gitlab.io                                      *
 *   gitlab   : https://gitlab.com/agrumery/agrum                           *
 *                                                                          *
 ****************************************************************************/

#pragma once


/**
 * @file
 * @brief Implementation of Loopy Belief Propagation in Bayesian networks.
 */
#include <agrum/BN/inference/loopyBeliefPropagation.h>   // to ease IDE parser
#ifndef DOXYGEN_SHOULD_SKIP_THIS

#  include <algorithm>
#  include <sstream>
#  include <string>

constexpr auto LBP_DEFAULT_MAXITER          = 100;
constexpr auto LBP_DEFAULT_EPSILON          = 1e-8;
constexpr auto LBP_DEFAULT_MIN_EPSILON_RATE = 1e-10;
constexpr auto LBP_DEFAULT_PERIOD_SIZE      = 1;
constexpr auto LBP_DEFAULT_VERBOSITY        = false;


// to ease parsing for IDE
#  include <agrum/BN/inference/loopyBeliefPropagation.h>

namespace gum {

  /// default constructor
  template < GUM_Numeric GUM_SCALAR >
  LoopyBeliefPropagation< GUM_SCALAR >::LoopyBeliefPropagation(const IBayesNet< GUM_SCALAR >* bn) :
      ApproximateInference< GUM_SCALAR >(bn) {
    // for debugging purposes
    GUM_CONSTRUCTOR(LoopyBeliefPropagation)

    this->setEpsilon(LBP_DEFAULT_EPSILON);
    this->setMinEpsilonRate(LBP_DEFAULT_MIN_EPSILON_RATE);
    this->setMaxIter(LBP_DEFAULT_MAXITER);
    this->setVerbosity(LBP_DEFAULT_VERBOSITY);
    this->setPeriodSize(LBP_DEFAULT_PERIOD_SIZE);

    _init_messages_();
  }

  /// destructor
  template < GUM_Numeric GUM_SCALAR >
  LoopyBeliefPropagation< GUM_SCALAR >::~LoopyBeliefPropagation() {
    GUM_DESTRUCTOR(LoopyBeliefPropagation)
  }

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::_init_messages_() {
    _messages_.clear();
    for (const auto& tail: this->BN().nodes()) {
      Tensor< GUM_SCALAR > p;
      p.add(this->BN().variable(tail));
      p.fill(static_cast< GUM_SCALAR >(1));

      for (const auto& head: this->BN().children(tail)) {
        _messages_.insert(Arc(head, tail), p);
        _messages_.insert(Arc(tail, head), p);
      }
    }
  }

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::updateOutdatedStructure_() {
    _init_messages_();
  }

  template < GUM_Numeric GUM_SCALAR >
  Tensor< GUM_SCALAR > LoopyBeliefPropagation< GUM_SCALAR >::_computeProdPi_(NodeId X) {
    const auto& varX = this->BN().variable(X);

    auto piX = this->BN().cpt(X);
    for (const auto& U: this->BN().parents(X)) {
      piX *= _messages_[Arc(U, X)];
    }
    piX = piX.sumIn({&varX});

    return piX;
  }

  template < GUM_Numeric GUM_SCALAR >
  Tensor< GUM_SCALAR > LoopyBeliefPropagation< GUM_SCALAR >::_computeProdPi_(NodeId X,
                                                                             NodeId except) {
    const auto& varX      = this->BN().variable(X);
    const auto& varExcept = this->BN().variable(except);
    auto        piXexcept = this->BN().cpt(X);
    for (const auto& U: this->BN().parents(X)) {
      if (U != except) { piXexcept *= _messages_[Arc(U, X)]; }
    }
    piXexcept = piXexcept.sumIn({&varX, &varExcept});
    return piXexcept;
  }

  template < GUM_Numeric GUM_SCALAR >
  Tensor< GUM_SCALAR > LoopyBeliefPropagation< GUM_SCALAR >::_computeProdLambda_(NodeId X) {
    Tensor< GUM_SCALAR > lamX;
    if (this->hasEvidence(X)) {
      lamX = *(this->evidence()[X]);
    } else {
      lamX.add(this->BN().variable(X));
      lamX.fill(1);
    }
    for (const auto& Y: this->BN().children(X)) {
      lamX *= _messages_[Arc(Y, X)];
    }

    return lamX;
  }

  template < GUM_Numeric GUM_SCALAR >
  Tensor< GUM_SCALAR > LoopyBeliefPropagation< GUM_SCALAR >::_computeProdLambda_(NodeId X,
                                                                                 NodeId except) {
    Tensor< GUM_SCALAR > lamXexcept;
    if (this->hasEvidence(X)) {   //
      lamXexcept = *this->evidence()[X];
    } else {
      lamXexcept.add(this->BN().variable(X));
      lamXexcept.fill(1);
    }
    for (const auto& Y: this->BN().children(X)) {
      if (Y != except) { lamXexcept *= _messages_[Arc(Y, X)]; }
    }

    return lamXexcept;
  }

  template < GUM_Numeric GUM_SCALAR >
  GUM_SCALAR LoopyBeliefPropagation< GUM_SCALAR >::_updateNodeMessage_(NodeId X) {
    auto piX  = _computeProdPi_(X);
    auto lamX = _computeProdLambda_(X);

    GUM_SCALAR KL = 0;
    Arc        argKL(0, 0);

    // update lambda_par (for arc U->x)
    for (const auto& U: this->BN().parents(X)) {
      auto newLambda = (_computeProdPi_(X, U) * lamX).sumIn({&this->BN().variable(U)});
      newLambda.normalize();
      auto ekl = static_cast< GUM_SCALAR >(0);
      try {
        ekl = _messages_[Arc(X, U)].KL(newLambda);
      } catch (InvalidArgument const&) {
        GUM_ERROR(InvalidArgument, "Not compatible pi during computation")
      } catch (FatalError const&) {   // 0 misplaced
        ekl = std::numeric_limits< GUM_SCALAR >::infinity();
      }
      if (ekl > KL) {
        KL    = ekl;
        argKL = Arc(X, U);
      }
      _messages_.set(Arc(X, U), newLambda);
    }

    // update pi_child (for arc x->child)
    for (const auto& Y: this->BN().children(X)) {
      auto newPi = (piX * _computeProdLambda_(X, Y));
      newPi.normalize();
      GUM_SCALAR ekl = KL;
      try {
        ekl = _messages_[Arc(X, Y)].KL(newPi);
      } catch (InvalidArgument const&) {
        GUM_ERROR(InvalidArgument, "Not compatible pi during computation")
      } catch (FatalError const&) {   // 0 misplaced
        ekl = std::numeric_limits< GUM_SCALAR >::infinity();
      }
      if (ekl > KL) {
        KL    = ekl;
        argKL = Arc(X, Y);
      }
      _messages_.set(Arc(X, Y), newPi);
    }

    return KL;
  }

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::_initStats_() {
    _init_messages_();
    for (const auto& node: this->BN().topologicalOrder()) {
      _updateNodeMessage_(node);
    }
  }

  /// Returns the probability of the variables.
  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::makeInference_() {
    _initStats_();
    this->initApproximationScheme();

    std::vector< NodeId > shuffleIds;
    for (const auto& node: this->BN().nodes())
      shuffleIds.push_back(node);

    auto engine = std::default_random_engine{};

    GUM_SCALAR error = 0.0;
    do {
      std::shuffle(std::begin(shuffleIds), std::end(shuffleIds), engine);
      this->updateApproximationScheme();
      for (const auto& node: shuffleIds) {
        GUM_SCALAR e = _updateNodeMessage_(node);
        if (e > error) error = e;
      }
    } while (this->continueApproximationScheme(error));
  }

  /// Returns the probability of the variable.
  template < GUM_Numeric GUM_SCALAR >
  const Tensor< GUM_SCALAR >& LoopyBeliefPropagation< GUM_SCALAR >::posterior_(NodeId id) {
    auto p = _computeProdPi_(id) * _computeProdLambda_(id);
    p.normalize();
    _posteriors_.set(id, p);

    return _posteriors_[id];
  }

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onStateChanged_() {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onEvidenceAdded_(const NodeId id,
                                                              bool         isHardEvidence) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onEvidenceErased_(const NodeId id,
                                                               bool         isHardEvidence) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onAllEvidenceErased_(bool contains_hard_evidence) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onEvidenceChanged_(const NodeId id,
                                                                bool         hasChangedSoftHard) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onModelChanged_(const GraphicalModel* bn) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::updateOutdatedTensors_() {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onMarginalTargetAdded_(const NodeId id) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onMarginalTargetErased_(const NodeId id) {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onAllMarginalTargetsAdded_() {}

  template < GUM_Numeric GUM_SCALAR >
  void LoopyBeliefPropagation< GUM_SCALAR >::onAllMarginalTargetsErased_() {}

} /* namespace gum */

#endif   // DOXYGEN_SHOULD_SKIP_THIS
