/*
 * This file is part of a parser for an extension of the PRISM language.
 *
 * This 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.
 *
 * The parser 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 the program this parser part of.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 * Copyright 2007-2010 Bjoern Wachter (Bjoern.Wachter@comlab.ox.ac.uk)
 * Copyright 2009-2010 Ernst Moritz Hahn (emh@cs.uni-sb.de)
 */

#ifndef Model_H
#define Model_H

#include <tr1/unordered_map>
#include <tr1/unordered_set>
#include <cvc3/vcl.h>

#include "Node.h"
#include <boost/shared_ptr.hpp>
#include "Property.h"
#include "SymbolTable.h"

namespace prismparser {
  extern CVC3::VCL vc;

  //  class Property;
  class Module;
  class Command;
  class Alternative;

  /*! container for set of Alternatives */
  class Alternatives : public std::vector < boost::shared_ptr < Alternative > >{};
  /*! container for set of Commands */
  class Commands : public std::vector < boost::shared_ptr < Command > >{};
  /*! container for set of synchronized Commands */
  class SynchronizedCommands : public std::map < std::string, Commands >{};
  /*! container for set of properties */
  class Properties : public std::vector < boost::shared_ptr < Property > > {};
  /*! \typedef action label */
  typedef std::string Action;
  /*! \typedef container for set of action labels */
  typedef std::set < std::string > Actions;
  /*! \typedef container for set of modules */
  typedef std::vector < boost::shared_ptr < Module > >Modules;

  typedef std::tr1::unordered_map < std::string, CVC3::Expr> Variables;

  enum ModelType {
    DTMC, MDP, CTMC, CTMDP
  };

  /*! \brief a model

    The getter functions are the most interesting part for
    people who want to use the class.
    The other functions are mainly used for model construction
    which is done automatically by the parser.
  */
  struct Model : Node {
    /*************************/
    /* The getter functions  */
    /*************************/
    /*! \brief Get module by name
      \return pointer to module if found, NULL otherwise
    */
    Module *getModule (const std::string &) const;

    const CVC3::Expr & getInitial() const;
    const CVC3::Expr &getDefaultInitialValue (const CVC3::Expr &) const;
    const Commands &getCommands() const;
    const std::vector < CVC3::Expr > &getUserPreds() const;
    const Properties &getProperties() const;
    const std::vector < CVC3::Expr > &getInvar() const;
    unsigned getMaxNrOfAlt() const;
    bool isInputVariable(const CVC3::Expr &) const;
    bool isParameterVariable(const CVC3::Expr &) const;
    ModelType getModelType() const;

    /*****************************************************************/
    /* The following part is only interesting for model construction */
    /*****************************************************************/

    void setDefaultInitialValue(const CVC3::Expr &, const CVC3::Expr &);
    void setModelType (ModelType);
    void flatten();
    void fixDeadlocks();
    void addModule(Module *);
    void addAction(const Action &);
    void setInitial(const CVC3::Expr &);
    void addVariable(const std::string &, const CVC3::Expr &, bool input =
		     false);
    void addVariable(const std::string & id, const CVC3::Expr & var,
		     const CVC3::Expr & lower, const CVC3::Expr & upper,
		     bool input = false);
    void addPredicate(const CVC3::Expr &);
    void addProperty(Property *);
    void addInvariant(const CVC3::Expr &);
    void addStateReward(const unsigned, const CVC3::Expr &, const CVC3::Expr &);
    std::vector < std::pair < CVC3::Expr, CVC3::Expr > >&getStateRewards();
    std::vector < std::pair < CVC3::Expr, CVC3::Expr > >&getStateRewards(const unsigned);
    void clearStateRewards();
    void clearTransRewards();
    void addTransReward(const unsigned, const CVC3::Expr &, const CVC3::Expr &);
    void addTransReward(const unsigned, const Action &, const CVC3::Expr &, const CVC3::Expr &);
    std::vector < std::pair < std::pair < Action, CVC3::Expr >, CVC3::Expr > >
      &getTransRewards();
    std::vector < std::pair < std::pair < Action, CVC3::Expr >, CVC3::Expr > >
      &getTransRewards(const unsigned);

