#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;
using namespace std;

template<typename T>
T neg(const T& t) {
  return -t;
}

string operator-(const string& str) {
  if (str[0] == '-')
    return str.substr(1);
  return '-' + str;
}

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

#include <ostream>
template<typename K,typename V>
ostream& operator<<(ostream& cout, const pair<K,V>& v) {
  cout << "(" << v.first << ", " << v.second << ")";
  return cout;
}
template<typename V>
ostream& operator<<(ostream& cout, const pair<Variable<Complete<int64_t> >*, V>& v) {
  cout << "(" << v.first->name() << ", " << v.second << ")";
  return cout;
}

template<typename V>
ostream& operator<<(ostream& cout, const vector<V>& v) {
  cout << "[";
  for(typename vector<V>::const_iterator it = v.begin(), ei = v.end();
      it != ei;
      ++it) {
    if (it != v.begin())
      cout << ", ";
    cout << *it;
  }
  cout << "]";
  return cout;
}

template<typename K,typename V>
ostream& operator<<(ostream& cout, const map<K,V>& v) {
  cout << "{";
  for (typename 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;  
}

// 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;
template<>
ZBar infinity() {
  return ZBar(1, true);
}





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

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

Vector operator+(const Vector& left, const Vector& right) {
  Vector::const_iterator
    left_iter = left.begin(),
    left_end = left.end(),
    right_iter = right.begin(),
    right_end = right.end();

  Vector result(left._val + right._val);
  while (left_iter != left_end && right_iter != right_end) {
    if (left_iter->first == right_iter->first) {
      result[left_iter->first] = left_iter->second + right_iter->second;
      left_iter++;
      right_iter++;
    } else {
      if (left_iter->first < right_iter->first) {
        result[left_iter->first] = left_iter->second;
        left_iter++;
      } else {
        result[right_iter->first] = right_iter->second;
        right_iter++;
      }
    }
  }
  Vector::const_iterator it = (right_iter == right_end ? left_iter : right_iter);
  Vector::const_iterator end = (right_iter == right_end ? left_end : right_end);
  for (; it != end; ++it)
    result[it->first] = it->second;

  return result;
}
Vector operator*(const ZBar& left, const Vector& right) {
  Vector result(left * right._val);

  for (Vector::const_iterator
         it = right.begin(),
         end = right.end();
       it != end;
       ++it) {
    result[it->first] = left * it->second;
  }

  return result;
}
Vector operator*(const Vector& left, const ZBar& right) {
  return right * left;
}

ostream& operator<<(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;
}



typedef pair<Vector, ZBar> Result; // a "slice" of an equation
Result operator-(const Result& r) {
  return Result(-r.first, -r.second);
}

typedef Vector Condition;

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









/* Expression functions */

Result fromExpr(const Expr*);

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

Result fromDeclExpr(const DeclRefExpr* expr) {
  Vector val(0);
  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), 0);
}

Result operator*(const ZBar& l, const Result& r) {
  return Result(l * r.first, l * r.second);
}

Result fromBinary(const BinaryOperator* op) {
  Result left = fromExpr(op->getLHS()->IgnoreParenCasts());
  Result right = fromExpr(op->getRHS()->IgnoreParenCasts());
  
  switch (op->getOpcode()) {
  case BO_Assign:
    return right;
  case BO_Sub:
    right = -right;
    //transform_values(negate<ZBar>, right.first);
    //right.second *= -1;
  case BO_Add:
    {
      Result result;
      result.first = 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), 0);
      }
      ZBar scalar = 0;
      Result value;
      if (left.first.empty()) {
        scalar = left.second;
        value = right;
      } else {
        scalar = right.second;
        value = left;
      }
      if (scalar >= 0) {
        return scalar * value;
      } else {
        return scalar * -value;
      }
    }
  case BO_LT:
  case BO_LE:
  case BO_GT:
  case BO_GE:
    break;
  } 
  return Result();
}

Result fromExpr(const Expr* 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::BinaryOperatorClass:
    return fromBinary(static_cast<const BinaryOperator*>(stmt));
  }
  const Expr* expr = stmt->IgnoreParenCasts();
  if (stmt != expr)
    return fromExpr(expr);
  assert(false);
  return Result();
}


/* Comparison stuff */

