#include "Log.hpp"
#include <iostream>
#include <vector>
#include <sstream>
#include "Complete.hpp"
#include "Expression.hpp"
#include "Operator.hpp"
#include "EquationSystem.hpp"
#include "MaxStrategy.hpp"
#include "ImprovementOperator.hpp"
#include "FixpointAlgorithm.hpp"

extern "C" {
#include "EquationSystemParser.h"
#include "EquationSystemLexer.h"
}

using namespace std;

template<typename T>
Expression<T>& treeToExpression(pANTLR3_BASE_TREE node, EquationSystem<T>& system) {
  ANTLR3_UINT32 num = node->getChildCount(node);
  string name = (char*) node->getText(node)->chars;

  if (num == 0) {
    // leaf node -> constant or variable
    if (name == "inf") {
      return system.constant(infinity<T>());
    } else {
      stringstream stream(name);
      T output;
      if (stream >> output) {
        return system.constant(output);
      } else {
        return system.variable(name);
      }
    }
  }

  // anything that's not a constant/variable
  std::vector<Expression<T>*> args;
  pANTLR3_BASE_TREE childNode;
  for (unsigned int i = 0; i < num; ++i) {
    childNode = (pANTLR3_BASE_TREE) node->getChild(node, i);
    args.push_back(&treeToExpression(childNode, system));
  }

  if (name == "max") {
    return system.maxExpression(args);
  } else {
    Operator<T>* op = NULL;
    if (name == "min") {
      op = new Minimum<T>();
    } else if (name == "+") {
      op = new Addition<T>();
    } else if (name == "-") {
      if (args.size() == 1) {
        op = new Negation<T>();
      } else {
        op = new Subtraction<T>();
      }
    } else if (name == "*") {
      op = new Multiplication<T>();
    } else if (name == ";") {
      op = new Comma<T>();
    } else if (name == "GUARD") {
      op = new Guard<T>();
    } else {
      throw "Parse error: Unknown operator";
    }
    return system.expression(op, args);
  }
}

template<typename T>
void treeToSystem(pANTLR3_BASE_TREE node, EquationSystem<T>& system) {
  ANTLR3_UINT32 num = node->getChildCount(node);

  if (num % 2 == 1)
    throw "Big problem here.";

  pANTLR3_BASE_TREE varNode;
  pANTLR3_BASE_TREE exprNode;
  for (unsigned int i = 0; i < num; i += 2) {
    varNode = (pANTLR3_BASE_TREE) node->getChild(node, i);
    exprNode = (pANTLR3_BASE_TREE) node->getChild(node, i+1);

    string varName = (char*) varNode->getText(varNode)->chars;
    Variable<T>& var = system.variable(varName);

    vector<Expression<T>*> args;
    args.push_back(&system.constant(-infinity<T>()));
    args.push_back(&treeToExpression(exprNode, system));
    system[var] = &system.maxExpression(args);
  }
}

typedef Complete<int> ZBar;
int main (int argc, char* argv[]) {
  if (argc <= 3) {
    cerr << "Usage: " << argv[0] << " filename fixpoint strategy [log]" << endl
         << "  fixpoint: naive | smart" << endl
         << "  strategy: naive | repeat | smart" << endl
         << "  log: section[,section[,section[...]]]" << endl;
    exit(1);
  }

  map<string,log::Logger*> loggers;
  loggers["info"] = &log::info;
  loggers["trace"] = &log::trace;
  loggers["strategy"] = &log::strategy;
  loggers["fixpoint"] = &log::fixpoint;
  loggers["debug"] = &log::debug;

  if (argc > 3) {
    char* arg = argv[4];
    char* str = strtok(arg, ",");
    do {
      if (str && loggers[str])
        loggers[str]->enabled(true);
    } while ((str = strtok(NULL, ",")) != NULL);
  }

  auto input = antlr3FileStreamNew((pANTLR3_UINT8)argv[1], ANTLR3_ENC_8BIT);
  auto lex = EquationSystemLexerNew(input);
  auto tokens = antlr3CommonTokenStreamSourceNew(ANTLR3_SIZE_HINT, TOKENSOURCE(lex));
  auto parser = EquationSystemParserNew(tokens);

  auto ret = parser -> equation_system(parser);

  EquationSystem<ZBar> system;
  treeToSystem<ZBar>(ret.tree, system);

  // PARSE ARGUMENTS - fixpoint
  FixpointAlgorithm<ZBar>* algorithm = NULL;
  if (!strcmp(argv[2], "naive")) {
    algorithm = new NaiveFixpointAlgorithm<ZBar>(system);
    log::info << "Naive fixpoint" << endl;
  } else if (!strcmp(argv[2], "smart")) {
    algorithm = new SmartFixpointAlgorithm<ZBar>(system);
    log::info << "Smart fixpoint" << endl;
  } else {
    cerr << "Unknown fixpoint algorithm." << endl;
  }

  // PARSE ARGUMENTS - strat improvement
  ImprovementOperator<ZBar>* naiveImprovement = new NaiveImprovementOperator<ZBar>();
  ImprovementOperator<ZBar>* improvement = NULL;
  if (!strcmp(argv[3], "repeat")) {
    improvement = new RepeatedImprovementOperator<ZBar>(*naiveImprovement);
    log::info << "Repeated strategy improvement" << endl;
  } else if (!strcmp(argv[3], "naive")) {
    improvement = naiveImprovement;
    log::info << "Naive strategy improvement" << endl;
  } else if (!strcmp(argv[3], "smart")) {
    improvement = new SmartImprovementOperator<ZBar>(system);
    log::info << "Smart strategy improvement" << endl;
  } else {
    cerr << "Unknown strategy improvement algorithm." << endl;
  }
  if (!improvement || !algorithm) {
    exit(1);
  }

  log::debug << system << std::endl;
  system.indexMaxExpressions(); // make reverse-lookup O(1) instead of O(n)
  StableVariableAssignment<ZBar> result(system.variableCount(), infinity<ZBar>());
  ConcreteMaxStrategy<ZBar> strategy(system);
  IdSet<Variable<ZBar>> s1(system.variableCount());
  IdSet<Variable<ZBar>> s2(system.variableCount());
  bool changed = false;
  s2.clear();
  s2.invert();
  do {
    s1.clear();
    algorithm->maxFixpoint(system, strategy, result, s2, s1);

    log::debug << result << std::endl;

    s2.clear();
    changed = improvement->improve(system, strategy, result, s1, s2);
    log::debug << "Changed: " << (changed ? "true" : "false") << std::endl;
  } while(changed);

  for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
    Variable<ZBar>& var = system.variable(i);
    cout << var.name() << " = " << result[var] << endl;
  }

  delete algorithm;
  if (naiveImprovement != improvement)
    delete improvement;
  delete naiveImprovement;

  parser -> free(parser);
  tokens -> free(tokens);
  lex -> free(lex);
  input -> close(input);

  return 0;
}