#ifndef EQUATION_SYSTEM_HPP
#define EQUATION_SYSTEM_HPP

#include <vector>
#include <set>
#include <map>
#include "Operator.hpp"
#include "Expression.hpp"
#include "VariableAssignment.hpp"
#include "IdMap.hpp"

template<typename Domain>
struct MaxStrategy;

template<typename Domain>
struct EquationSystem {
  EquationSystem()
    : _expr_to_var(NULL) { }

  virtual ~EquationSystem() {
    for (typename std::set<Expression<Domain>*>::iterator it = _expressions.begin();
         it != _expressions.end();
         ++it) {
      delete *it;
    }
    for (typename std::set<Operator<Domain>*>::iterator it = _operators.begin();
         it != _operators.end();
         ++it) {
      delete *it;
    }
    delete _expr_to_var;
  }

  MaxExpression<Domain>& maxExpression(const std::vector<Expression<Domain>*>& arguments) {
    unsigned int id = _max_expressions.size();
    Maximum<Domain>* max = new Maximum<Domain>();
    MaxExpression<Domain>* expr = new MaxExpression<Domain>(id, *max, arguments);
    _operators.insert(max);
    _max_expressions.push_back(expr);
    _expressions.insert(expr);
    return *expr;
  }
  MaxExpression<Domain>& maxExpression(unsigned int i) const {
    return *_max_expressions[i];
  }
  unsigned int maxExpressionCount() const {
    return _max_expressions.size();
  }

  Expression<Domain>& expression(Operator<Domain>* op, const std::vector<Expression<Domain>*>& arguments) {
    Expression<Domain>* expr = new OperatorExpression<Domain>(*op, arguments);
    _operators.insert(op);
    _expressions.insert(expr);
    return *expr;
  }

  Variable<Domain>& variable(const std::string& name) {
    if (_variable_names.find(name) == _variable_names.end()) {
      // not found - create a new variable and whatnot
      unsigned int id = _variables.size();
      Variable<Domain>* var = new Variable<Domain>(id, name);
      _variables.push_back(var);
      _right_sides.push_back(NULL);
      _expressions.insert(var);
      _variable_names[name] = var;
      return *var;
    } else {
      return *_variable_names[name];
    }
  }
  Variable<Domain>& variable(unsigned int id) const {
    return *_variables[id];
  }
  unsigned int variableCount() const {
    return _variables.size();
  }

  Constant<Domain>& constant(const Domain& value) {
    Constant<Domain>* constant = new Constant<Domain>(value);
    _expressions.insert(constant);
    return *constant;
  }

  MaxExpression<Domain>* operator[](const Variable<Domain>& var) const {
    return _right_sides[var.id()];
  }
  MaxExpression<Domain>*& operator[](const Variable<Domain>& var) {
    return _right_sides[var.id()];
  }

  void indexMaxExpressions() {
    _expr_to_var = new IdMap<MaxExpression<Domain>,Variable<Domain>*>(maxExpressionCount(), NULL); 
    for (unsigned int i = 0, length = _right_sides.size(); i < length; ++i) {
      _right_sides[i]->mapTo(*_variables[i], *_expr_to_var);
    }
  }

  Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const {
    assert(_expr_to_var); // make sure we've indexed
    const MaxExpression<Domain>* maxExpr = expr.toMaxExpression();//dynamic_cast<const MaxExpression<Domain>*>(&expr);
    if (maxExpr) {
      return (*_expr_to_var)[*maxExpr];
    } else {
      return NULL;
    }
  }

  virtual bool equalAssignments(VariableAssignment<Domain>& l, VariableAssignment<Domain>& r) const {
    for (unsigned int i = 0, length = _variables.size();
         i < length;
         ++i) {
      const Variable<Domain>& var = *_variables[i];
      if (l[var] != r[var])
        return false;
    }
    return true;
  }

  void print(std::ostream& cout) const {
    for (unsigned int i = 0, length = _variables.size();
         i < length;
         ++i) {
      cout << *_variables[i] << " = " << *_right_sides[i] << std::endl;
    }
  }

  private:
  std::set<Operator<Domain>*> _operators;
  std::set<Expression<Domain>*> _expressions;
  std::vector<Variable<Domain>*> _variables;
  std::map<std::string, Variable<Domain>*> _variable_names;
  IdMap<MaxExpression<Domain>, Variable<Domain>*>* _expr_to_var;
  std::vector<MaxExpression<Domain>*> _max_expressions;
  std::vector<MaxExpression<Domain>*> _right_sides;
};

template<typename T>
std::ostream& operator<<(std::ostream& cout, const EquationSystem<T>& system) {
  system.print(cout);
  return cout;
}

#include "MaxStrategy.hpp"

#endif