#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>
std::pair<Expression<T>*, Expression<T>*> treeToExpression(pANTLR3_BASE_TREE node, EquationSystem<T>& system, bool negative=false) {
  ANTLR3_UINT32 num = node->getChildCount(node);

  string name = (char*) node->getText(node)->chars;

  // leaf node - variable
  if (num == 0) {
    if (negative) {
      vector<Expression<T>*> firstArgs;
      firstArgs.push_back(system.newExpression(system.newVariable(name + "_u")));
      vector<Expression<T>*> secondArgs;
      secondArgs.push_back(system.newExpression(system.newVariable(name + "_l")));
      return std::pair<Expression<T>*, Expression<T>*>(
          system.newExpression(new Negate<T>(),
              firstArgs),
          system.newExpression(new Negate<T>(),
              secondArgs));
    } else {
      return std::pair<Expression<T>*, Expression<T>*>(
          system.newExpression(system.newVariable(name + "_l")),
          system.newExpression(system.newVariable(name + "_u")));
    }
  }

  // a range itself
  if (name == "RANGE") {
    pANTLR3_BASE_TREE childNode;
    string childName;
    T firstValue, secondValue;

    childNode = (pANTLR3_BASE_TREE) node->getChild(node, 0);
    childName = (char*) childNode->getText(childNode)->chars;
    stringstream firstStream(childName);
    if (!(firstStream >> firstValue)) {
      throw "Invalid number.";
    }
    string firstChildName = childName;

    childNode = (pANTLR3_BASE_TREE) node->getChild(node, 1);
    childName = (char*) childNode->getText(childNode)->chars;
    stringstream secondStream(childName);
    if (!(secondStream >> secondValue)) {
      throw "Invalid number.";
    }
    string secondChildName = childName;

    if (negative) {
      T temp = firstValue;
      string tempname = firstChildName;

      firstValue = -secondValue;
      firstChildName = "-" + secondChildName;

      secondValue = -temp;
      secondChildName = "-" + tempname;
    }

    return std::pair<Expression<T>*, Expression<T>*>(
        system.newExpression(new Constant<T>(firstChildName, firstValue)),
        system.newExpression(new Constant<T>(secondChildName, secondValue)));
  }

  if (name == "-") {
    negative = true;
  }

  // subexpressions
  std::vector<Expression<T>*> firstArgs;
  std::vector<Expression<T>*> secondArgs;
  pANTLR3_BASE_TREE childNode;
  for (unsigned int i = 0; i < num; ++i) {
    childNode = (pANTLR3_BASE_TREE) node->getChild(node, i);
    std::pair<Expression<T>*, Expression<T>*> expr = treeToExpression(childNode, system, negative);
    firstArgs.push_back(expr.first);
    secondArgs.push_back(expr.second);
  }

  // other operators
  if (name == "-") {
    //if (firstArgs.size() == 1) { // secondArgs.size() == firstArgs.size()
    return std::pair<Expression<T>*, Expression<T>*>(
        firstArgs[0],
        secondArgs[0]);
  }

  if (name == "max") {
    return std::pair<Expression<T>*, Expression<T>*>(
        system.newMaxExpression(firstArgs),
        system.newMaxExpression(secondArgs));
  } else {
    // we need two because expressions delete their operator
    // bit of a silly design by me
    Operator<T>* firstOp = NULL;
    Operator<T>* secondOp = NULL;
    if (name == "min") {
      firstOp = new Minimum<T>();
      secondOp = new Minimum<T>();
    } else if (name == "+") {
      firstOp = new Addition<T>();
      secondOp = new Addition<T>();
    } else if (name == ";") {
      firstOp = new Comma<T>();
      secondOp = new Comma<T>();
    } else if (name == "GUARD") {
      firstOp = new Guard<T>();
      secondOp = new Guard<T>();
    } else {
      std::cerr << "Unknown leaf node type: " << name << std::endl;
      throw "Unknown leaf node type";
    }
    return std::pair<Expression<T>*, Expression<T>*>(
        system.newExpression(firstOp, firstArgs),
        system.newExpression(secondOp, secondArgs));
  }
}

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);

    Variable<T>* var;
    vector<Expression<T>*> args;

    string varName = (char*) varNode->getText(varNode)->chars;
    std::pair<Expression<T>*, Expression<T>*> exprs = treeToExpression(exprNode, system);

    var = system.newVariable(varName + "_l");
    args.push_back(system.newExpression(new Constant<T>("-inf", -infinity<T>())));
    args.push_back(exprs.first);
    system[*var] = system.newMaxExpression(args);
    std::cout << var->op_name << " = " << *system[*var] << std::endl;

    args.clear();
    var = system.newVariable(varName + "_u");
    args.push_back(system.newExpression(new Constant<T>("-inf", -infinity<T>())));
    args.push_back(exprs.second);
    system[*var] = system.newMaxExpression(args);
    std::cout << var->op_name << " = " << *system[*var] << std::endl;
  }
}

typedef Complete<int> ZBar;
int main (int argc, char* argv[]) {
  FixpointAlgorithm<ZBar>* algorithm = NULL;
  if (argc > 2 && !strcmp(argv[2], "naive")) {
    algorithm = new NaiveFixpoint<ZBar>();
    cout << "Naive fixpoint" << endl;
  } else {
    algorithm = new RecursiveFixpoint<ZBar>();
    cout << "Recursive fixpoint" << endl;
  }

  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(*algorithm);
  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;
}