    Model();
    ~Model();

    /*! \brief symbol table of all variables */
    struct hash {
      size_t operator()(CVC3::Expr x) const {
	return x.hash();
      }
    };
    std::tr1::unordered_map<std::string, CVC3::Expr > variables;
    std::tr1::unordered_set<CVC3::Expr,hash> input_variables;
    std::tr1::unordered_set<CVC3::Expr,hash> parameter_variables;

    /* Node functions */
    virtual std::string toString() const;

    Model(const Model &);
    Model &operator=(const Model &);

    /*! brief check if the model only uses constant rates */
    bool usesOnlyConstantRates();

    /**     \brief compute implicit default initial values of variables
	    \note PRISM language implicitly assumes that the lower bound of the variable range is initial value
	    if no explicit initial expression is given
	    \post the initial condition is initialized to default expression
    */
    void computeDefaultInitialValue();

    bool hasBounds(const CVC3::Expr &) const;
    std::pair < CVC3::Expr, CVC3::Expr > getBounds(const CVC3::Expr &);
  private:
    ModelType model_type;	//! DTMC, DTMDP, CTMC, CMTDP,...
    CVC3::Expr init;		//! initial states
    Commands guarded_transitions;	//! guarded transitions
    Properties properties;	//! properties
    Modules modules;		//! modules
    std::vector < CVC3::Expr > user_predicates;	//! user-added predicates
    std::vector < CVC3::Expr > invariants;	//! invariants
    CVC3::ExprHashMap < std::pair < CVC3::Expr, CVC3::Expr > >variable_bounds;	//! bound on a variable
    CVC3::ExprHashMap < CVC3::Expr > default_initial_value;	//! default initial value if no explicit initial state given
    std::vector < std::string > var_names;	//! variables of global namespace
    Actions actions;		//! set of actions
    unsigned maxNrOfAlt;	//! maximal number of alternatives in assignment
    std::vector<std::vector<std::pair<CVC3::Expr, CVC3::Expr> > > stateRewards;
    std::vector<std::vector<std::pair<std::pair<Action, CVC3::Expr>,
      CVC3::Expr> > > transRewards;
  };

  /*! \brief A model is typically the parallel composition of different modules */
  struct Module : Node {
    /*****************************************************************/
    /* The following part is only interesting for model construction */
    /*****************************************************************/
    Module(const std::string &);
    Module();
    /*! \brief Instantiate with existing module
      \par existing module in which to substitute
      \par lhs left-hand sides of substitution
      \par rhs right-hand sides of substitution
      \par al  mapping of action labels
    */
    Module(const std::string & name,
	   const Module & existing,
	   const CVC3::ExprHashMap < CVC3::Expr > &repl,
	   const HashMap < std::string, std::string > &al);

    /*! \brief add a command */
    void addCommand(Command *);

    /* Getters */
    const std::string &getName() const;

    /* Node functions */
    virtual std::string toString() const;

  private:
    friend class Model;

    SynchronizedCommands sync_guarded_transitions;	//! synchronized guarded transtions
    Commands guarded_transitions;	//! interleaved guarded transitions
    std::tr1::unordered_map < std::string, CVC3::Expr > local_vars;	//! local variables
    std::string name;		//! name of the module
  };

  class Alternative : Node {
  public:
    Alternative();

    //! \brief create copy of guarded transition with substitution[lhs/rhs]
    //Assignment ( const Assignment&, const std::vector<Expr>& lhs, const std::vector<Expr>& rhs);
    //! \brief create copy of guarded transition with substitution[lhs/rhs]
    Alternative(const Alternative &, const CVC3::ExprHashMap<CVC3::Expr> &);

    void setWeight(const CVC3::Expr &);
    const CVC3::Expr &getWeight();

    //! \brief return rate as double (if possible!!)
    //! \brief returns whether rate is a double
    bool isRateDouble();
    double getRateAsDouble() const;

    /*! \brief does any right-hand side share an identifier with given expr */
    bool LvalueCompatible( const CVC3::Expr& ) const;

