//============================================================================
//
// This file is part of the Shapes Generative Model Library (SGML) project.
//
// Copyright (c) 2010-2014 - Evangelos Kalogerakis, Siddhartha Chaudhuri (authors of the code) / Stanford University, UMass-Amherst
//
// SGML is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// SGML is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with SGML.  If not, see <http://www.gnu.org/licenses/>.
//
//============================================================================

#ifndef __S12_SHAPEBNET_NODE_HPP
#define __S12_SHAPEBNET_NODE_HPP

#include "../../Node.hpp"
#include "S12_Common.hpp"
#include "S12_KMeans.hpp"
#include "S12_MatrixOps.hpp"
#include "S12_SigmoidCPD.hpp"
#include "S12_Gamma.hpp"

namespace SNL {
namespace S12 {

class RootNode;
class LayoutNode;
class ExistsNode;
class CardinalityNode;
class AdjacencyNumPartsNode;
class GaussianNode;
class SigmoidCPD;
class Gaussian;
//class Factor;
class Node;


// A multi-variate gaussian distribution.
class SNL_S12_API Gaussian
{
public:
  SNL_VECTOR(double)    b_0;  // used for continuous cond. gaussian case
  SNL_MATRIX(double)    b;    // used for continuous cond. gaussian case
  SNL_MATRIX(double)    covariance; // used for both cases
  SNL_MATRIX(double)    bak_covariance; // used for both cases (backup covariance, when the node is dependent)
  SNL_MATRIX(double)    inverse_covariance; // used for both cases
  double                determinant_covariance; // used for both cases
  SNL_MATRIX(double)    covariance_xx; // used for continuous cond. gaussian case

  SNL_VECTOR(double)    mean;
  SNL_VECTOR(double)    bak_mean;
  size_t                rank;
  size_t                num_data_points;

  struct Canonical
  {
    SNL_MATRIX(double)  K;
    SNL_VECTOR(double)  h;
    double              g;
  };
  Canonical canonical;

  Gaussian(): determinant_covariance(0), rank(0), num_data_points(0)  {};

  inline size_t num_params() const
  {
    if ( !b_0.empty() )
    {
      return b_0.size() + numNonZeroEntries( b ) + numUniqueNonZeroEntries( covariance );
    }
    else
    {
      return mean.size() + numUniqueNonZeroEntries( covariance );
    }
  }

  void initialize(size_t _d);
  double evaluateGaussian(SNL_VECTOR(double) x);
  double evaluateCondGaussian(SNL_VECTOR(double) y, SNL_VECTOR(double) x);
  void getCanonicalForm();
	double getRoughProbability(double, double, size_t num_intervals = 5);


}; // class Gaussian



// Basic Bayes Net node.
class SNL_S12_API Node : public SNL::Node
{
public:
  //Factor               factor;          // used for inference
  SNL_VECTOR(double)   observation;
  string               name;
  size_t               domain_size;
  SNL_VECTOR(double)   cpt;
  SNL_VECTOR(double)   MS;              // expected counts
  SNL_VECTOR(size_t)   values;

  Node(int type_ = -1, long id_ = -1, string _name = "") : SNL::Node(type_, id_), name(_name), domain_size(0) {};

  int getType() const { return type; }
  long getID() const { return id; }
  string getName() const { ostringstream oss; oss << "Node (" << (char)type << '-' << id << ") " << name; return oss.str(); }
  size_t getObservedDomainSize() const { return observation.empty()? domain_size: 1; }
  inline size_t getValuePos(size_t & value) const
  {
    for (size_t i = 0; i < values.size(); i++)
      if (values[i] == value)
        return i;
     return values.size();
  }


  void initialize();
  void clearMS();
  size_t ac2(size_t d, size_t p);                       // access 2D-cpt
  size_t acn(size_t d, SNL_VECTOR(size_t) x);               // access cpt
  size_t getValues(size_t offset, SNL_VECTOR(size_t) & x);  // get cpt coordinates given an offset value
  //void getFactor();
  void printcpt(std::ostream &os = std::cout);

}; // class Node


class SNL_S12_API RootNode : public Node
{
public:
  float loglik;
  float cheesemanstutz;
  int num_indep_parameters;
  int empty_clusters;
  bool istemp;
  SNL_VECTOR( size_t )  add_layout_domain_size;  // initialized in init