Condition fromComparison(const BinaryOperator* op, bool negate) {
  Condition cond(infinity<ZBar>());
  if (!op) {
    if (negate)
      return -cond;
    else
      return cond;
  }
  if (op->isRelationalOp()) {
    const Expr* left = op->getLHS()->IgnoreParenCasts();
    const Expr* right = op->getRHS()->IgnoreParenCasts();

    bool flip = false;
    string name;
    int64_t value;
    if (left->getStmtClass() == Stmt::DeclRefExprClass) {
      name = static_cast<const DeclRefExpr*>(left)->getNameInfo().getAsString();
    } else if (right->getStmtClass() == Stmt::DeclRefExprClass) {
      name = static_cast<const DeclRefExpr*>(right)->getNameInfo().getAsString();
      flip = true;
    } else {
      return cond;
    }

    if (right->getStmtClass() == Stmt::IntegerLiteralClass) {
      value = *static_cast<const IntegerLiteral*>(right)->getValue().getRawData();
    } else if (left->getStmtClass() == Stmt::IntegerLiteralClass) {
      value = *static_cast<const IntegerLiteral*>(left)->getValue().getRawData();
    } else {
      return cond;
    }

    BinaryOperator::Opcode operation = op->getOpcode();
    if (flip) {
      switch (operation) {
      case BO_LT: operation = BO_GT; break;
      case BO_GT: operation = BO_LT; break;
      case BO_LE: operation = BO_GE; break;
      case BO_GE: operation = BO_LE; break;
      }
    }

    switch (operation) {
    case BO_LT:
      if (negate)
	cond[-name] = -value;
      else
	cond[name] = value - 1;
      break;
    case BO_LE:
      if (negate)
	cond[-name] = -(value + 1);
      else
	cond[name] = value;
      break;
    case BO_GE:
      if (negate)
	cond[name] = value - 1;
      else
	cond[-name] = -value;
      break;
    case BO_GT:
      if (negate)
	cond[name] = value;
      else
	cond[-name] = -(value + 1);
      break;
    }
  }
  return cond;
}




/* Blocks */

typedef map<string, unsigned int> Counters;
typedef map<string, EqnVar*> VarMap;
typedef map<const CFGBlock*, set<string> > BlockVars;

