#ifndef MAX_EXPRESSION_HPP
#define MAX_EXPRESSION_HPP

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

template<typename Domain>
struct MaxStrategyExpression : public Expression<Domain> {
  MaxStrategyExpression(const Expression<Domain>& expr, const IdMap<MaxExpression<Domain>,unsigned int>& 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[*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 IdMap<MaxExpression<Domain>,unsigned int>& _strategy;
};

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

  const Expression<Domain>* operator[](const Variable<Domain>& v) const {
    if (_expressions[v] == NULL) {
      Expression<Domain>* expression = new MaxStrategyExpression<Domain>(*_system[v], _strategy);
      _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);
  }

  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 improve(const VariableAssignment<Domain>& rho) {
    for (unsigned int i = 0, length = _system.maxExpressionCount();
         i < length;
         ++i) {
      MaxExpression<Domain>& expr = _system.maxExpression(i);
      Domain bestValue = MaxStrategyExpression<Domain>(expr, _strategy).eval(rho);
      unsigned int bestIndex = this->get(expr);

      // this relies on the fact that an expression will only be proessed after the expressions
      // it depends on (which should always be true, as they form a DAG)
      const std::vector<Expression<Domain>*> args = expr.arguments();
      for (unsigned int j = 0, length = args.size();
           j < length;
           ++j) {
        const Domain value = MaxStrategyExpression<Domain>(*args[j], _strategy).eval(rho);
        if (bestValue < value) {
          bestValue = value;
          bestIndex = j;
        }
      }
      this->set(expr, bestIndex);
    }
  }

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

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

#endif