#ifndef IMPROVEMENT_OPERATOR_HPP
#define IMPROVEMENT_OPERATOR_HPP

#include "IdSet.hpp"
#include "MaxStrategy.hpp"

template<typename Domain>
struct ImprovementOperator {
  virtual ~ImprovementOperator() { }
  virtual bool improve(
      EquationSystem<Domain>& system,
      ConcreteMaxStrategy<Domain>& strat,
      StableVariableAssignment<Domain>& rho,
      IdSet<Variable<Domain> >& changedIn,
      IdSet<Variable<Domain> >& changedOut
  ) const  = 0;
};

template<typename Domain>
struct NaiveImprovementOperator : public ImprovementOperator<Domain> {
  bool improve(
      EquationSystem<Domain>& system,
      ConcreteMaxStrategy<Domain>& strat,
      StableVariableAssignment<Domain>& rho,
      IdSet<Variable<Domain> >&,
      IdSet<Variable<Domain> >&
  ) const {
    bool changed = false;
    for (unsigned int i = 0, length = system.maxExpressionCount();
         i < length;
         ++i) {
      MaxExpression<Domain>& expr = system.maxExpression(i);
      // 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)
      unsigned int index = expr.bestStrategy(rho, strat);

      if (index != strat.get(expr)) {
        changed = true;
        strat.set(expr, index);
      }
    }
    //log::strategy << strat << std::endl;
    return changed;
  }
};
  
template<typename Domain>
struct RepeatedImprovementOperator : public ImprovementOperator<Domain> {
  RepeatedImprovementOperator(const ImprovementOperator<Domain>& op)
    : _subImprovement(op) { }

  bool improve(
      EquationSystem<Domain>& system,
      ConcreteMaxStrategy<Domain>& strat,
      StableVariableAssignment<Domain>& rho,
      IdSet<Variable<Domain> >& changedIn,
      IdSet<Variable<Domain> >& changedOut
  ) const {
    if (_subImprovement.improve(system, strat, rho, changedIn, changedOut)) {
      StableVariableAssignment<Domain>* rho2 = system.eval(rho, strat);
      improve(system, strat, *rho2, changedIn, changedOut);
      delete rho2;
      return true;
    }
    return false;
  }
  private:
  const ImprovementOperator<Domain>& _subImprovement;
};



#include "VariableAssignment.hpp"

template<typename Domain>
struct DynamicMaxStrategy : public MaxStrategy<Domain> {
  DynamicMaxStrategy(
    const EquationSystem<Domain>& system
  ) : _system(system),
      _rho(NULL),
      _values(system),
      _stable(system.maxExpressionCount()),
      _influence(system.maxExpressionCount(),
                 IdSet<MaxExpression<Domain>>(system.maxExpressionCount())),
      _var_influence(system.variableCount(),
                     IdSet<MaxExpression<Domain>>(system.maxExpressionCount()))
  {}

  void setRho(DynamicVariableAssignment<Domain>& rho) {
    _rho = &rho;
  }

  unsigned int get(const MaxExpression<Domain>& e) const {
    solve(e);
    return _values.get(e);
  }

  void invalidate(const Variable<Domain>& v) const {
    log::strategy << "invalidating " << v << " - " << *_system[v] << std::endl;
    log::strategy << "\t" << _var_influence[v] << std::endl;
    _stable.filter(_var_influence[v]);

    for (auto it = _var_influence[v].begin();
         it != _var_influence[v].end();
         ++it) {
      solve(_system.maxExpression(*it));
    }
  }

private:
  void solve(const MaxExpression<Domain>& x) const {
    if (!_stable.contains(x)) {
      _stable.insert(x);
      
      unsigned int val = x.bestStrategy(DependencyAssignment(*this, *_rho, x),
                                        DependencyStrategy(*this, x)); 

      if (val != _values.get(x)) {
        log::strategy << "Updating strategy for " << x
                      << " to " << val << std::endl;

        auto oldInfluence = _influence[x];
        _influence[x].clear();
        _values.set(x, val);

        _rho->invalidate(*_system.varFromExpr(x));

        _stable.filter(oldInfluence);

        for (auto it = oldInfluence.begin();
             it != oldInfluence.end();
             ++it) {
          solve(_system.maxExpression(*it));
        }
      } else {
        log::strategy << x << " did not change" << std::endl;
      } 
    } else {
      log::strategy << x << " is stable" << std::endl;
    }
  }

  struct DependencyAssignment : public VariableAssignment<Domain>{
    DependencyAssignment(const DynamicMaxStrategy& strat,
                         VariableAssignment<Domain>& rho,
                         const MaxExpression<Domain>& expr)
      : _strat(strat),
        _rho(rho),
        _expr(expr) {
    }
    const Domain& operator[](const Variable<Domain>& var) const {
      _strat._var_influence[var].insert(_expr);
      return _rho[var];
    }
  private:
    const DynamicMaxStrategy& _strat;
    VariableAssignment<Domain>& _rho;
    const MaxExpression<Domain>& _expr;
  };

  struct DependencyStrategy : public MaxStrategy<Domain> {
    DependencyStrategy(const DynamicMaxStrategy& strat, const MaxExpression<Domain>& expr)
      : _strat(strat),
        _expr(expr) {
    }
    unsigned int get(const MaxExpression<Domain>& e) const {
      _strat.solve(e);
      if (&_expr != &e) {
        _strat._influence[e].insert(_expr);
      }
      return _strat._values.get(e);
    }
  private:
    const DynamicMaxStrategy& _strat;
    const MaxExpression<Domain>& _expr;
  };

private:
  const EquationSystem<Domain>& _system;
  mutable DynamicVariableAssignment<Domain>* _rho;
  mutable ConcreteMaxStrategy<Domain> _values;
  mutable IdSet<MaxExpression<Domain>> _stable;
  mutable IdMap<MaxExpression<Domain>,IdSet<MaxExpression<Domain>>> _influence;
  mutable IdMap<Variable<Domain>,IdSet<MaxExpression<Domain>>> _var_influence;
};


#endif