void runOnBlock(const CFGBlock* block, EqnSys& system, BlockVars& block_vars) {
  Counters counters;
  string block_id = toString(block->getBlockID());
  VarMap vars;

  for (set<string>::iterator it = block_vars[block].begin(),
	 ei = block_vars[block].end();
       it != ei;
       ++it) {
    vars[*it] = &system.variable(*it + '-' + block_id + "-pre");
  }
  
  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();

    string name = "";
    Result result;
    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();
          result = fromExpr(right);
        }
      }
    } else if (stmt->getStmtClass() == Stmt::DeclStmtClass) {
      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();
          result = fromExpr(decl->getInit());
          jt++;
          if (jt != ej) {
            llvm::errs() << "Only the first declaration in a multi-declaration statement is used.\n";
          }
          break; // only take the first one, for now
        }
      }
    }
    if (name == "")
      continue;

    string count = toString(counters[name]);
    EqnVar* var = &system.variable(name + '-' + block_id + '[' + count + ']');
    EqnVar* negative_var = &system.variable(-name + '-' + block_id + '[' + count + ']');
    counters[name]++;
    for (int negative = 0; negative < 2; ++negative) { // one loop for positive, the other for negative
      if (negative) {
	result = -result;
      }
      
      EqnExpr* expression;

      if (result.first.size() > 0) {
        // make sure all our variables exist in vars
        for (Vector::iterator
               it = result.first.begin(),
               ei = result.first.end();
             it != ei;
             ++it) {
          if (!vars[it->first])
            vars[it->first] = &system.variable(it->first + '-' + block_id + "-pre");
        }

        // set up the min-cost-flow operator
        vector<ZBar> supplies;
        vector<pair<int,int> > arcs;
        vector<EqnExpr*> minCostArgs;
        ZBar dummy_value = 0;
        supplies.push_back(dummy_value); // dummy node to suck up flow
        int index = 1; // the solver uses 1-indexing, for some reason
        for (map<std::string,EqnVar*>::iterator
               it = vars.begin(),
               ei = vars.end();
             it != ei;
             it++) {
          index++;
          supplies.push_back(result.first[it->first]);
          dummy_value -= result.first[it->first];
          if (it->first[0] == '-')
            arcs.push_back(pair<int,int>(1,index));
          else 
            arcs.push_back(pair<int,int>(index,1));
          minCostArgs.push_back(vars[it->first]);
        }
        supplies[0] = dummy_value;

        EqnExpr* minCostExpr = &system.expression(new MinCostFlow<ZBar>(supplies, arcs), minCostArgs);
        
        // add the constant factor to the min-cost bit
        vector<EqnExpr*> additionArgs;
        additionArgs.push_back(&system.constant(result.second));
        additionArgs.push_back(minCostExpr);
        expression = &system.expression(new Addition<ZBar>(), additionArgs);
      } else {
        expression = &system.constant(result.second);
      }

      // max(-inf, expr), so strategy iteration will work
      vector<EqnExpr*> maxArgs;
      maxArgs.push_back(&system.constant(-infinity<ZBar>()));
      maxArgs.push_back(expression);
      if (negative)
	system[*negative_var] = &system.maxExpression(maxArgs);
      else
	system[*var] = &system.maxExpression(maxArgs);
    }
    vars[name] = var;
    vars[-name] = negative_var;
    block_vars[block].insert(name);
    block_vars[block].insert(-name);
  }

  // add to our successor entry values
  for (CFGBlock::const_succ_iterator
	 it = block->succ_begin(),
	 ei = block->succ_end();
       it != ei;
       ++it) {
    bool negate_terminator = it != block->succ_begin(); // not the first means `false` branch
    Condition cond = fromComparison(static_cast<const BinaryOperator*>(block->getTerminatorCondition()), negate_terminator);
    for (VarMap::iterator jt = vars.begin(),
	   ej = vars.end();
	 jt != ej;
	 ++jt) {
      block_vars[*it].insert(jt->first);

      ZBar val = cond[jt->first];
      EqnVar* var = &system.variable(jt->first + '-' + toString((*it)->getBlockID()) + "-pre");
      if (system[*var] == NULL) {
	vector<EqnExpr*> maxArgs;
	maxArgs.push_back(&system.constant(-infinity<ZBar>()));
	system[*var] = &system.maxExpression(maxArgs);
      }
      
      EqnExpr* expr = NULL;
      if (val == -infinity<ZBar>()) {
	// don't do anything here: min(-inf, x) = -inf  (for all x)
        expr = &system.constant(-infinity<ZBar>());
      } else if (val == infinity<ZBar>()) {
	// no need to have a min here: min(inf, x) = x  (for all x)
	expr = jt->second;
      } else {
	// need a min here
	vector<EqnExpr*> minArgs;
	minArgs.push_back(&system.constant(val));
	minArgs.push_back(jt->second);
	expr = &system.expression(new Minimum<ZBar>(), minArgs);
      }

      system[*var]->arguments().push_back(expr);
    }
  }
}






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

IntervalAnalysis :: ~IntervalAnalysis() {
}

void IntervalAnalysis::runOnAllBlocks(const Decl& decl) {
  const CFG *cfg = this->context->getCFG();

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

  EqnSys system; 
  BlockVars block_vars;

  vector<EqnExpr*> infArg; 
  infArg.push_back(&system.constant(-infinity<ZBar>())); // left-most argument has to be -infinity
  infArg.push_back(&system.constant(infinity<ZBar>()));
  set<string>& function_arguments = block_vars[&cfg->getEntry()];
  string block_id = toString(cfg->getEntry().getBlockID());
  if (const FunctionDecl* func = dyn_cast<const FunctionDecl>(&decl)) {
    for (unsigned int i = func->getNumParams(); i > 0; i--) {
      string name = func->getParamDecl(i-1)->getNameAsString();

      // add the variables to the first block
      function_arguments.insert(name);
      function_arguments.insert(neg(name));
      
      // set the vars to infinity (unconstrained)
      system[system.variable(name + '-' + block_id + "-pre")] = &system.maxExpression(infArg);
      system[system.variable(neg(name) + '-' + block_id + "-pre")] = &system.maxExpression(infArg);
    }
  }

  set<const CFGBlock*> seen;
  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;
    }
    seen.insert(block);
    todo.pop_front();
    runOnBlock(block, system, block_vars);
    for (CFGBlock::const_succ_iterator it = block->succ_begin(),
                                       ei = block->succ_end();
         it != ei;
         it++ ) {
      todo.push_back(*it);
    }
  }

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

  system.indexMaxExpressions();

  Solver<ZBar> solver(system);

  for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
    EqnVar& var = system.variable(i);
    llvm::errs() << toString(var.name()) << " = " << toString(solver.solve(var)) << "\n";
  }
}


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