#include "clang/Analysis/Analyses/Interval.h"
#include "clang/AST/Stmt.h"
#include "clang/Analysis/CFG.h"
#include "clang/Analysis/AnalysisContext.h"
#include "clang/AST/StmtVisitor.h"

#include "clang/Analysis/Analyses/IntervalSolver/Log.hpp"
#include "clang/Analysis/Analyses/IntervalSolver/Complete.hpp"
#include "clang/Analysis/Analyses/IntervalSolver/VariableAssignment.hpp"
#include "clang/Analysis/Analyses/IntervalSolver/EquationSystem.hpp"
  
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Process.h"
  
#include <deque>
#include <algorithm>
#include <vector>
#include <map>
#include <set>

using namespace clang;

#include <sstream>
template<typename T>
std::string toString(const T& obj) {
  std::stringstream stream;
  stream << obj;
  return stream.str();
}

#include <ostream>
template<typename K,typename V>
std::ostream& operator<<(std::ostream& cout, const std::pair<K,V>& v) {
  cout << "(" << v.first << ", " << v.second << ")";
  return cout;
}

template<typename K,typename V>
std::ostream& operator<<(std::ostream& cout, const std::map<K,V>& v) {
  cout << "{";
  for (typename std::map<K,V>::const_iterator it = v.begin(), ei = v.end();
       it != ei;
       ++it) {
    if (it != v.begin())
      cout << ", ";
    cout << it->first << ": " << it->second;
  }
  cout << "}";
  return cout;  
}


IntervalAnalysis :: IntervalAnalysis(AnalysisDeclContext &context)
  : context(&context) {
}

IntervalAnalysis :: ~IntervalAnalysis() {
}

// Two pieces of state:
//  -> condition protecting a node
//  -> node's expression itself
// We can then combine these in a straightforward way to
// get out equation system, whereupon we can solve for what
// we want to know. Then we can have program invariants!
//
// Hooray!

typedef Complete<int64_t> ZBar;
//typedef std::map<std::string, ZBar> Vector;

struct Vector : public std::map<std::string, ZBar> {
  Vector(const ZBar& val=infinity<ZBar>())
    : _val(val) { }
  ZBar operator[](const std::string& key) const {
    if (this->find(key) != this->end())
      return this->find(key)->second;
    return _val;
  }
  ZBar& operator[](const std::string& key) {
    if (this->find(key) != this->end())
      return this->find(key)->second;
    std::pair<iterator,bool> p = this->insert(std::pair<const std::string, ZBar>(key, _val));
    return p.first->second;
  }
  ZBar _val;
};

typedef std::pair<Vector, ZBar> Result; // a "slice" of an equation

//typedef std::map<std::string, Result> LinearEquation; // one `Result` per variable
struct LinearEquation : public std::map<std::string, Result> {
  Result operator[](const std::string& key) const {
    if (this->find(key) != this->end())
      return this->find(key)->second;
    Result r;
    r.first[key] = 1;
    r.second = 0;
    return r;
  }
  Result& operator[](const std::string& key) {
    if (this->find(key) != this->end())
      return this->find(key)->second;
    Result r;
    r.first[key] = 1;
    r.second = 0;
    std::pair<iterator,bool> p = this->insert(std::pair<const std::string, Result>(key, r));
    return p.first->second;
  }
};

typedef Vector Condition;

typedef EquationSystem<Vector> EqnSys;
typedef Expression<Vector> EqnExpr;
typedef Variable<Vector> EqnVar;


struct LinearOperator : public Operator<Vector> {
  LinearOperator(const LinearEquation& result)
    : _values(result) {}

  Vector eval(const std::vector<Vector>& vector) const {
    assert(vector.size() == 1);
    const Vector& v = vector[0];
    Vector result = v;
    for (LinearEquation::const_iterator it = _values.begin(),
           ei = _values.end();
         it != ei;
         ++it) {
      ZBar subresult = 0;
      for (Vector::const_iterator jt = it->second.first.begin(),
             ej = it->second.first.end();
           jt != ej;
           ++jt) {
        subresult += jt->second * v[jt->first];
      }
      subresult += it->second.second;
      result[it->first] = subresult;
    }
    return result;
  }

