#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 "VariableAssignment.hpp"

extern "C" {
#include "parser/EquationSystemParser.h"
#include "parser/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
  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 == "+" || name == "add") {
      op = new Addition<T>();
    } else if (name == "-") {
      if (args.size() == 1) {
        op = new Negation<T>();
      } else {
        op = new Subtraction<T>();
      }
    } else if (name == "sub") {
      op = new Subtraction<T>();
    } else if (name == "*" || name == "mult") {
      op = new Multiplication<T>();
    } else if (name == ";") {
      op = new Comma<T>();
    } else if (name == "GUARD" || name == "guard") {
      op = new Guard<T>();
    } else {
      std::cout << "throw exception" << *(char*)NULL;
      //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) {
    std::cout << "throw exception" << *(char*)NULL;
    //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[]) {
  map<string,log::Logger*> loggers;
  loggers["strategy"] = &log::strategy;
  loggers["fixpoint"] = &log::fixpoint;
  loggers["debug"] = &log::debug;

  set<string> variables;
  if (argc > 2) {
    int i = 2;
    while (argv[i] != NULL) {
      if (string(argv[i]) == "-v") {
        ++i;
        if (i < argc) {
          char* arg = argv[i];
          char* str = strtok(arg, ",");
          do {
            if (str && loggers[str])
              loggers[str]->enabled(true);
          } while ((str = strtok(NULL, ",")) != NULL);
        }
      } else {
        variables.insert(argv[i]);
      }
      ++i;
    }
  }

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

  EquationSystemParser_equation_system_return ret = parser -> equation_system(parser);

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

  log::debug << system << endl;
  system.indexMaxExpressions(); // make reverse-lookup O(1) instead of O(n)

  IdMap<Variable<ZBar>,ZBar> result = solve_for(system);
  if (variables.size() > 0) {
    for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
      Variable<ZBar>& var = system.variable(i);
      if (variables.find(var.name()) != variables.end())
        cout << var.name() << " = " << result[var] << endl;
    }
  } else {
    for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
      Variable<ZBar>& var = system.variable(i);
      cout << var.name() << " = " << result[var] << endl;
    }
  }
  /*
  DynamicMaxStrategy<ZBar> strategy(system);
  DynamicVariableAssignment<ZBar> rho(system, strategy);
  strategy.setRho(rho);

  if (variables.size() > 0) {
    for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
      Variable<ZBar>& var = system.variable(i);
      if (variables.find(var.name()) != variables.end()) {
        cout << var.name() << " = " << rho[var].asKnown() << endl;
      }
    }
  } else {
    for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
      Variable<ZBar>& var = system.variable(i);
      cout << var.name() << " = " << rho[var].asKnown() << endl;
    }
  }
  */

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

  return 0;
}