#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/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!



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 scalar_mult(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;
}

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) {
  assert(false); // unary operations aren't supported. Sorry!
  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(scalar_mult(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;
      }
    }
  }
  assert(false); // Unknown binary operation. Only assignment, addition, subtraction and multiplication are permitted
  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 */

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


vector<Block> splitBlock(const CFGBlock* block) {
  vector<Block> subBlocks;

  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 != "") {
      if (subBlocks.empty())
        subBlocks.push_back(Block());
      Block& subBlock = subBlocks.back();

      bool make_new = subBlock.find(name) != subBlock.end();
      if (!make_new) {
        for (Vector::iterator
               it = result.first.begin(),
               ei = result.first.end();
             it != ei;
             ++it) {
          if (subBlock.find(it->first) != subBlock.end()) {
            make_new = true;
            break;
          }
        }
      }
      if (make_new) {
        Block newBlock;
        newBlock[name] = result;
        subBlocks.push_back(newBlock);
      } else {
        subBlock[name] = result;
      }
    }
  }

  return subBlocks;
}

typedef map<string, unsigned int> Counters;

string var_name(const int& i, const string& block_id) {
  return toString(i) + '-' + block_id;
}

void runOnBlock(const ConstraintMatrix& T, const CFGBlock* block, EqnSys& system) {

  int size = T.size();
  string block_id = toString(block->getBlockID());
  vector<unsigned int> counters(size);
  vector<EqnVar*> vars(size);

  vector<EqnExpr*> infArgs;
  infArgs.push_back(&system.constant(-infinity<ZBar>()));
  infArgs.push_back(&system.constant(infinity<ZBar>()));
  for (int i = 0; i < size; ++i) {
    vars[i] = &system.variable(var_name(i, block_id)+"-pre");
    if (system[*vars[i]] == NULL) {
      system[*vars[i]] = &system.maxExpression(infArgs);
    }
  }

  vector<Block> subBlocks = splitBlock(block);

  map<string,int> mcf_indices; // 1-indexed, for the mcf stuff
  int mcf_vert_count = 1;
  vector<pair<int,int> > mcf_edges;
  {
    for (int j = 0; j < size; ++j) {
      string from = "";
      string to = "";
      for (map<string,int>::const_iterator
             jt = T[j].begin(),
             ej = T[j].end();
           jt != ej;
           ++jt) {
        if (mcf_indices.find(jt->first) == mcf_indices.end()) {
          mcf_indices[jt->first] = ++mcf_vert_count;
        }
        assert(jt->second == 1 || jt->second == -1);
        assert(jt->second != 1 || from == "");
        assert(jt->second != -1 || to == "");
        if (jt->second == 1)
          from = jt->first;
        if (jt->second == -1)
          to = jt->first;
      }
      int source = (from == "" ? 1 : mcf_indices[from]);
      int dest = (to == "" ? 1 : mcf_indices[to]);
      
      mcf_edges.push_back(pair<int,int>(source,dest));
    }
  }

  

  for (int blockIndex = 0, blockSize = subBlocks.size(); blockIndex < blockSize; ++blockIndex) {
    Block& subBlock = subBlocks[blockIndex];

    vector<EqnVar*> nextVars(size);
    for (int i = 0; i < size; ++i) {
      const map<string,int>& tk = T[i];

      ZBar b = 0;
      map<string,int> indices;
      for (map<string,int>::const_iterator
             tk_it = tk.begin(),
             tk_end = tk.end();
           tk_it != tk_end;
           tk_it++) {
        b += ZBar(tk_it->second) * subBlock[tk_it->first].second;
      }

      vector<ZBar> g(mcf_vert_count);
      for (map<string,int>::const_iterator
             tk_it = tk.begin(),
             tk_end = tk.end();
           tk_it != tk_end;
           tk_it++) { // each column of Tk
        const map<string,ZBar>& A = subBlock[tk_it->first].first;
        for (map<string,ZBar>::const_iterator
               A_it = A.begin(),
               A_end = A.end();
             A_it != A_end;
             A_it++) { // each row of A
          ZBar value = A_it->second * ZBar(tk_it->second);
          g[mcf_indices[A_it->first]-1] += value;
          g[0] -= value;
        }
      }

      vector<EqnExpr*> args(size);
      for (int j = 0; j < size; ++j) {
        assert(vars[j] != NULL);
        args[j] = vars[j];
      }
      
      EqnExpr* minCostExpr = &system.expression(new MinCostFlow<ZBar>(g, mcf_edges), args);

      EqnExpr* expression;
      if (b == 0) {
        expression = minCostExpr;
      } else {
        // add the constant factor to the min-cost bit
        vector<EqnExpr*> additionArgs;
        additionArgs.push_back(&system.constant(b));
        additionArgs.push_back(minCostExpr);
        expression = &system.expression(new Addition<ZBar>(), additionArgs);
      }

      string count = toString(counters[i]);
      counters[i]++;
      EqnVar* newVar = &system.variable(var_name(i, block_id)+'-'+count);
      
      // make it max(-inf, expr), so strategy iteration will work
      vector<EqnExpr*> maxArgs;
      maxArgs.push_back(&system.constant(-infinity<ZBar>()));
      maxArgs.push_back(expression);
      system[*newVar] = &system.maxExpression(maxArgs);
      nextVars[i] = newVar;
    }
    vars = nextVars; // update our variable listings
  }

  // 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 (int i = 0; i < size; ++i) {
      EqnVar* var = &system.variable(var_name(i, 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 (T[i].size() == 1) { // only one value in this row, so it's okay for a test
        ZBar val = infinity<ZBar>();
        Vector::iterator cond_finder;
        if (T[i].begin()->second < 0)
          cond_finder = cond.find('-' + T[i].begin()->first);
        else
          cond_finder = cond.find(T[i].begin()->first);
        if (cond_finder != cond.end()) {
          val = cond_finder->second;
          //val *= ; // potentially change it's sign
        }
        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 = vars[i];
        } else {
          // need a min here
          vector<EqnExpr*> minArgs;
          minArgs.push_back(&system.constant(val));
          minArgs.push_back(vars[i]);
          expr = &system.expression(new Minimum<ZBar>(), minArgs);
        }
      } else { // more than one value in this row, so leave it for now...
        expr = vars[i];
      }
      system[*var]->arguments().push_back(expr);
    }
  }
}






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

IntervalAnalysis :: ~IntervalAnalysis() {
}

map<CFGBlock*,vector<ZBar> > IntervalAnalysis::runOnAllBlocks(const Decl& decl, const ConstraintMatrix& T) {
  const CFG *cfg = this->context->getCFG();

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

  EqnSys system;

  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(T, block, system);
    for (CFGBlock::const_succ_iterator it = block->succ_begin(),
           ei = block->succ_end();
         it != ei;
         it++ ) {
      todo.push_back(*it);
    }
  }

  system.indexMaxExpressions();

  Solver<ZBar> solver(system);

  map<CFGBlock*, vector<ZBar> > resultMap;

  for (int i = 0; i < system.variableCount(); ++i) {
    solver.solve(system.variable(i));
  }

  for (CFG::const_iterator
         it = cfg->begin(),
         ei = cfg->end();
       it != ei;
       ++it) {
    vector<ZBar>& vector = resultMap[*it];
    string block_id = toString((*it)->getBlockID());
    for (int i = 0, size = T.size(); i < size; ++i) {
      EqnVar& var = system.variable(var_name(i, block_id) + "-pre");
      vector.push_back(solver.solve(var));
    }
  }

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

  return resultMap;
}


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