//============================================================================
//
// 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_HPP
#define __S12_SHAPEBNET_HPP

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


namespace SNL {
	namespace S12 {


class SNL_S12_API ShapeBNet;



		class SNL_S12_API TreeNode
		{
		private:
			double  loglik_value;
			double  adj_penalty;
		public:
			char        ntype; // 'R', 'L', 'N', 'A', 'B', 'G', '-'
			size_t      nid;   // node id
			size_t      value; // value of node
			vector<int> partial_assignment;	

			TreeNode(): ntype('-'), nid(0), value(0), loglik_value(0.0), adj_penalty(0.0) {}
			TreeNode( TreeNode& prev, double prob, char _ntype, size_t _nid, size_t _value ): ntype(_ntype), nid(_nid), value(_value)
			{
				partial_assignment = prev.partial_assignment;
				partial_assignment.push_back( _value );
				loglik_value = prev.loglik_value + log( min(prob + SUPEREPSILON, 1.0) );
				adj_penalty = prev.adj_penalty;
			}
			TreeNode( TreeNode& prev, double prob, char _ntype, size_t _nid, size_t _value, double _adj_penalty ): ntype(_ntype), nid(_nid), value(_value)
			{
				partial_assignment = prev.partial_assignment;
				partial_assignment.push_back( _value );
				loglik_value = prev.loglik_value + log( min(prob + SUPEREPSILON, 1.0) );
				adj_penalty = prev.adj_penalty + _adj_penalty;
			}
			double getLogLikValue() 
			{ 
				return loglik_value;
			}
			double getAdjPenalty() 
			{ 
				return adj_penalty;
			}
			bool operator < (const TreeNode& t) const 
			{
				return loglik_value < t.loglik_value;
			}
		};


		class SNL_S12_API ShapeBNetSample
		{
		public:
			double											prob;
			ShapeBNet*									bnet;	
			vector<size_t>							cardinality_nodes_value;
			vector< vector< double > >	scales;
			size_t											root_node_value;        
			vector< int >							  use_mesh_ids;           // used only in deterministic setting

			vector<size_t>							layout_nodes_value;     // used only in random sampling	
			vector<size_t>							adjacencynumslots_nodes_value; // used only in random sampling	
			vector<size_t>							adjacencynumparts_nodes_value; // used only in random sampling	
			vector< vector< double > >  parts_features;                // used only in random sampling	
			float												distToClosestMesh, distToClosestSample;
		};

		class SNL_S12_API ShapeBNet : public SNL::Network
		{
		public:
			SNL_VECTOR(Label *) labels;
			SNL_VECTOR(MeshPart *) parts;
			SNL_VECTOR(Mesh *) meshes;
			SNL_VECTOR( SNL_VECTOR( SNL_VECTOR(double) ) )   features_training_data;       // M meshes x L virtual labels x F features
			SNL_VECTOR( SNL_VECTOR( SNL_VECTOR(size_t) ) )   adjacentparts_training_data;  // M meshes x L virtual labels x L' adj labels x 1
			SNL_VECTOR( SNL_VECTOR( size_t  ) )              cardinality_training_data;    // M meshes x L virtual labels x 1 features
			SNL_VECTOR( SNL_VECTOR( size_t  ) )              exists_training_data;         // M meshes x L virtual labels x 1 features
			string                                           name;

			double internal_map_multiplier;
			map< pair<char, size_t>, size_t >		treemap;
			vector<size_t>	style_id; // same size as meshes


			void init(SNL_VECTOR(SNL::Node *)* _nodes );  // init used for training
			void init(size_t root_domain_size, SNL_VECTOR( size_t ) & add_layout_domain_size, bool prepare_for_training = false);    // init used for training
			void internal_clean();
			void internal_clean( SNL_VECTOR(SNL::Node *)* );
			void learn(size_t, bool full_structure_search = false);
			void EM(size_t);
			void CheesemanStutzScore();
			SNL_VECTOR(float) E();
			void M();
			inline string getName() { return name; }
			void doAnalyticSpecificInference();
			MeshPart* getMeshRefPart( size_t m, size_t l );


		public:
			SNL::Options                           options;
			RootNode         *                     root_node;
			SNL_VECTOR( LayoutNode *)              layout_nodes;
			SNL_VECTOR( CardinalityNode * )        cardinality_nodes, ordered_cardinality_nodes;
			SNL_VECTOR( AdjacencyNumPartsNode * )  adjacencynumparts_nodes;
			SNL_VECTOR( GaussianNode * )           gaussian_nodes, ordered_gaussian_nodes;


			ShapeBNet(const string& _name = "unnamed", const SNL::Options &_options = SNL::Options()): name(_name)
			{
				options.set("deactivate_adjacencies", 0);
				options.set("use_cpts_everywhere", 0);
				options.set("dirichlet_hyperparameter_alpha", 0.1); // hyperparameter alpha
				options.set("wishart_hyperparameter_beta", 1e-3); // hyperparameter alpha
				options.set("wishart_hyperparameter_nu", 1); // hyperparameter alpha
				options.set("num_max_em_iterations", 30);
				options.set("learn_tree_only", 0);
				options.set("min_root_domain_size", 2);
				options.set("max_root_domain_size", 2);
				options.set("min_add_layout_domain_size", 0);
				options.set("max_add_layout_domain_size", 5);
				options.set("kmeans_iterations", 100);
				options.set("kmeans_replicates", 1000);
				options.set("em_initialization_randomness", 0.01);
				options.set("consider_empty_clusters", 1);
				options.set("max_parents_lateral", 3);
				options.set("max_seconds_for_inference", 0.15);
				options.set("minimum_iterations_for_inference", 3.0);
				options.merge(_options);
				root_node = new RootNode();
				root_node->istemp = true;
				nodes = NULL;
			}
			~ShapeBNet();
			void setRootDomainSize(size_t root_domain_size) { root_node->domain_size = root_domain_size; }
			void setLayoutNodeDomainSize(std::vector<size_t>& add_layout_domain_size) {root_node->add_layout_domain_size = add_layout_domain_size; }
			void init(SNL::Options const & _options = SNL::Options() );
			bool topologicalSort(); // returns false if topolofical sort finds cycles.
			vector<ShapeBNetSample> getAllInstances(char *, int limit_output_samples = INT_MAX);
			void setTrainingData( SNL_VECTOR(Label *) & labels, SNL_VECTOR(MeshPart *) & parts, SNL_VECTOR(Mesh *) & meshes );
			void train(SNL::Options const & _options = SNL::Options() );
			void extend(SNL::Options const & options) {}
			void load(std::string const & in_path);
			void save(std::string const & out_path) const;
			void printStats(std::ostream & os = std::cout);

			void setObservation(SNL::Node* node, const std::vector<double> & value);
			void setObservation(SNL::Node* node, const double value);
			void clearObservation(SNL::Node* node);
			void clearAllObservations();
			void doInference();
		}; // class ShapeBNet

	} // namespace S12
} // namespace SNL

#endif
