/****************************************************************************
 *   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 PRMSystem.
 *
 * @author Lionel TORTI and Pierre-Henri WUILLEMIN(_at_LIP6)
 *
 */
#include <agrum/base/multidim/aggregators/exists.h>
#include <agrum/base/multidim/aggregators/forall.h>
#include <agrum/base/multidim/aggregators/max.h>
#include <agrum/base/multidim/aggregators/min.h>
#include <agrum/PRM/elements/PRMInstance.h>
#include <agrum/PRM/elements/PRMSystem.h>

namespace gum {
  namespace prm {
    template < GUM_Numeric GUM_SCALAR >
    PRMSystem< GUM_SCALAR >::PRMSystem(std::string_view name) : PRMObject(name) {
      GUM_CONSTRUCTOR(PRMSystem);
    }

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

      for (const auto& elt: *this)
        delete elt.second;

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

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

    template < GUM_Numeric GUM_SCALAR >
    PRMSystem< GUM_SCALAR >::PRMSystem(const PRMSystem& from) :
        PRMObject(from), skeleton_(from.skeleton_), nodeIdMap_(from.nodeIdMap_),
        nameMap_(from.nameMap_), instanceMap_(from.instanceMap_), arrayMap_(from.arrayMap_) {
      GUM_CONS_CPY(PRMSystem);
    }