  void print(std::ostream& cout) const {
    cout << "linear[" << _values << "]";
  }

  LinearEquation _values;
};



template<class F, class M>
void transform_values(const F& f, M& map) {
  for (typename M::iterator it = map.begin(),
         ei = map.end();
       it != ei;
       ++it) {
    it->second = f(it->second);
  }
}

template<class M, class F> 
M merge_maps_with(const F& f, const M& left, const M& right) {
  M result;
  typename M::const_iterator first1 = left.begin(), last1 = left.end(),
    first2 = right.begin(), last2 = right.end();
  for (; first1 != last1 && first2 != last2;) {
    if (first2->first < first1->first) {
      result[first2->first] = first2->second;
      ++first2;
    } else if (first1->first == first2->first) {
      result[first1->first] = f(first1->second, first2->second);
      ++first1;
      ++first2;
    } else {
      result[first1->first] = first1->second;
      ++first1;
    }
  }
  while (first1 != last1) {
    result[first1->first] = first1->second;
    ++first1;
  }
  while (first2 != last2) {
    result[first2->first] = first2->second;
    ++first2;
  }
  return result;
}

template<>
Vector minimum(const Vector& l, const Vector& r) {
  return merge_maps_with(minimum<ZBar>, l, r);
}
template<class T>
T max(const T& l, const T& r) {
  return (l < r ? l : r);
}
template<class T>
T negate(const T& v) {
  return -v;
}
template<class T>
T addValues(const T& l, const T& r) {
  return l + r;
}

Vector operator-(const Vector& vector) {
  Vector result(-vector._val);
  for (Vector::const_iterator it = vector.begin(),
         ei = vector.end();
       it != ei;
       ++it) {
    result[it->first] = -it->second;
  }
  return result;
}

Vector operator+(const Vector& left, const Vector& right) {
  return merge_maps_with(addValues<ZBar>, left, right);
}

Vector operator-(const Vector& left, const Vector& right) {
  return merge_maps_with(addValues<ZBar>, left, -right);
}

Vector operator*(const Vector& left, const ZBar& right) {
  Vector result;
  for (Vector::const_iterator it = left.begin(),
         ei = left.end();
       it != ei;
       ++it) {
    result[it->first] = (it->second * right);
  }
  return result;
}
Vector operator*(const ZBar& left, const Vector& right) {
  return right * left;
}
bool operator<(const Vector& left, const Vector& right) {
  for (Vector::const_iterator it = left.begin(),
         ei = left.end();
       it != ei;
       ++it) {
    if (it->second < right[it->first]) {
      return true;
    }
  }
  for (Vector::const_iterator it = right.begin(),
         ei = right.end();
       it != ei;
       ++it) {
    if (left[it->first] < it->second) {
      return true;
    }
  }
  return false;
}

template<>
Vector infinity<Vector>() {
  return Vector(infinity<ZBar>());
}

std::ostream& operator<<(std::ostream& cout, const Vector& v) {
  cout << "{";
  for (Vector::const_iterator it = v.begin(), ei = v.end();
       it != ei;
       ++it) {
    cout << it->first << ": " << it->second << ", ";
  }
  cout << "_: " << v._val;
  cout << "}";
  return cout;
}


Result fromStmt(const Stmt*);

Result fromInteger(const IntegerLiteral* expr) {
  return Result(Vector(), *expr->getValue().getRawData());
}

Result fromDeclExpr(const DeclRefExpr* expr) {
  Vector val;
  val[expr->getNameInfo().getAsString()] = 1;
  return Result(val, 0);
}

Result fromUnary(const UnaryOperator* op) {
  switch (op->getOpcode()) {
  case UO_PreInc:
    break;
  case UO_PostInc:
    break;
  }
  return Result(Vector(), 0);
}

