/****************************************************************************
 *   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 Inline implementation of StructuredBayesBall.
 *
 * @author Lionel TORTI and Pierre-Henri WUILLEMIN(_at_LIP6)
 */

#include <agrum/PRM/inference/structuredBayesBall.h>

namespace gum {
  namespace prm {

    template < GUM_Numeric GUM_SCALAR >
    StructuredBayesBall< GUM_SCALAR >::~StructuredBayesBall() {
      GUM_DESTRUCTOR(StructuredBayesBall);

      for (const auto& elt: _reqMap_)
        delete elt.second.first;
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::_clean_() {
      for (const auto& elt: _reqMap_)
        delete elt.second.first;

      _keyMap_.clear();
      _reqMap_.clear();
    }

    template < GUM_Numeric GUM_SCALAR >
    bool StructuredBayesBall< GUM_SCALAR >::_isHardEvidence_(const PRMInstance< GUM_SCALAR >* i,
                                                             NodeId                           n) {
      if (!i->exists(n)) return false;
      typename PRMInference< GUM_SCALAR >::Chain chain = std::make_pair(i, &(i->get(n)));

      if (_inf_->hasEvidence(chain)) {
        const Tensor< GUM_SCALAR >* e = _inf_->evidence(i)[n];
        Instantiation               inst(e);
        Size                        count = 0;

        for (inst.setFirst(); !inst.end(); inst.inc()) {
          if ((e->get(inst) == (GUM_SCALAR)1.0)) ++count;
          else if (e->get(inst) != (GUM_SCALAR)0.0) return false;
        }

        return (count == 1);
      }

      return false;
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::_compute_(const PRMInstance< GUM_SCALAR >* i,
                                                      NodeId                           n) {
      _clean_();
      /// Key = instance.PRMClassElement<GUM_DATA>
      /// pair = <upper mark, lower mark>
      StructuredBayesBall< GUM_SCALAR >::InstanceMap marks;
      _fromChild_(i, n, marks);
      _fillMaps_(marks);

      for (const auto& elt: marks)
        delete elt.second;
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::_fromChild_(const PRMInstance< GUM_SCALAR >* i,
                                                        NodeId                           n,
                                                        InstanceMap&                     marks) {
      auto p_marks = marks.tryGet(i);
      if (!p_marks) {
        marks.insert(i, new StructuredBayesBall< GUM_SCALAR >::MarkMap());
        p_marks = marks.tryGet(i);
      }

      if (!(*p_marks)->exists(n)) { (*p_marks)->insert(n, std::pair< bool, bool >(false, false)); }

      // Sending message to parents
      switch (i->type().get(n).elt_type()) {
        case PRMClassElement< GUM_SCALAR >::prm_slotchain : {
          if (!_getMark_(marks, i, n).first) {
            _getMark_(marks, i, n).first = true;

            for (const auto inst: i->getInstances(n))
              _fromChild_(inst, inst->get(_getSC_(i, n).lastElt().safeName()).id(), marks);
          }

          if (!_getMark_(marks, i, n).second) {
            _getMark_(marks, i, n).second = true;

            for (const auto chi: i->type().containerDag().children(n))
              _fromParent_(i, chi, marks);
          }

          break;
        }

        case PRMClassElement< GUM_SCALAR >::prm_aggregate :
        case PRMClassElement< GUM_SCALAR >::prm_attribute : {
          if (!_getMark_(marks, i, n).first) {
            _getMark_(marks, i, n).first = true;

            if (!_isHardEvidence_(i, n))
              for (const auto par: i->type().containerDag().parents(n))
                _fromChild_(i, par, marks);
          }

          if (!_getMark_(marks, i, n).second) {
            _getMark_(marks, i, n).second = true;

            // In i.
            for (const auto chi: i->type().containerDag().children(n))
              _fromParent_(i, chi, marks);

            // Out of i.
            if (i->hasRefAttr(n)) {
              const auto& refs = i->getRefAttr(n);

              for (auto iter = refs.begin(); iter != refs.end(); ++iter)
                _fromParent_(iter->first, iter->first->type().get(iter->second).id(), marks);
            }
          }

          break;
        }

        default : {
          // We shouldn't reach any other PRMClassElement<GUM_DATA> than
          // PRMAttribute
          // or
          // PRMSlotChain<GUM_SCALAR>.
          GUM_ERROR(FatalError, "This case is impossible.")
        }
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::_fromParent_(const PRMInstance< GUM_SCALAR >* i,
                                                         NodeId                           n,
                                                         InstanceMap&                     marks) {
      auto p_marks = marks.tryGet(i);
      if (!p_marks) {
        marks.insert(i, new StructuredBayesBall< GUM_SCALAR >::MarkMap());
        p_marks = marks.tryGet(i);
      }

      if (!(*p_marks)->exists(n)) { (*p_marks)->insert(n, std::pair< bool, bool >(false, false)); }

      // Concerns only PRMAttribute (because of the hard evidence)
      if ((_isHardEvidence_(i, n)) && (!_getMark_(marks, i, n).first)) {
        _getMark_(marks, i, n).first = true;

        for (const auto par: i->type().containerDag().parents(n))
          _fromChild_(i, par, marks);
      } else if (!_getMark_(marks, i, n).second) {
        _getMark_(marks, i, n).second = true;

        // In i.
        for (const auto chi: i->type().containerDag().children(n))
          _fromParent_(i, chi, marks);

        // Out of i.
        if (i->hasRefAttr(n)) {
          for (auto iter = i->getRefAttr(n).begin(); iter != i->getRefAttr(n).end(); ++iter)
            _fromParent_(iter->first, iter->first->type().get(iter->second).id(), marks);
        }
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::_fillMaps_(InstanceMap& marks) {
      // First find for each instance it's requisite nodes
      HashTable< const PRMInstance< GUM_SCALAR >*, Set< NodeId >* > req_map;

      for (const auto& elt: marks) {
        Set< NodeId >* req_set = new Set< NodeId >();

        for (const auto& elt2: *elt.second)
          if (elt2.second.first) req_set->insert(elt2.first);

        req_map.insert(elt.first, req_set);
      }

      // Remove all instances with 0 requisite nodes
      Set< const PRMInstance< GUM_SCALAR >* > to_remove;

      for (const auto& elt: req_map)
        if (elt.second->size() == 0) to_remove.insert(elt.first);

      for (const auto remo: to_remove) {
        delete req_map[remo];
        req_map.erase(remo);
      }

      // Fill  _reqMap_ and  _keyMap_
      for (const auto& elt: req_map) {
        std::string key = _buildHashKey_(elt.first, *elt.second);

        if (auto p_req = _reqMap_.tryGet(key)) {
          _keyMap_.insert(elt.first, std::pair< std::string, Set< NodeId >* >(key, p_req->first));
          p_req->second += 1;
          delete elt.second;
          req_map[elt.first] = 0;
        } else {
          _reqMap_.insert(key, std::pair< Set< NodeId >*, Size >(elt.second, 1));
          _keyMap_.insert(elt.first, std::pair< std::string, Set< NodeId >* >(key, elt.second));
        }
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    std::string
        StructuredBayesBall< GUM_SCALAR >::_buildHashKey_(const PRMInstance< GUM_SCALAR >* i,
                                                          Set< NodeId >& req_nodes) {
      std::string result(i->type().name());

      for (const auto node: i->type().containerDag().nodes())
        if (req_nodes.exists(node)) result += std::format("-{}", node);

      return result;
    }

    template < GUM_Numeric GUM_SCALAR >
    StructuredBayesBall< GUM_SCALAR >::StructuredBayesBall(
        const PRMInference< GUM_SCALAR >& inference) : _inf_(&inference) {
      GUM_CONSTRUCTOR(StructuredBayesBall);
    }

    template < GUM_Numeric GUM_SCALAR >
    StructuredBayesBall< GUM_SCALAR >::StructuredBayesBall(
        const StructuredBayesBall< GUM_SCALAR >& source) : _inf_(0) {
      GUM_CONS_CPY(StructuredBayesBall);
      GUM_ERROR(FatalError, "Not allowed.")
    }

    template < GUM_Numeric GUM_SCALAR >
    StructuredBayesBall< GUM_SCALAR >& StructuredBayesBall< GUM_SCALAR >::operator=(
        const StructuredBayesBall< GUM_SCALAR >& source) {
      GUM_ERROR(FatalError, "Not allowed.")
    }

    template < GUM_Numeric GUM_SCALAR >
    const std::string&
        StructuredBayesBall< GUM_SCALAR >::key(const PRMInstance< GUM_SCALAR >* i) const {
      return _keyMap_[i].first;
    }

    template < GUM_Numeric GUM_SCALAR >
    const std::string&
        StructuredBayesBall< GUM_SCALAR >::key(const PRMInstance< GUM_SCALAR >& i) const {
      return _keyMap_[&i].first;
    }

    template < GUM_Numeric GUM_SCALAR >
    const Set< NodeId >& StructuredBayesBall< GUM_SCALAR >::requisiteNodes(
        const PRMInstance< GUM_SCALAR >* i) const {
      return *(_keyMap_[i].second);
    }

    template < GUM_Numeric GUM_SCALAR >
    const Set< NodeId >& StructuredBayesBall< GUM_SCALAR >::requisiteNodes(
        const PRMInstance< GUM_SCALAR >& i) const {
      return *(_keyMap_[&i].second);
    }

    template < GUM_Numeric GUM_SCALAR >
    Size StructuredBayesBall< GUM_SCALAR >::occurrence(std::string_view key) const {
      return _reqMap_[std::string{key}].second;
    }

    template < GUM_Numeric GUM_SCALAR >
    float StructuredBayesBall< GUM_SCALAR >::liftRatio() const {
      return ((float)_reqMap_.size()) / ((float)_keyMap_.size());
    }

    template < GUM_Numeric GUM_SCALAR >
    bool StructuredBayesBall< GUM_SCALAR >::exists(const PRMInstance< GUM_SCALAR >* i) const {
      return _keyMap_.exists(i);
    }

    template < GUM_Numeric GUM_SCALAR >
    bool StructuredBayesBall< GUM_SCALAR >::exists(const PRMInstance< GUM_SCALAR >& i) const {
      return _keyMap_.exists(&i);
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::compute(const PRMInstance< GUM_SCALAR >* i, NodeId n) {
      _compute_(i, n);
    }

    template < GUM_Numeric GUM_SCALAR >
    void StructuredBayesBall< GUM_SCALAR >::compute(const PRMInstance< GUM_SCALAR >& i, NodeId n) {
      _compute_(&i, n);
    }

    template < GUM_Numeric GUM_SCALAR >
    const PRMSlotChain< GUM_SCALAR >&
        StructuredBayesBall< GUM_SCALAR >::_getSC_(const PRMInstance< GUM_SCALAR >* i, NodeId n) {
      return static_cast< const PRMSlotChain< GUM_SCALAR >& >(i->type().get(n));
    }

    template < GUM_Numeric GUM_SCALAR >
    std::pair< bool, bool >&
        StructuredBayesBall< GUM_SCALAR >::_getMark_(InstanceMap&                     marks,
                                                     const PRMInstance< GUM_SCALAR >* i,
                                                     NodeId                           n) {
      return (*(marks[i]))[n];
    }

  } /* namespace prm */
} /* namespace gum */
