#ifndef MAX_EXPRESSION_HPP
#define MAX_EXPRESSION_HPP

#include "Expression.hpp"
#include "EquationSystem.hpp"
#include "IdSet.hpp"

template<typename Domain>
struct MaxStrategy : public EquationSystem<Domain> {
  MaxStrategy(const EquationSystem<Domain>& sub)
    : _sub(sub) { }

  virtual ~MaxStrategy() { }
  virtual const Expression<Domain>* operator[](const Variable<Domain>& var) const {
    return _sub[var];
  }
  virtual VariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho) const {
    return _sub.eval(rho);
  }
  virtual unsigned int variableCount() const {
    return _sub.variableCount();
  }
  virtual Variable<Domain>& variable(unsigned int i) const {
    return _sub.variable(i);
  }
  virtual StableVariableAssignment<Domain>* assignment(const Domain& value) const {
    return _sub.assignment(value);
  }
  virtual bool equalAssignments(const VariableAssignment<Domain>& r1, const VariableAssignment<Domain>& r2) const {
    return _sub.equalAssignments(r1, r2);
  }
  virtual Variable<Domain>* varFromExpr(const Expression<Domain>& expr) const {
    return _sub.varFromExpr(expr);
  }
  virtual void print(std::ostream& cout) const {
    return _sub.print(cout);
  }

  virtual unsigned int get(const MaxExpression<Domain>& e) const = 0;
  private:
  const EquationSystem<Domain>& _sub;
};

template<typename Domain>
struct MaxStrategyExpression : public Expression<Domain> {
  MaxStrategyExpression(const Expression<Domain>& expr, const MaxStrategy<Domain>& strategy)
    : _expr(expr), _strategy(strategy) { }

  virtual Domain eval(const VariableAssignment<Domain>& rho) const {
    // relies on implementation details - BAD BAD BAD, maybe
    const OperatorExpression<Domain>* opExpr = dynamic_cast<const OperatorExpression<Domain>*>(&_expr);
    if (opExpr) {
      const MaxExpression<Domain>* maxExpr = dynamic_cast<const MaxExpression<Domain>*>(opExpr);
      const std::vector<Expression<Domain>*> args = opExpr->arguments();
      if (maxExpr) {
        unsigned int i = _strategy.get(*maxExpr);
        return MaxStrategyExpression(*args[i], _strategy).eval(rho);
      } else {
        std::vector<Domain> argumentValues;
        for (typename std::vector<Expression<Domain>*>::const_iterator it = args.begin();
             it != args.end();
             ++it) {
          argumentValues.push_back(MaxStrategyExpression(**it, _strategy).eval(rho));
        }
        return opExpr->op().eval(argumentValues);
      }
    } else {
      return _expr.eval(rho);
    }
  }

  void print(std::ostream& cout) const {
    cout << _expr;
  }
  private:
  const Expression<Domain>& _expr;
  const MaxStrategy<Domain>& _strategy;
};

template<typename Domain>
struct ConcreteMaxStrategy : public MaxStrategy<Domain> {
  ConcreteMaxStrategy(const ConcreteEquationSystem<Domain>& system)
    : MaxStrategy<Domain>(system),
      _system(system),
      _expressions(system.variableCount(), NULL),
      _strategy(system.maxExpressionCount(), 0) {
  }

  ~ConcreteMaxStrategy() {
    for (int i = 0, length = _system.variableCount();
         i < length;
         ++i) {
      Expression<Domain>* expr = _expressions[_system.variable(i)];
      if (expr)
        delete expr;
    }
  }

  const Expression<Domain>* operator[](const Variable<Domain>& v) const {
    if (_expressions[v] == NULL) {
      Expression<Domain>* expression = new MaxStrategyExpression<Domain>(*_system[v], *this);
      _expressions[v] = expression;
      return expression;
    } else {
      return _expressions[v];
    }
  }

  unsigned int get(const MaxExpression<Domain>& e) const {
    return _strategy[e];
  }
  unsigned int set(const MaxExpression<Domain>& e, unsigned int i) {
    _strategy[e] = i;
    return i;
  }

  VariableAssignment<Domain>* eval(const VariableAssignment<Domain>& rho) const {
    StableVariableAssignment<Domain>* result = this->assignment(-infinity<Domain>());
    for (unsigned int i = 0, length = _system.variableCount();
         i < length;
         ++i) {
      const Variable<Domain>& var = _system.variable(i);
      const Expression<Domain>& expr = * (*this)[var];
      (*result)[var] = expr.eval(rho);
    }
    return result;
  }

  unsigned int variableCount() const {
    return _system.variableCount();
  }

  Variable<Domain>& variable(unsigned int i) const {
    return _system.variable(i);
  }
  
  unsigned int maxExpressionCount() const {
    return _system.maxExpressionCount();
  }

  MaxExpression<Domain>& maxExpression(unsigned int i) const {
    return _system.maxExpression(i);
  }

  StableVariableAssignment<Domain>* assignment(const Domain& v) const {
    return _system.assignment(v);
  }

  bool equalAssignments(const VariableAssignment<Domain>& l, const VariableAssignment<Domain>& r) const {
    return _system.equalAssignments(l, r);
  }

  void print(std::ostream& cout) const {
    cout << _system << std::endl;
  }

  const ConcreteEquationSystem<Domain>& system() const {
    return _system;
  }

  const IdMap<MaxExpression<Domain>,unsigned int>& strategy() const {
    return _strategy;
  }

  private:
  const ConcreteEquationSystem<Domain>& _system;
  mutable IdMap<Variable<Domain>,Expression<Domain>*> _expressions;
  IdMap<MaxExpression<Domain>,unsigned int> _strategy;
};

#endif