  SNL_MATRIX(double)                   Pr;                          // num samples x root domain size
  SNL_MATRIX(double)                   old_Pr;                      // num samples x root domain size
  SNL_VECTOR( SNL_VECTOR( Mesh * )  )  most_likely_meshes;          // most likely meshes for each value of root node
  SNL_VECTOR( SNL_VECTOR( Mesh * )  )  most_likely_initial_meshes;  // most likely meshes for each value of root node before EM
  SNL_VECTOR( size_t )                 style_id;

  RootNode() : Node('R', 0, "root_node"), loglik(-FLT_MAX), cheesemanstutz(FLT_MAX), num_indep_parameters(-1), empty_clusters(-1), istemp(false)
  {
    values = increasingNumberVector((size_t)0, domain_size);
  }
  inline size_t num_param() const
  {
    return domain_size - 1;
  }
  void initialize();
  void clearMS();
  void learn(double);
  bool isDiscrete() const { return true; }
  size_t dim() const { return 1; }
  //void getFactor();
  size_t getValuePos(size_t & value) const;
}; // class RootNode


class SNL_S12_API LayoutNode : public Node
{
public:
  RootNode              *                  parent;
  Label                 *                  label;
  SNL_MATRIX(double)                       Pl;                          // num samples x domain size
  SNL_VECTOR( SNL_MATRIX(double) )         Plr;                         // num samples x domain size x root domain size
  SNL_MATRIX(double)                       old_Pl;                      // num samples x domain size
  SNL_VECTOR( SNL_MATRIX(double) )         old_Plr;                     // num samples x domain size x root domain size
  SNL_VECTOR( SNL_VECTOR( Mesh * )  )      most_likely_meshes;          // most likely meshes for each value of the layout node
  SNL_VECTOR( SNL_VECTOR( Mesh * )  )      most_likely_initial_meshes;  // most likely meshes for each value of the layout node

  LayoutNode(Label * _label, RootNode * r, long _id) : Node('L', _id, "layout_node_" + _label->name), parent(r), label(_label)
  {
    addParent(parent);
    values = increasingNumberVector((size_t)0, domain_size);
  }
  inline size_t num_param() const
  {
    return (domain_size - 1) * parent->domain_size;
  }

  void initialize();
  void clearMS();
  void learn(double);
  bool isDiscrete() const { return true; }
  size_t dim() const { return 1; }
  //void getFactor();
  size_t getValuePos(size_t & value) const;

}; // class LayoutNode


class SNL_S12_API ExistsNode : public Node
{
public:
  LayoutNode  *  parent;

  ExistsNode(LayoutNode * l): Node('E', -1, "internal_exists_node_" + l->label->name)
  {
    domain_size = 2;
    parent = l;
    addParent(l);
    values = increasingNumberVector((size_t)0, domain_size);
  }

  void initialize();
  void clearMS();
  void learn();
  bool isDiscrete() const { return true; }
  size_t dim() const { return 1; }
  //void getFactor();
  size_t getValuePos(size_t & value) const;

}; // class ExistsNode


class SNL_S12_API CardinalityNode : public Node
{
public:
  RootNode                    *      parent;
  Label                       *      label;
  ExistsNode                  *      e;
  SNL_VECTOR( CardinalityNode *)     parents_CardinalityNode;
  SNL_VECTOR(SigmoidCPD)             sigmoids;

  CardinalityNode(Label * _label, RootNode * r, LayoutNode * l, long _id) : Node('C', _id, "cardinality_node_" + _label->name), parent(r), label(_label)
  {
    e = new ExistsNode(l);
    parents.clear();
    parents.push_back( parent );
    parents.push_back( e );
  }
  ~CardinalityNode()
  {
    delete e;
  }

  void registerValue(size_t & value)
  {
    for (size_t i = 0; i < values.size(); i++)
      if (values[i] == value)
        return;

      values.push_back(value);
    domain_size = values.size();
  }

