#include <iostream>
#include <vector>
#include <sstream>
#include "Complete.hpp"
#include "Expression.hpp"
#include "Operator.hpp"
#include "EquationSystem.hpp"
#include "MaxStrategy.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, ConcreteEquationSystem<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 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, ConcreteEquationSystem<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) {
    cout << "Usage: " << argv[0] << " filename naive|smart naive|repeat" << endl;
    exit(1);
  }


  FixpointAlgorithm<ZBar>* algorithm = NULL;
  if (!strcmp(argv[2], "naive")) {
    algorithm = new NaiveFixpointAlgorithm<ZBar>();
    cout << "Naive fixpoint" << endl;
  } else if (!strcmp(argv[2], "smart")) {
    algorithm = new SmartFixpointAlgorithm<ZBar>();
    cout << "Smart fixpoint" << endl;
  } else {
    cout << "Unknown fixpoint algorithm." << endl;
  }


  MaxStrategy<ZBar>::ImprovementOperator* naiveImprovement = new MaxStrategy<ZBar>::NaiveImprovementOperator();
  MaxStrategy<ZBar>::ImprovementOperator* improvement = NULL;
  if (!strcmp(argv[3], "repeat")) {
    improvement = new MaxStrategy<ZBar>::RepeatedImprovementOperator(*naiveImprovement);
    cout << "Repeated strategy improvement" << endl;
  } else if (!strcmp(argv[3], "naive")) {
    improvement = naiveImprovement;
    cout << "Naive strategy improvement" << endl;
  } else {
    cout << "Unknown strategy improvement algorithm." << endl;
  }
  if (!improvement || !algorithm) {
    exit(1);
  }

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

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

  cout << system << endl;
  VariableAssignment<ZBar>* result = NULL;
  MaxStrategy<ZBar> strategy(system);
  bool changed;
  do {
    if (result) delete result;
    result = algorithm->maxFixpoint(strategy);
    /*for (unsigned int i = 0, size = system.variableCount(); i < size; ++i) {
      Variable<ZBar>& var = system.variable(i);
      cout << var.name() << " = " << (*result)[var] << ", ";
    }
    cout << endl;*/
    changed = strategy.improve(*improvement, *result);
  } while(changed);

  VariableAssignment<ZBar>* rho = result;

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

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

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

  return 0;
}