Result fromBinary(const BinaryOperator* op) {
  Result left = fromStmt(op->getLHS()->IgnoreParenCasts());
  Result right = fromStmt(op->getRHS()->IgnoreParenCasts());
  
  switch (op->getOpcode()) {
  case BO_Assign:
    return right;
  case BO_Sub:
    transform_values(negate<ZBar>, right.first);
    right.second *= -1;
  case BO_Add:
    {
      Result result;
      result.first = merge_maps_with(addValues<ZBar>,
                                     left.first, right.first);
      result.second = left.second + right.second;
      return result;
    }
  case BO_Mul:
    {
      if (!left.first.empty() && !right.first.empty()) {
        return Result(Vector(), 0);
      }
      ZBar scalar = 0;
      Result value;
      if (left.first.empty()) {
        scalar = left.second;
        value = right;
      } else {
        scalar = right.second;
        value = left;
      }
      if (scalar > 0) {
        for (Vector::iterator it = value.first.begin(),
               ei = value.first.end();
             it != ei;
             ++it) {
          it->second *= scalar;
        }
      } else {
        Vector newValue;
        for (Vector::iterator it = value.first.begin(),
               ei = value.first.end();
             it != ei;
             ++it) {
          std::string name;
          if (it->first[0] == '-') {
            name = it->first.substr(1);
          } else {
            name = '-' + it->first;
          }
          newValue[name] = (-scalar) * it->second;
        }
      }
      right.second *= scalar;
      return value;
    }
  case BO_LT:
  case BO_LE:
  case BO_GT:
  case BO_GE:
    break;
  } 
  return Result();
}

Result fromDeclStmt(const DeclStmt* stmt) {
  for (DeclStmt::const_decl_iterator it = stmt->decl_begin(),
                                     ei = stmt->decl_end();
       it != ei;
       ++it) {
    if ((*it)->getKind() == Decl::Var) {
      const VarDecl* decl = static_cast<const VarDecl*>(*it);
      llvm::errs() << decl->getNameAsString() << " = ";

      Result expr = fromStmt(decl->getInit());
      for (Vector::iterator it = expr.first.begin(),
             ei = expr.first.end();
           it != ei;
           ++it) {
        if (it != expr.first.begin())
          llvm::errs() << " + ";
        llvm::errs() << toString(it->second) << "*" << toString(it->first);
      }
      llvm::errs() << " + " << toString(expr.second) << "\n";


      llvm::errs() << "-" << decl->getNameAsString() << " = ";
      for (Vector::iterator it = expr.first.begin(),
             ei = expr.first.end();
           it != ei;
           ++it) {
        if (it != expr.first.begin())
          llvm::errs() << " + ";
        llvm::errs() << toString(it->second) << "*" << toString(it->first);
      }
      llvm::errs() << " - " << toString(expr.second) << "\n";
    }
  }

  return Result();
}

Result fromAssignment(const BinaryOperator* op) {
  const Expr* left = op->getLHS()->IgnoreParenCasts();
  if (left->getStmtClass() == Stmt::DeclRefExprClass) {
    std::string name = static_cast<const DeclRefExpr*>(left)->getNameInfo().getAsString();

    Result expr = fromStmt(op->getRHS()->IgnoreParenCasts());
    llvm::errs() << name << " = ";
    for (Vector::iterator it = expr.first.begin(),
           ei = expr.first.end();
         it != ei;
         ++it) {
      if (it != expr.first.begin())
        llvm::errs() << " + ";
      llvm::errs() << toString(it->second) << "*" << toString(it->first);
    }
    llvm::errs() << " + " << toString(expr.second) << "\n";
    return expr;
  }
  return Result();
}

Result fromStmt(const Stmt* stmt) {
  if (!stmt)
    return Result();
  //stmt->dump();
  switch (stmt->getStmtClass()) {
  case Stmt::IntegerLiteralClass:
    return fromInteger(static_cast<const IntegerLiteral*>(stmt));
  case Stmt::DeclRefExprClass:
    return fromDeclExpr(static_cast<const DeclRefExpr*>(stmt));
  case Stmt::UnaryOperatorClass:
    return fromUnary(static_cast<const UnaryOperator*>(stmt));
  case Stmt::DeclStmtClass:
    return fromDeclStmt(static_cast<const DeclStmt*>(stmt));
  case Stmt::BinaryOperatorClass:
    return fromBinary(static_cast<const BinaryOperator*>(stmt));
  }
  const Expr* expr = dyn_cast<const Expr>(stmt);
  if (expr) {
    const Expr* expr2 = expr->IgnoreParenCasts();
    if (expr != expr2)
      return fromStmt(expr2);
  }
  llvm::errs() << "we shouldn't get here...\n";
  return Result();
}