    template < GUM_Numeric GUM_SCALAR >
    NodeId PRMSystem< GUM_SCALAR >::add(PRMInstance< GUM_SCALAR >* i) {
      if (nameMap_.exists(i->name())) {
        GUM_ERROR(DuplicateElement,
                  "an Instance<GUM_SCALAR> with the same is already in this System");
      }

      NodeId id = skeleton_.addNode();
      nodeIdMap_.insert(id, i);
      nameMap_.insert(i->name(), i);

      auto p_inst = instanceMap_.tryGet(&(i->type()));
      if (!p_inst) {
        instanceMap_.insert(&(i->type()), new Set< PRMInstance< GUM_SCALAR >* >());
        p_inst = instanceMap_.tryGet(&(i->type()));
      }
      (*p_inst)->insert(i);

      return id;
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::groundedBN(BayesNetFactory< GUM_SCALAR >& factory) const {
      factory.startNetworkDeclaration();
      factory.addNetworkProperty("name", name());
      factory.endNetworkDeclaration();

      // Adding nodes
      for (PRMSystem< GUM_SCALAR >::const_iterator iter = begin(); iter != end(); ++iter) {
        _groundAttr_(*(iter.val()), factory);
      }

      // Adding arcs and filling CPTs
      for (PRMSystem< GUM_SCALAR >::const_iterator iter = begin(); iter != end(); ++iter) {
        _groundRef_(*(iter.val()), factory);
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::_groundAttr_(const PRMInstance< GUM_SCALAR >& instance,
                                               BayesNetFactory< GUM_SCALAR >&   factory) const {
      for (const auto node: instance.type().containerDag()) {
        // Working a Class<GUM_SCALAR> level because PRMAggregate<GUM_SCALAR>
        // are
        // instantiated as PRMAttribute<GUM_SCALAR> in an
        // PRMInstance<GUM_SCALAR>
        switch (instance.type().get(node).elt_type()) {
          case PRMClassElement< GUM_SCALAR >::prm_attribute : {
            // TODO: make a special case for noisy-or
            DiscreteVariable* var = instance.get(node).type().variable().clone();
            var->setName(
                std::format("{}.{}", instance.name(), instance.type().get(node).safeName()));
            factory.setVariable(*var);   // var is copied by the factory
            delete var;
            break;
          }

          case PRMClassElement< GUM_SCALAR >::prm_aggregate : {
            _groundAgg_(instance.type().get(node),
                        std::format("{}.{}", instance.name(), instance.type().get(node).safeName()),
                        factory);
            break;
          }

          default :
            break;
            /* Do nothing */
            ;
        }
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::_groundAgg_(const PRMClassElement< GUM_SCALAR >& elt,
                                              std::string_view                     name,
                                              BayesNetFactory< GUM_SCALAR >&       factory) const {
      factory.startVariableDeclaration();
      factory.variableName(name);

      const DiscreteVariable& agg_var = elt.type().variable();

      VarType var_type = agg_var.varType();

      factory.variableType(var_type);

      switch (var_type) {
        case VarType::LABELIZED : {
          const auto l = static_cast< const LabelizedVariable& >(agg_var);
          for (Idx i = 0; i < l.domainSize(); ++i) {
            factory.addModality(l.label(i));
          }
          break;
        }
        case VarType::INTEGER : {
          const auto l = static_cast< const IntegerVariable& >(agg_var);
          for (Idx i = 0; i < l.domainSize(); ++i) {
            factory.addModality(l.label(i));
          }
          break;
        }
        case VarType::NUMERICAL : {
          const auto l = static_cast< const NumericalDiscreteVariable& >(agg_var);
          for (Idx i = 0; i < l.domainSize(); ++i) {
            factory.addModality(l.label(i));
          }
          break;
        }
        case VarType::DISCRETIZED : {
          const auto d = static_cast< const DiscretizedVariable< double >& >(agg_var);
          for (Idx i = 0; i < d.domainSize(); ++i) {
            factory.addTick(d.tick(i));
          }
          break;
        }
        case VarType::RANGE : {
          const auto r = static_cast< const RangeVariable& >(agg_var);
          factory.addMin(r.minVal());
          factory.addMax(r.maxVal());
          break;
        }
        case VarType::CONTINUOUS : {
          GUM_ERROR(NotImplementedYet,
                    "PRM aggregator grounding does not support yet continuous variables");
        }
      }

      switch (auto& agg = static_cast< const PRMAggregate< GUM_SCALAR >& >(elt); agg.agg_type()) {
        case PRMAggregate< GUM_SCALAR >::AggregateType::MIN : {
          factory.setVariableCPTImplementation(new aggregator::Min< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::MAX : {
          factory.setVariableCPTImplementation(new aggregator::Max< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::EXISTS : {
          factory.setVariableCPTImplementation(new aggregator::Exists< GUM_SCALAR >(agg.label()));
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::FORALL : {
          factory.setVariableCPTImplementation(new aggregator::Forall< GUM_SCALAR >(agg.label()));
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::COUNT : {
          factory.setVariableCPTImplementation(new aggregator::Count< GUM_SCALAR >(agg.label()));
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::MEDIAN : {
          factory.setVariableCPTImplementation(new aggregator::Median< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::AMPLITUDE : {
          factory.setVariableCPTImplementation(new aggregator::Amplitude< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::OR : {
          factory.setVariableCPTImplementation(new aggregator::Or< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::AND : {
          factory.setVariableCPTImplementation(new aggregator::And< GUM_SCALAR >());
          break;
        }

        case PRMAggregate< GUM_SCALAR >::AggregateType::SUM : {
          factory.setVariableCPTImplementation(new aggregator::Sum< GUM_SCALAR >());
          break;
        }

        default : GUM_ERROR(OperationNotAllowed, "Aggregator not handled yet for " << agg.name())
      }

      factory.endVariableDeclaration();
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::_groundRef_(const PRMInstance< GUM_SCALAR >& instance,
                                              BayesNetFactory< GUM_SCALAR >&   factory) const {
      for (const auto& elt: instance) {
        factory.startParentsDeclaration(
            std::format("{}.{}", instance.name(), elt.second->safeName()));

        for (const auto par: instance.type().containerDag().parents(elt.second->id())) {
          switch (instance.type().get(par).elt_type()) {
            case PRMClassElement< GUM_SCALAR >::prm_aggregate :
            case PRMClassElement< GUM_SCALAR >::prm_attribute : {
              factory.addParent(
                  std::format("{}.{}", instance.name(), instance.get(par).safeName()));
              break;
            }

            case PRMClassElement< GUM_SCALAR >::prm_slotchain : {
              std::string parent_name
                  = static_cast< const PRMSlotChain< GUM_SCALAR >& >(instance.type().get(par))
                        .lastElt()
                        .safeName();

              try {
                for (const auto ref: instance.getInstances(par)) {
                  factory.addParent(std::format("{}.{}", ref->name(), parent_name));
                }
              } catch (NotFound const&) {
                // No instances for this slotchain
              }

              break;
            }

            default :
              break;
              /* nothing to do by default */
          }
        }

        factory.endParentsDeclaration();

        // Checking if we need to ground the Tensor (only for class level
        // attributes since
        // aggregates Tensors are generated)
        if (PRMClassElement< GUM_SCALAR >::isAttribute(instance.type().get(elt.second->safeName())))
          _groundTensor_(instance, *elt.second, factory);
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::_groundTensor_(const PRMInstance< GUM_SCALAR >&  instance,
                                                 const PRMAttribute< GUM_SCALAR >& attr,
                                                 BayesNetFactory< GUM_SCALAR >&    factory) const {
      Bijection< const DiscreteVariable*, const DiscreteVariable* > bijection;
      std::string var_name = std::format("{}.{}", instance.name(), attr.safeName());
      bijection.insert(&(attr.type().variable()), &(factory.variable(var_name)));

      for (const auto parent: instance.type().containerDag().parents(attr.id())) {
        switch (instance.type().get(parent).elt_type()) {
          case PRMClassElement< GUM_SCALAR >::prm_aggregate :
          case PRMClassElement< GUM_SCALAR >::prm_attribute : {
            auto parent_name
                = std::format("{}.{}", instance.name(), instance.get(parent).safeName());
            bijection.insert(&(instance.get(parent).type().variable()),
                             &(factory.variable(parent_name)));
            break;
          }

          case PRMClassElement< GUM_SCALAR >::prm_slotchain : {
            const PRMSlotChain< GUM_SCALAR >& sc
                = static_cast< const PRMSlotChain< GUM_SCALAR >& >(instance.type().get(parent));
            auto parent_name = std::format("{}.{}",
                                           instance.getInstance(sc.id()).name(),
                                           sc.lastElt().safeName());
            bijection.insert(
                &(instance.getInstance(sc.id()).get(sc.lastElt().safeName()).type().variable()),
                &(factory.variable(parent_name)));
            break;
          }

          default : {
            GUM_ERROR(FatalError, "invalid ClassElement<GUM_SCALAR> type as parent.")
            break;
          }
        }
      }

      // Copy Tensor
      // DO NOT USE MultiDimBijArray as they will wreck havok if you delete
      // the prm befor its grounded BN (happens a lot in pyAgrum)
      Tensor< GUM_SCALAR >* p = new Tensor< GUM_SCALAR >();
      for (auto var: attr.cpf().variablesSequence()) {
        p->add(*(bijection.second(var)));
      }
      Instantiation inst(attr.cpf()), jnst(*p);
      for (inst.setFirst(), jnst.setFirst(); !(inst.end() || jnst.end()); inst.inc(), jnst.inc()) {
        p->set(jnst, attr.cpf().get(inst));
      }
      GUM_ASSERT(inst.end() && jnst.end());
      factory.setVariableCPT(var_name, p, false);
    }

    template < GUM_Numeric GUM_SCALAR >
    NodeId PRMSystem< GUM_SCALAR >::add(std::string_view array, PRMInstance< GUM_SCALAR >& i) {
      return add(array, &i);
    }

    template < GUM_Numeric GUM_SCALAR >
    const DiGraph& PRMSystem< GUM_SCALAR >::skeleton() const {
      return skeleton_;
    }

    template < GUM_Numeric GUM_SCALAR >
    PRMInstance< GUM_SCALAR >& PRMSystem< GUM_SCALAR >::get(NodeId id) {
      auto p = nodeIdMap_.tryGet(id);
      if (!p) GUM_ERROR(NotFound, "found no Instance<GUM_SCALAR> matching the given id")
      return *(*p);
    }

    template < GUM_Numeric GUM_SCALAR >
    const PRMInstance< GUM_SCALAR >& PRMSystem< GUM_SCALAR >::get(NodeId id) const {
      auto p = nodeIdMap_.tryGet(id);
      if (!p) GUM_ERROR(NotFound, "found no Instance<GUM_SCALAR> matching the given id")
      return *(*p);
    }

    template < GUM_Numeric GUM_SCALAR >
    NodeId PRMSystem< GUM_SCALAR >::get(const PRMInstance< GUM_SCALAR >& i) const {
      try {
        return nodeIdMap_.keyByVal(const_cast< PRMInstance< GUM_SCALAR >* >(&i));
      } catch (NotFound const&) {
        GUM_ERROR(NotFound, "found no Instance<GUM_SCALAR> matching the given id")
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    PRMObject::prm_type PRMSystem< GUM_SCALAR >::obj_type() const {
      return PRMObject::prm_type::SYSTEM;
    }

    template < GUM_Numeric GUM_SCALAR >
    Size PRMSystem< GUM_SCALAR >::size() const {
      return nodeIdMap_.size();
    }

    template < GUM_Numeric GUM_SCALAR >
    bool PRMSystem< GUM_SCALAR >::isInstantiated(const PRMClass< GUM_SCALAR >& c) const {
      return instanceMap_.exists(const_cast< PRMClass< GUM_SCALAR >* >(&c));
    }

    template < GUM_Numeric GUM_SCALAR >
    bool PRMSystem< GUM_SCALAR >::isInstance(std::string_view name) const {
      return nameMap_.exists(name);
    }

    template < GUM_Numeric GUM_SCALAR >
    bool PRMSystem< GUM_SCALAR >::isArray(std::string_view name) const {
      return arrayMap_.exists(name);
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::instantiate() {
      for (auto iter = begin(); iter != end(); ++iter) {
        (*(iter.val())).instantiate();
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    PRMInstance< GUM_SCALAR >& PRMSystem< GUM_SCALAR >::get(std::string_view name) {
      auto p = nameMap_.tryGet(name);
      if (!p) GUM_ERROR(NotFound, "found no Instance<GUM_SCALAR> matching the given name")
      return *(*p);
    }

    template < GUM_Numeric GUM_SCALAR >
    const PRMInstance< GUM_SCALAR >& PRMSystem< GUM_SCALAR >::get(std::string_view name) const {
      auto p = nameMap_.tryGet(name);
      if (!p) GUM_ERROR(NotFound, "found no Instance<GUM_SCALAR> matching the given name")
      return *(*p);
    }

    template < GUM_Numeric GUM_SCALAR >
    const Set< PRMInstance< GUM_SCALAR >* >&

        PRMSystem< GUM_SCALAR >::get(const PRMClass< GUM_SCALAR >& type) const {
      auto p = instanceMap_.tryGet(const_cast< PRMClass< GUM_SCALAR >* >(&type));
      if (!p) GUM_ERROR(NotFound, "the given Class<GUM_SCALAR> has no instantiation in this System")
      return *(*p);
    }

    template < GUM_Numeric GUM_SCALAR >
    const Sequence< PRMInstance< GUM_SCALAR >* >&

        PRMSystem< GUM_SCALAR >::getArray(std::string_view name) const {
      auto p = arrayMap_.tryGet(name);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return *(p->second);
    }

    template < GUM_Numeric GUM_SCALAR >
    PRMClassElementContainer< GUM_SCALAR >&
        PRMSystem< GUM_SCALAR >::getArrayType(std::string_view name) {
      auto p = arrayMap_.tryGet(name);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return *(p->first);
    }

    template < GUM_Numeric GUM_SCALAR >
    const PRMClassElementContainer< GUM_SCALAR >&
        PRMSystem< GUM_SCALAR >::getArrayType(std::string_view name) const {
      auto p = arrayMap_.tryGet(name);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return *(p->first);
    }

    template < GUM_Numeric GUM_SCALAR >
    NodeId PRMSystem< GUM_SCALAR >::add(std::string_view array, PRMInstance< GUM_SCALAR >* i) {
      auto p_arr = arrayMap_.tryGet(array);
      if (!p_arr) GUM_ERROR(NotFound, "found no array matching the given name")
      if (i->type().isSubTypeOf(*(p_arr->first))) {
        NodeId id = add(i);
        p_arr->second->insert(i);
        return id;
      } else {
        GUM_ERROR(PRMTypeError,
                  "the given Instance<GUM_SCALAR> is of an incorrect "
                  "Class<GUM_SCALAR> type");
      }
    }

    template < GUM_Numeric GUM_SCALAR >
    void PRMSystem< GUM_SCALAR >::addArray(std::string_view                        array,
                                           PRMClassElementContainer< GUM_SCALAR >& type) {
      if (arrayMap_.exists(array)) {
        GUM_ERROR(DuplicateElement, "an array '" << array << "' is already in this System")
      }

      arrayMap_.insert(
          std::string{array},
          PRMSystem< GUM_SCALAR >::model_pair(&type, new Sequence< PRMInstance< GUM_SCALAR >* >()));
    }

    template < GUM_Numeric GUM_SCALAR >
    typename PRMSystem< GUM_SCALAR >::iterator PRMSystem< GUM_SCALAR >::begin() {
      return nodeIdMap_.begin();
    }

    template < GUM_Numeric GUM_SCALAR >
    const typename PRMSystem< GUM_SCALAR >::iterator& PRMSystem< GUM_SCALAR >::end() {
      return nodeIdMap_.end();
    }

    template < GUM_Numeric GUM_SCALAR >
    typename PRMSystem< GUM_SCALAR >::const_iterator PRMSystem< GUM_SCALAR >::begin() const {
      return nodeIdMap_.begin();
    }

    template < GUM_Numeric GUM_SCALAR >
    const typename PRMSystem< GUM_SCALAR >::const_iterator& PRMSystem< GUM_SCALAR >::end() const {
      return nodeIdMap_.end();
    }

    template < GUM_Numeric GUM_SCALAR >
    typename PRMSystem< GUM_SCALAR >::array_iterator
        PRMSystem< GUM_SCALAR >::begin(std::string_view a) {
      auto p = arrayMap_.tryGet(a);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return p->second->begin();
    }

    template < GUM_Numeric GUM_SCALAR >
    const typename PRMSystem< GUM_SCALAR >::array_iterator&
        PRMSystem< GUM_SCALAR >::end(std::string_view a) {
      auto p = arrayMap_.tryGet(a);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return p->second->end();
    }

    template < GUM_Numeric GUM_SCALAR >
    typename PRMSystem< GUM_SCALAR >::const_array_iterator
        PRMSystem< GUM_SCALAR >::begin(std::string_view a) const {
      auto p = arrayMap_.tryGet(a);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return p->second->begin();
    }

    template < GUM_Numeric GUM_SCALAR >
    const typename PRMSystem< GUM_SCALAR >::const_array_iterator&
        PRMSystem< GUM_SCALAR >::end(std::string_view a) const {
      auto p = arrayMap_.tryGet(a);
      if (!p) GUM_ERROR(NotFound, "found no array matching the given name")
      return p->second->end();
    }

    template < GUM_Numeric GUM_SCALAR >
    bool PRMSystem< GUM_SCALAR >::exists(std::string_view name) const {
      return nameMap_.exists(name) || arrayMap_.exists(name);
    }
  } /* namespace prm */
} /* namespace gum */