  void registerParents()
  {
    sigmoids.clear();

    clearParents();
    addParent(parent);
    addParent(e);

    for (size_t i = 0; i < parents_CardinalityNode.size(); ++i)
      addParent(parents_CardinalityNode[i]);
  }

  size_t num_param() const
  {
    size_t np = 0;

    for (size_t d2 = 0; d2 < sigmoids.size(); d2++)
      np += sigmoids[d2].weights.numRows();

    return np;
  }

  void initialize();
  void clearMS();
  void learn(const SNL_VECTOR( SNL_VECTOR(size_t) ) &, double, bool use_cpt = false);
  bool isDiscrete() const { return true; }
  size_t dim() const { return 1; }
  //void getFactor();
  size_t getValuePos(size_t & value) const;

}; // class CardinalityNode


class SNL_S12_API AdjacencyNumPartsNode : public Node
{
public:
  LayoutNode                  *      parent;
  SNL_VECTOR(SigmoidCPD)             sigmoids;
  Label                       *      label;
  size_t                             adj_label_pos;
  Label                       *      adj_label;
  SNL_VECTOR( CardinalityNode *)     parents_CardinalityNode;

  AdjacencyNumPartsNode(Label * _label, Label * _adj_label, size_t _adj_label_pos, LayoutNode * l, long _id)
  : Node('A', _id, "adjacency_node_" + _label->name + "-" + _adj_label->name), parent(l), label(_label), adj_label_pos(_adj_label_pos), adj_label(_adj_label)
  {
    addParent(parent);
  }

  inline void registerValue(size_t & value)
  {
    for (size_t i = 0; i < values.size(); i++)
      if (values[i] == value)
        return;

      values.push_back(value);
    domain_size = values.size();
  }

  inline void registerParents()
  {
    sigmoids.clear();

    clearParents();
    addParent(parent);

    for (size_t i = 0; i < parents_CardinalityNode.size(); ++i)
      addParent(parents_CardinalityNode[i]);
  }

  inline size_t num_param() const
  {
    size_t np = 0;

    for (size_t d2 = 0; d2 < sigmoids.size(); d2++)
      np += sigmoids[d2].weights.numRows();

    return np;
  }

  void initialize();
  void clearMS();
  void learn(const SNL_VECTOR( SNL_VECTOR( SNL_VECTOR(size_t) ) ) &, const SNL_VECTOR( SNL_VECTOR(size_t) ) &, double, bool use_cpt = false);
  bool isDiscrete() const { return true; }
  size_t dim() const { return 1; }
  //void getFactor();
  size_t getValuePos(size_t & value) const;

}; // class AdjacencyNumPartsNode


class SNL_S12_API GaussianNode: public Node
{
private:
  void reduceCanonicalForm( Gaussian::Canonical&, SNL_VECTOR(GaussianNode*)& );
public:
  LayoutNode               *    layout_node;
  SNL_VECTOR(Gaussian)          gaussians;
  SNL_VECTOR( GaussianNode *)   parents_gaussians;

  GaussianNode(LayoutNode * _layout_node, long _id) : Node('G', _id, "gaussian_node_" + _layout_node->label->name), layout_node(_layout_node)
  { }

  inline size_t num_param() const
  {
    size_t np = 0;

    for (size_t d2 = 0; d2 < gaussians.size(); d2++)
    {
      np += gaussians[d2].num_params();
    }

    return np;
  }

  inline size_t num_param(size_t d2) const
  {
    return gaussians[d2].num_params();
  }

  inline void registerParents()
  {
    clearParents();
    addParent(layout_node);

    for (size_t i = 0; i < parents_gaussians.size(); ++i)
      addParent(parents_gaussians[i]);
  }

  bool learn(const SNL_VECTOR( SNL_VECTOR( SNL_VECTOR(double) ) ) &, double, double);
  bool isDiscrete() const { return false; }
  size_t dim() const
  {
    return gaussians.size() <= 1 ? 0 : max( gaussians[1].mean.size(), gaussians[1].b_0.size() );
  }
  size_t getObservedDomainSize() const { return observation.empty()? dim(): 0; }
  //void getFactor();
  friend ostream& operator<<(std::ostream& os, const GaussianNode& gn);

}; // class GaussianNode

} // namespace S12
} // namespace SNL

#endif