Condition fromComparison(const BinaryOperator* op) {
  if (op->isRelationalOp()) {
    const Expr* left = op->getLHS()->IgnoreParenCasts();
    const Expr* right = op->getRHS()->IgnoreParenCasts();

    Condition cond;
    std::string name;
    int64_t value;
    if (left->getStmtClass() == Stmt::DeclRefExprClass) {
      if (right->getStmtClass() == Stmt::IntegerLiteralClass) {
        name = static_cast<const DeclRefExpr*>(left)->getNameInfo().getAsString();
        value = *static_cast<const IntegerLiteral*>(right)->getValue().getRawData();
        switch (op->getOpcode()) {
        case BO_LT:
          cond[name] = value - 1;
          break;
        case BO_LE:
          cond[name] = value;
          break;
        case BO_GE:
          cond['-' + name] = -value;
          break;
        case BO_GT:
          cond['-' + name] = -(value + 1);
          break;
        }
        return cond;
      } else {
        return Condition();
      }
    } else if (right->getStmtClass() == Stmt::DeclRefExprClass) {
      if (left->getStmtClass() == Stmt::IntegerLiteralClass) {
        name = static_cast<const DeclRefExpr*>(right)->getNameInfo().getAsString();
        value = *static_cast<const IntegerLiteral*>(left)->getValue().getRawData();
        switch (op->getOpcode()) {
        case BO_LT:
          cond['-' + name] = -(value + 1);
          break;
        case BO_LE:
          cond['-' + name] = -value;
          break;
        case BO_GE:
          cond[name] = value;
          break;
        case BO_GT:
          cond[name] = value - 1;
          break;
        }
        return cond;
      } else {
        return Condition();
      }
    }
    return Condition();
  }
  return Condition();
}

