#include <iostream.h>
#include <stdlib.h>
#include <vector>
#include <string.h>
#include <sys/types.h>
#include <time.h>

char* color_table[] = {"red", "blue", "green", "cyan", "magenta"};
int max_colors = 5;

int random_int(int range)
{
  return rand()%range;
}

char* random_color(int max)
{
  return color_table[random_int(max)];
}

int getBlockCount(float prob)
{

  int count = 1;
  while (drand48() > prob)
    count++;
  return count;
}

int getColorCount(float prob)
{
  return 3;
  /*
  int count = 1;
  while (drand48() > prob)
    count++;
  return count;
  */
}

int color_count;
int block_count;

int bc_count;
char* *bc_vec;

void printDomain();
void printActions();
void printProblem();
void printInitStack();
void printInitColors();
void printGoal();
void printRecursiveStack(int block_index, int tower_index);

void printDomain()
{
  cout << "(define" << endl
       << " (domain blocks-domain)" << endl
       << " (:requirements :negative-preconditions" << endl
       << "  :universal-preconditions :existential-preconditions" << endl
       << "  :disjunctive-preconditions :probabilistic-effects" << endl
       << "  :conditional-effects :rewards :equality)" << endl;

  cout << " (:predicates" << endl;
  for (int i=0; i<color_count; i++)
    cout << "  (is-" << color_table[i] << " ?block)" << endl;
  cout << "  (holding ?block)" << endl
       << "  (is-table ?block)" << endl
       << "  (on-top-of ?block1 ?block2)" << endl
       << " )" << endl;

  printActions();

  cout << ")" << endl;
}

void printActions()
{
  cout << " (:action pick-up-block-from" << endl
       << "  :parameters (?block_top ?block_bottom)" << endl
       << "  :precondition (and (not (= ?block_top ?block_bottom))" << endl
       << "                     (forall (?b) (not (holding ?b)))" << endl
       << "                     (on-top-of ?block_top ?block_bottom)" << endl
       << "                     (forall (?b2) (not (on-top-of ?b2 ?block_top))))" << endl
       << "  :effect (probabilistic 0.75 (and (holding ?block_top)" << endl
       << "                                   (not (on-top-of ?block_top ?block_bottom)))" << endl
       << "                         0.25 (when (not (= ?block_bottom table))" << endl
       << "                                    (and (not (on-top-of ?block_top ?block_bottom))" << endl
       << "                                         (on-top-of ?block_top table))))" << endl
       << " )" << endl;

  cout << " (:action put-down-block-on" << endl
       << "  :parameters (?block_top ?block_bottom)" << endl
       << "  :precondition (and (not (= ?block_top ?block_bottom))" << endl
       << "                     (holding ?block_top)" << endl
       << "                     (not (holding ?block_bottom))" << endl
       << "                     (or (is-table ?block_bottom)" << endl
       << "                         (forall (?b) (not (on-top-of ?b ?block_bottom)))))" << endl
       << "  :effect (probabilistic 0.75 (and (not (holding ?block_top))" << endl
       << "                                   (on-top-of ?block_top ?block_bottom))" << endl
       << "                         0.25 (and (not (holding ?block_top))" << endl
       << "                                   (on-top-of ?block_top table)))" << endl
       << " )" << endl;
}

void printProblem()
{
  cout << "(define" << endl
       << " (problem blocks-problem)" << endl
       << " (:domain blocks-domain)" << endl
       << " (:objects";
  for (int i=0; i<block_count; i++)
    cout << " block" << i;
  cout << ")" << endl;
  cout << " (:init" << endl
       << "  (is-table table)" << endl;
  printInitStack();
  printInitColors();
  cout << " )" << endl;
  printGoal();
  cout << ")" << endl;
}

void printInitStack()
{
  for (int block_index = 0; block_index < block_count;) {
    int tower_size = random_int(block_count-block_index)+1;
    for (int i=0; i<tower_size-2; i++) {
      cout << "  (on-top-of block" << block_index
           << " block" << block_index+1 << ")" << endl;
      block_index++;
    }
    cout << "  (on-top-of block" << block_index << " table)" << endl;
    block_index++;
  }
}

void printInitColors()
{
  bc_count = block_count;
  bc_vec = new char*[bc_count];
  for (int i=0; i<block_count; i++)
    bc_vec[i] = random_color(color_count);
  for (int i=0; i<block_count; i++)
    cout << "  (is-" << bc_vec[i] << " block" << i << ")" << endl;
}

// (exists (?f1) (and (is-red ?f1) next)

void printGoal()
{
  cout << " (:goal";
  printRecursiveStack(0, random_int(block_count)+1);
  cout << ")" << endl;
}

void printRecursiveStack(int block_index, int tower_index)
{
  if (block_index == block_count) {
    cout << " (on-top-of ?fb" << block_index-1 << " table)";
    return;
  }

  cout << " (exists (?fb" << block_index << ") (and";


  int color_index = random_int(bc_count);
  char *color = bc_vec[color_index];
  cout << " (is-" << color << " ?fb" << block_index << ")";
  for (int i=color_index; i<bc_count-1; i++)
    bc_vec[i] = bc_vec[i+1];
  bc_count--;

  if (block_index) {
    cout << " (on-top-of ?fb" << block_index-1;
    if (tower_index)
      cout << " ?fb" << block_index;
    else
      cout << " table";
    cout << ")";
  }

  if (!tower_index)
    tower_index = random_int(block_count-block_index)+2;

  printRecursiveStack(block_index+1, tower_index-1);
  cout << "))";
}

void usage(char *name) 
{
  cout << "usage: " << name << " [-b block-count] [-c color-count]" << endl;
}

int main(int argc, char **argv)
{
  srand48(time(0));
  srand(time(0));

  color_count = -1;
  block_count = -1;
  
  for (int i=1; i<argc; i++) {
    if (!strcmp(argv[i], "-b")) {
      if (i==argc-1) {
	usage(*argv);
	return 1;
      }
      block_count = atoi(argv[++i]);
    }
    if (!strcmp(argv[i], "-c")) {
      if (i==argc-1) {
	usage(*argv);
	return 1;
      }
      color_count = atoi(argv[++i]);
    }
    if (!strcmp(argv[i], "-?")) {
      usage(*argv);
      return 0;
    }
  }

  if (color_count == -1)
    color_count = getColorCount(0.5);
  if (block_count == -1)
    block_count = getBlockCount(0.1);

  printDomain();
  printProblem();

  return 0;

}