#include <iostream>
#include <vector>
#include <map>
#include <string>
#include <sstream>
#include <math.h>

#include "Operator.hpp"
#include "Expression.hpp"
#include "MaxStrategy.hpp"
#include "EquationSystem.hpp"
#include "Complete.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;

  // leaf node - constant or variable
  if (num == 0) {
    if (name == "inf") {
      return system.newExpression(new Constant<T>(infinity<T>()), vector<Expression<T>*>());
    } else {
      stringstream stream(name);
      T output;
      if (stream >> output) {
        return system.newExpression(new Constant<T>(output), vector<Expression<T>*>());
      } else {
        return system.newExpression(system.newVariable(name), vector<Expression<T>*>());
      }
    }
  }

  // other operators
  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.newMaxExpression(args);
  } else {
    Operator<T>* op = NULL;
    if (name == "min") {
      op = new Minimum<T>();
    } else if (name == "+") {
      op = new Addition<T>();
    } else if (name == "-") {
      op = new Subtraction<T>();
    } else if (name == ";") {
      op = new Comma<T>();
    } else if (name == "GUARD") {
      op = new Guard<T>();
    } else {
      std::cerr << "Unknown leaf node type: " << name << std::endl;
      throw "Unknown leaf node type";
    }
    return system.newExpression(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.newVariable(varName);

    system[*var] = treeToExpression(exprNode, system);
  }
}

typedef Complete<int> ZBar;
int main (int argc, char* argv[]) {
  pANTLR3_INPUT_STREAM input;
  pEquationSystemLexer lex;
  pANTLR3_COMMON_TOKEN_STREAM tokens;
  pEquationSystemParser parser;

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

  EquationSystemParser_equation_system_return ret = parser -> equation_system(parser);

  EquationSystem<ZBar> system;
  treeToSystem<ZBar>(ret.tree, system);
  // DO MORE HERE

  VariableAssignment<ZBar> rho = system.minFixpoint(NaiveFixpoint<ZBar>());
  std::cout << rho << std::endl;
  rho = system.minFixpoint(RecursiveFixpoint<ZBar>());
  std::cout << rho << std::endl;

  const std::vector<Variable<ZBar>*> vars = system.vars();
  for (unsigned int i = 0, size = vars.size(); i < size; ++i) {
    cout << vars[i]->name() << " = " << rho[*vars[i]] << endl;
  }

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

  return 0;
}