std::map<const CFGBlock*, EqnVar*> seen_blocks;
EqnVar* runOnBlock(const CFGBlock* block, EqnSys& system) {
  if (seen_blocks.find(block) != seen_blocks.end())
    return seen_blocks.find(block)->second;

  std::string id = toString(block->getBlockID());
  unsigned int counter = 0;
  
  // detemine equations for this block
  LinearEquation eqn;
  EqnVar* var;
  Result current;

  var = &system.variable(id);
  for (CFGBlock::const_iterator it = block->begin(),
                                ei = block->end();
       it != ei;
       ++it) {
    const CFGStmt* cfg_stmt = it->getAs<CFGStmt>();
    const Stmt* stmt = cfg_stmt->getStmt();

    std::string name = "";
    Result expr;
    if (stmt->getStmtClass() == Stmt::BinaryOperatorClass) {
      const BinaryOperator* binop = static_cast<const BinaryOperator*>(stmt);
      if (binop->isAssignmentOp()) {
        const Expr* left = binop->getLHS()->IgnoreParenCasts();
        const Expr* right = binop->getRHS()->IgnoreParenCasts();
        if (left->getStmtClass() == Stmt::DeclRefExprClass) {
          name = static_cast<const DeclRefExpr*>(left)->getNameInfo().getAsString();
          expr = fromStmt(right);
        }
      }
    } else if (stmt->getStmtClass() == Stmt::DeclStmtClass) {
      stmt->dump();
      const DeclStmt* decl_stmt = static_cast<const DeclStmt*>(stmt);
      for (DeclStmt::const_decl_iterator jt = decl_stmt->decl_begin(),
             ej = decl_stmt->decl_end();
           jt != ej;
           ++jt) {
        if ((*jt)->getKind() == Decl::Var) {
          const VarDecl* decl = static_cast<const VarDecl*>(*jt);
          name = decl->getNameAsString();
          expr = fromStmt(decl->getInit());
          if (jt+1 != ej) {
            llvm::errs() << "Only the first declaration in a multi-declaration statement is used.\n";
          }
          llvm::errs() << "Name: " << name << "\n";
          llvm::errs() << toString(expr) << "\n";
          break; // only take the first one, for now
        } else {
          llvm::errs() << (*jt)->getDeclKindName() << "\n";
        }
      }
    }
    if (name == "")
      continue;

    // by here we expect `name` and `expr` to be set correctly

    bool newBlock = false;
    for (Vector::const_iterator jt = expr.first.begin(),
           ej = expr.first.end();
         jt != ej;
         ++jt) {
      // new block if we read from any variable we've already written to
      newBlock |= (eqn.find(jt->first) != eqn.end());
    }

    if (newBlock) {
      EqnVar* lastVar = var;
      var = &system.variable(id + "-" + toString(++counter));

      // the linear expression
      std::vector<EqnExpr*> args;
      args.push_back(lastVar);
      EqnExpr* expr = &system.expression(new LinearOperator(eqn), args);

      // the max expression
      std::vector<EqnExpr*> maxArgs;
      maxArgs.push_back(&system.constant(-infinity<Vector>()));
      maxArgs.push_back(expr);
      system[*var] = &system.maxExpression(args);
      eqn = LinearEquation();
    }

    eqn[name] = expr;
    expr.first = expr.first;
    expr.second = -expr.second;
    eqn["-"+name] = expr;
  }

  EqnVar* lastVar = var;
  var = &system.variable(id + "-" + toString(++counter));

  // the linear expression
  std::vector<EqnExpr*> args;
  args.push_back(lastVar);
  EqnExpr* expr = &system.expression(new LinearOperator(eqn), args);

  // the max expression
  std::vector<EqnExpr*> maxArgs;
  maxArgs.push_back(&system.constant(-infinity<Vector>()));
  maxArgs.push_back(expr);
  system[*var] = &system.maxExpression(maxArgs);

  // save the last variable we used (the final values from this block)
  seen_blocks[block] = var;

  // determine predecessor variables/conditions
  // we have to do this last to prevent an infinite loop situation
  std::vector<EqnExpr*> preArgs;
  for (CFGBlock::const_pred_iterator it = block->pred_begin(),
         ei = block->pred_end();
       it != ei;
       ++it) {
    if ((*it)->getTerminatorCondition()) {
      const BinaryOperator* binop = static_cast<const BinaryOperator*>((*it)->getTerminatorCondition());
      std::vector<EqnExpr*> args;
      args.push_back(&system.constant(fromComparison(binop)));
      args.push_back(runOnBlock(*it, system));
      preArgs.push_back(&system.expression(new Minimum<Vector>(), args));
    } else {
      preArgs.push_back(runOnBlock(*it, system));
    }
  }
  EqnVar* preVar = &system.variable(id);
  if (preArgs.empty())
    preArgs.push_back(&system.constant(infinity<Vector>()));
  system[*preVar] = &system.maxExpression(preArgs);

  // return the final value we used
  return var;
}

void IntervalAnalysis::runOnAllBlocks() {
  llvm::errs() << "Enter run on all blocks\n";

  const CFG *cfg = this->context->getCFG();

  EqnSys system;

  std::set<const CFGBlock*> seen;
  std::deque<const CFGBlock*> todo;
  todo.push_back(&cfg->getEntry());

  while (!todo.empty()) {
    const CFGBlock* block = todo.front();
    if (seen.find(todo.front()) != seen.end()) {
      todo.pop_front();
      continue;
    }
    llvm::errs() << block->getBlockID() << "\n";
    seen.insert(block);
    todo.pop_front();
    runOnBlock(block, system);
    llvm::errs() << "-> ";
    for (CFGBlock::const_succ_iterator it = block->succ_begin(),
                                       ei = block->succ_end();
         it != ei;
         it++ ) {
      llvm::errs() << (*it)->getBlockID() << ", ";
      todo.push_back(*it);
    }
    llvm::errs() << "\n\n";
  }

  llvm::errs() << "Exit run on all blocks\n";

  llvm::errs() << toString(system) << "\n";

  system.indexMaxExpressions();
  DynamicMaxStrategy<Vector> strategy(system);
  DynamicVariableAssignment<Vector> rho(system, strategy);
  strategy.setRho(rho);

  for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
    Variable<Vector>& var = system.variable(i);
    llvm::errs() << toString(var.name()) << " = " << toString(rho[var]) << "\n";
  }

  //  cfg->dump(context->getASTContext().getLangOpts(),
  //            llvm::sys::Process::StandardErrHasColors());
}


const void *IntervalAnalysis::getTag() { static int x; return &x; }