    /* Node functions */
    virtual std::string toString() const;
    virtual Alternative* Clone() const;
    virtual void Cleanup();
    //void Assign(BasicExpr* left, BasicExpr* right);

    void Assign(const CVC3::Expr &, const CVC3::Expr &);

    /*! \brief multiply with another assignment */
    Alternative* operator*(const Alternative &) const;
    /*! \brief access to lhs given a right-hand side */
    //Expr operator[](const Expr&) const;

    CVC3::Expr operator[](const CVC3::Expr &) const;

    // Apply an assignment to an expression
    CVC3::Expr operator()(const CVC3::Expr &) const;

    // Apply an assignment to a state
    void getSuccessor(CVC3::ExprHashMap<CVC3::Expr> &,
		      CVC3::ExprHashMap<CVC3::Expr> &) const;

    void CollectLhs(std::set<CVC3::Expr> &) const;

    typedef CVC3::ExprHashMap<CVC3::Expr> Map;

    const Map &getMap() const;
    const std::set<CVC3::Expr>& getSupport() const;
    
  private:
    CVC3::Expr p; //! weight of the assignment
    bool is_rate_double;
    double p_double; //! weight of assignment in case it's constant
    Map map;
    Map wp_cache;
    std::set<CVC3::Expr> support;
  };

  /*! \brief A command
    \note A command consists of an action label, a guard, and different weighted alternatives
  */
  struct Command : public Node {
    const CVC3::Expr &getGuard() const;
    const std::string & getAction() const;
    const Alternatives &getAlternatives() const;

    /*! \brief does any right-hand side share an identifier with given expr */
    bool LvalueCompatible(const CVC3::Expr &) const;

    unsigned getNrOfAlt() const;
    /*****************************************************************/
    /* The following part is only interesting for model construction */
    /*****************************************************************/
    /*! \brief setter for Command::guard */
    void setGuard(const CVC3::Expr & __g);

    /*! \brief setter for Command::action */
    void setAction(const std::string & __action) {
      action = __action;
    }
    /*! \brief add another Assignment */
    void addAlternative(Alternative *);
    /*! \brief Perform consistency check on guarded transition (probabilities sum up to 1) */
    void checkSanity();

    //! \brief plain constructor
    Command();
    /*! \brief create copy of guarded transition with some variables replaced
      \par lhs left-hand sides of substitution
      \par rhs right-hand sides of substitution
      \par al  mapping of action labels
    */
    Command(const Command & gt,
	    const CVC3::ExprHashMap < CVC3::Expr > &repl);

    /* Node functionality */
    virtual std::string toString () const;

    const Alternative& operator[](unsigned) const;

    /*! \brief compute sum of two commands
      \note ([] guard_1 -> D_1) + ([] guard_2 -> D_2) = [] guard_1 & guard_2 -> D_1 + D_2 */
    Command *operator+(const Command &) const;

    /*! \brief compute product for synchronous parallel composition
      \note ([] guard_1 -> D_1) * ([] guard_2 -> D_2) = [] guard_1 & guard_2 -> D_1 * D_2
      @see Model::Flatten()
    */
    Command *operator*(const Command &) const;

    /*!
     * get the updates in which the alternatives
     * of the command differ
     * Example:
     * for command
     * [] x > 2 -> 0.1: (x'=1) & (y'=1) + 0.8: (x'=2) & (y'=1) + 0.1: (x'=3) & (y'=1)
     * prob_choices 0, 2
     * => ( base = [y'=1] , alt = { [x -> 1], [x -> 3] } )
     */
    void factorize(const std::vector<int>& prob_choices,
		   Alternative::Map& base,
		   std::vector<Alternative::Map>& alt
		   ) const;

    const std::set<CVC3::Expr>& getSupport() const;


    CVC3::Expr WP(const CVC3::Expr& e, unsigned i) const;
  private:
    Action action;		//! action on which transition is triggered
    CVC3::Expr g;		//! guard
    Alternatives alternatives;	//! weighted alternatives
    std::set<CVC3::Expr> support;

    friend class Model;
    friend class AbsModel;
  };
}

#endif
