#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,
      const VariableAssignment<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,
      const VariableAssignment<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;
    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,
      const VariableAssignment<Domain>& rho,
      IdSet<Variable<Domain> >& changedIn,
      IdSet<Variable<Domain> >& changedOut
  ) const {
    if (_subImprovement.improve(system, strat, rho, changedIn, changedOut)) {
      VariableAssignment<Domain>* rho2 = system.eval(rho, strat);
      improve(system, strat, *rho2, changedIn, changedOut);
      delete rho2;
      return true;
    }
    return false;
  }
  private:
  const ImprovementOperator<Domain>& _subImprovement;
};

template<typename Domain>
struct SmartImprovementOperator : public ImprovementOperator<Domain> {
  SmartImprovementOperator(const EquationSystem<Domain>& system)
    : _system(system),
      _influence(system.maxExpressionCount(), IdSet<MaxExpression<Domain> >(system.maxExpressionCount())),
      _first_run(true) {
  }

  struct DynamicMaxStrategy : public MaxStrategy<Domain> {
    DynamicMaxStrategy(
        const EquationSystem<Domain>& system,
        ConcreteMaxStrategy<Domain>& strat,
        IdMap<MaxExpression<Domain>,IdSet<MaxExpression<Domain>>>& influence,
        const VariableAssignment<Domain>& rho,
        IdSet<MaxExpression<Domain>>& stable,
        IdSet<MaxExpression<Domain>>& changed
    ) : _system(system),
        _rho(rho),
        _values(strat),
        _stable(stable),
        _changed(changed),
        _influence(influence) {
    }

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

    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)) {
          auto oldInfluence = _influence[x];
          _influence[x].clear();
          _values.set(x, val);
          _changed.insert(x);

          _stable.filter(oldInfluence);

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

    struct DependencyAssignment : public VariableAssignment<Domain>{
      DependencyAssignment(const DynamicMaxStrategy& strat,
                           const VariableAssignment<Domain>& rho,
                           const MaxExpression<Domain>& expr)
        : _strat(strat),
          _rho(rho),
          _expr(expr) {
      }
      const Domain& operator[](const Variable<Domain>& var) const {
        _strat._influence[*_strat._system[var]].insert(_expr);
        return _rho[var];
      }
      private:
      const DynamicMaxStrategy& _strat;
      const 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);
        _strat._influence[e].insert(_expr);
        return _strat._values.get(e);
      }
      private:
      const DynamicMaxStrategy& _strat;
      const MaxExpression<Domain>& _expr;
    };

    private:
    const EquationSystem<Domain>& _system;
    const VariableAssignment<Domain>& _rho;
    ConcreteMaxStrategy<Domain>& _values;
    IdSet<MaxExpression<Domain>>& _stable;
    IdSet<MaxExpression<Domain>>& _changed;
    IdMap<MaxExpression<Domain>,IdSet<MaxExpression<Domain>>>& _influence;
  };

  void invalidateSubexpressions(IdSet<MaxExpression<Domain>>& set, Expression<Domain>& expr) const {
    MaxExpression<Domain>* maxExpr = static_cast<MaxExpression<Domain>*>(&expr);
    if (maxExpr != NULL && maxExpr->id() < _system.maxExpressionCount()) {
      invalidateAll(set, *maxExpr);
      for (auto it = maxExpr->arguments().begin(),
             end = maxExpr->arguments().end();
           it != end;
           ++it) {
        invalidateSubexpressions(set, **it);
      }
    }
  }

  void invalidateAll(IdSet<MaxExpression<Domain>>& set, MaxExpression<Domain>& expr) const {
    if (!set.contains(expr)) {
      set.insert(expr);
      for (auto it = _influence[expr].begin(),
             end = _influence[expr].end();
           it != end;
           ++it) {
        invalidateSubexpressions(set, _system.maxExpression(*it));
      }
    }
  }

  void markChanged(IdSet<Variable<Domain>>& set, MaxExpression<Domain>& expr, IdSet<MaxExpression<Domain>>& seen) const {
    if (seen.contains(expr)) return;
    seen.insert(expr);
    auto var = _system.varFromExpr(expr);
    if (var == NULL || !set.contains(*var)) {
      if (var != NULL) {
        set.insert(*var);
      }
      for (auto it = _influence[expr].begin(),
                end = _influence[expr].end();
           it != end;
           ++it) {
        markChanged(set, _system.maxExpression(*it), seen);
      }
    }
  }

  bool improve(
      EquationSystem<Domain>&,
      ConcreteMaxStrategy<Domain>& stratOut,
      const VariableAssignment<Domain>& rho,
      IdSet<Variable<Domain> >& changedIn,
      IdSet<Variable<Domain> >& changedOut
  ) const {
    IdSet<MaxExpression<Domain>> invalidSet(_system.maxExpressionCount());
    IdSet<MaxExpression<Domain>> changedSet(_system.maxExpressionCount());
    if (_first_run) {
      _first_run = false;
    } else {
      for (auto it = changedIn.begin(),
                end = changedIn.end();
           it != end;
           ++it) {
        invalidateSubexpressions(invalidSet, *_system[_system.variable(*it)]);
      }
      invalidSet.invert();
    }

    log::debug << "stable: " << invalidSet;
    log::debug << "infl: " << _influence;
    DynamicMaxStrategy strat(_system, stratOut, _influence, rho, invalidSet, changedSet);
    auto inv = invalidSet;
    for (unsigned int i = 0, length = _system.maxExpressionCount();
         i < length;
         ++i) {
      strat.get(_system.maxExpression(i));
    }
    IdSet<MaxExpression<Domain>> seen;
    for (auto it = changedSet.begin(),
              end = changedSet.end();
         it != end;
         ++it) {
      markChanged(changedOut, _system.maxExpression(*it), seen);
    }
    return changedOut.begin() != changedOut.end(); // true iff changedOut is non-empty
  }
  private:
  const EquationSystem<Domain>& _system;
  mutable IdMap<MaxExpression<Domain>,IdSet<MaxExpression<Domain>>> _influence;
  mutable bool _first_run;
};

#endif