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

using namespace clang;

typedef EquationSystem<Complete<int> > EqnSys;
typedef Expression<Complete<int> > EqnExpr;

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

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;
typedef std::pair<Vector, ZBar> Result; // a "slice" of an equation

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<class T>
T negate(const T& v) {
  return -v;
}
template<class T>
T addValues(const T& l, const T& r) {
  return l + r;
}
template<class T>
T subValues(const T& l, const T& r) {
  return l - r;
}


Result fromStmt(const Stmt*, EqnSys&);

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

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

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


Addition<Complete<int> >* add = new Addition<Complete<int> >();
Subtraction<Complete<int> >* sub = new Subtraction<Complete<int> >();
Multiplication<Complete<int> >* mul = new Multiplication<Complete<int> >();

Result fromBinary(const BinaryOperator* op, EqnSys& system) {
  Result left = fromStmt(op->getLHS()->IgnoreParenCasts(), system);
  Result right = fromStmt(op->getRHS()->IgnoreParenCasts(), system);
  
  switch (op->getOpcode()) {
  case BO_Sub:
    transform_values(negate<ZBar>, right.first);
  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;
      }
      for (Vector::iterator it = value.first.begin(),
             ei = value.first.end();
           it != ei;
           ++it) {
        it->second *= scalar;
      }
      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, EqnSys& system) {
  /*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);
      Variable<Complete<int> > var = system.variable(decl->getNameAsString());
      std::vector<EqnExpr*> args;
      args.push_back(&system.constant(-infinity<Complete<int> >()));
      args.push_back(fromStmt(decl->getInit(), system));
      system[var] = &system.maxExpression(args);
    }
  }*/

  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(), system);
      llvm::errs() << "<{ ";
      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->first) << " = " << toString(it->second);
      }
      llvm::errs() << " }, " << toString(expr.second) << ">\n";
    }
  }

  return Result();
}

Result fromAssignment(const BinaryOperator* op, EqnSys& system) {
  /*EqnExpr* left = fromStmt(op->getLHS()->IgnoreParenCasts(), system);
  EqnExpr* right = fromStmt(op->getRHS()->IgnoreParenCasts(), system);
  Variable<Complete<int> >* var = static_cast<Variable<Complete<int> >*>(left);

  std::vector<EqnExpr*> args;
  args.push_back(&system.constant(-infinity<Complete<int> >()));
  args.push_back(right);
  if (system[*var] != NULL)
    args.push_back(system[*var]);
    system[*var] = &system.maxExpression(args);*/
  
  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(), system);
    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->first) << " = " << toString(it->second);
    }
    llvm::errs() << " }, " << toString(expr.second) << ">\n";
    return expr;
  }
  return Result();
}

Result fromStmt(const Stmt* stmt, EqnSys& system) {
  if (!stmt)
    return Result();
  switch (stmt->getStmtClass()) {
  case Stmt::IntegerLiteralClass:
    return fromInteger(static_cast<const IntegerLiteral*>(stmt), system);
  case Stmt::DeclRefExprClass:
    return fromDeclExpr(static_cast<const DeclRefExpr*>(stmt), system);
  case Stmt::UnaryOperatorClass:
    return fromUnary(static_cast<const UnaryOperator*>(stmt), system);
  case Stmt::DeclStmtClass:
    return fromDeclStmt(static_cast<const DeclStmt*>(stmt), system);
  case Stmt::BinaryOperatorClass:
    {
      const BinaryOperator* binop = static_cast<const BinaryOperator*>(stmt);
      if (binop->isAssignmentOp())
        return fromAssignment(binop, system);
      else
        return fromBinary(binop, system);
    }
  }
  return Result();
}

void runOnBlock(std::string id, const CFGBlock* block, EqnSys& system) {
  for (CFGBlock::const_iterator it = block->begin(),
                                ei = block->end();
       it != ei;
       ++it) {
    const CFGStmt* cfg_stmt = it->getAs<CFGStmt>();
    Result expr = fromStmt(cfg_stmt->getStmt(), system);
    /*llvm::errs() << "<{ ";
    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->first) << " = " << toString(it->second);
    }
    llvm::errs() << "}, " << toString(expr.second) << ">\n";*/
  }
  fromStmt(block->getTerminatorCondition(), system);

  /*if (terminator.getStmt() != NULL) {
    if (terminator.getStmt()->getStmtClass() == Stmt::IfStmtClass) {
      const IfStmt* if_stmt = static_cast<const IfStmt*>(terminator.getStmt());
      llvm::errs() << "If: \n";
      if_stmt->dump();
    } else {
      llvm::errs() << "\n";
      terminator.getStmt()->dump();
    }
    }*/
  return; // TODO: return a generated expression
}

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() << (void*)block << "\n";
    seen.insert(block);
    todo.pop_front();
    runOnBlock(toString(block), block, system);
    llvm::errs() << "-> ";
    for (CFGBlock::const_succ_iterator it = block->succ_begin(),
                                       ei = block->succ_end();
         it != ei;
         it++ ) {
      llvm::errs() << (void*) *it << ", ";
      todo.push_back(*it);
    }
    llvm::errs() << "\n\n";
  }

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

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

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

  for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
    Variable<Complete<int> >& 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; }