#include <stdio.h>
#include "Fields.h"
#include "dList.h"
#include "rbTree.h"

int num_dfs = 0;

class Vertex {
  public:
    string name;
    rbTree<string> edges;
    int visited;
  public:
    Vertex(string n): name(n), visited(0) {}
};

class Edge {
  public:
    string name;   // Names are "v1->name v2->name"
    Vertex *v1;
    Vertex *v2;
    double cap;
    Edge *backedge;
    Edge *fgedge;
  public:
    Edge(Vertex *from, Vertex *to): v1(from), v2(to), cap(0.0),
          backedge(0), fgedge(0) {
      name = v1->name;
      name += " ";
      name += v2->name;
    }
    ~Edge() {
      if (v1->edges.find(v2->name))
        v1->edges.deleteNode();
      if (backedge != 0) { backedge->backedge = 0; }
    }
}; 

class Graph {
  public:
    double flow;
    double maxcap;
    rbTree<string> vertices;
    Vertex *source;
    Vertex *sink;

  public:
    Graph(): source(0), sink(0), maxcap(0.0) {}
    Vertex *get_vertex(string name);
    Edge *get_edge(Vertex *v1, Vertex *v2);
    void print_graph();
    void read_graph(string fname);
};
  
void Graph::print_graph()
{
  Vertex *v;
  Edge *e;

  printf("  Source %s, Sink %s", source->name.c_str(),
                                 sink->name.c_str());
  if (flow >= 0) printf(", Flow: %8.3lf", flow); else printf(", No Flow");
  printf("\n\n");

  for (vertices.first(); !vertices.endOfList(); vertices.next()) {
    v = (Vertex *) vertices.getVal().v;
    printf("  Vertex %s\n", v->name.c_str());

    rbTree<string> *edges = &(v->edges);
    for (edges->first(); !edges->endOfList(); edges->next()) {
      e = (Edge *) edges->getVal().v;
      printf("    Edge to %s.  Capacity %7.3lf\n", e->v2->name.c_str(), e->cap);
    }
    printf("\n");
  }
}

void Graph::read_graph(string fname)
{

  Fields *f;
  int nf;
  double d;
  Edge *e;
  Vertex *v1, *v2;
  string s;

  f = new Fields(fname);

  while (f->get_line() >= 0) {
    nf = f->get_NF();
    if (nf > 0) {
      s = f->get_field(0);
      if (s == "EDGE") {
        if (nf != 4 || sscanf(f->get_field(3).c_str(), "%lf", &d) != 1 || 
            d <= 0) {
          fprintf(stderr, "%s %d: EDGE v1 v2 cap (cap > 0)\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        v1 = get_vertex(f->get_field(1));
        v2 = get_vertex(f->get_field(2));
        e = get_edge(v1, v2);
        if (e->cap > 0) {
          fprintf(stderr, "%s %d: Edge already exists\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        e->cap = d;
        if (d > maxcap) maxcap = d;
        if (e->backedge == 0) {               /* Create the backedge */
          e->backedge = get_edge(v2, v1);
          e->backedge->backedge = e;
        }
      } else if (s == "SOURCE") {
        if (nf != 2) {
          fprintf(stderr, "%s %d: SOURCE v\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        if (source != 0) {
          fprintf(stderr, "%s %d: Duplicate Source\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        source = get_vertex(f->get_field(1));
      } else if (s == "SINK") {
        if (nf != 2) {
          fprintf(stderr, "%s %d: SINK v\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        if (sink != 0) {
          fprintf(stderr, "%s %d: Duplicate Sink\n", 
                          f->get_name().c_str(), f->get_line_number());
          exit(1);
        }
        sink = get_vertex(f->get_field(1));
      } else {
        fprintf(stderr, "%s %d: Bad line\n", 
                        f->get_name().c_str(), f->get_line_number());
        exit(1);
      }
    }
  }
        
  if (source == 0) {
    fprintf(stderr, "No source specified\n");
    exit(1);
  } else if (sink == 0) {
    fprintf(stderr, "No sink specified\n");
    exit(1);
  }

}
  
Vertex *Graph::get_vertex(string name)
{
  Vertex *v;

  if (!vertices.find(name)) {
    v = new Vertex(name);
    vertices.insert(name, new_jval_v(v));
  } else {
    v = (Vertex *) vertices.getVal().v;
  }
  return v;
}

Edge *Graph::get_edge(Vertex *v1, Vertex *v2)
{
  Edge *e;

  if (!(v1->edges.find(v2->name))) {
    e = new Edge(v1, v2);
    v1->edges.insert(v2->name, new_jval_v(e));
  } else {
    e = (Edge *) v1->edges.getVal().v;
  }
  return e;
}

bool visit(Vertex *v, Vertex *sink, dList<Edge *> *path)
{
  Edge *e;

  if (v->visited == num_dfs) return false;
  if (v == sink) return true;
  v->visited = num_dfs;

  rbTree<string> *edges = &(v->edges);
  for (edges->first(); !edges->endOfList(); edges->next()) {  
    e = (Edge *) edges->getVal().v;
    if (e->cap > 0.0 && visit(e->v2, sink, path)) {
      path->prepend(e);
      return true;
    }
  }
  return false;
}

main(int argc, char **argv)
{
  Graph *g;
  Graph *fg; // Flow graph
  Vertex *v, *v1, *v2;
  dList<Edge *> *path;
  Edge *e;
  bool ok;
  double mincap;
  int print;

  if (argc != 3) {
    fprintf(stderr, "usage: augpath1 filename print(yes|no)\n");
    exit(1);
  }

  if (strcmp(argv[2], "yes") == 0) {
    print = 1;
  } else if (strcmp(argv[2], "no") == 0) {
    print = 0;
  } else {
    fprintf(stderr, "usage: augpath1 filename print(yes|no)\n");
    exit(1);
  }

  g = new Graph();
  g->read_graph(argv[1]);

  fg = new Graph();
  g->flow = -1;
  
  /* Make empty flow graph */

  for (g->vertices.first(); !g->vertices.endOfList(); g->vertices.next()) {
    v = (Vertex *) g->vertices.getVal().v;
    v1 = fg->get_vertex(v->name);
    rbTree<string> *edges = &(v->edges);
    for (edges->first(); !edges->endOfList(); edges->next()) {
      e = (Edge *) edges->getVal().v;
      v2 = fg->get_vertex(e->v2->name);
      e->fgedge = fg->get_edge(v1, v2);
    }
  }
  fg->source = fg->get_vertex(g->source->name);
  fg->sink = fg->get_vertex(g->sink->name);
  fg->flow = 0.0;

  if (print) {
    printf("Original Graph:\n");
    g->print_graph();
    printf("\n");
    fg->print_graph();
  }
  
  ok = true;
  while (ok) {
    num_dfs++;

    /* Find a path */

    path = new dList<Edge *>;
    ok = visit(g->source, g->sink, path);
    mincap = g->maxcap;

    /* If a path is found */
    if (ok) {
      if (print) {
        printf("---------------------------------------------------------\n");
        printf("Path found -- %s", g->source->name.c_str());
      }

      /* Find the maximum capacity of the path */
      for (path->first(); !path->endOfList(); path->next()) {
        e = path->get();
        if (e->cap < mincap) mincap = e->cap;
        if (print) printf(" -> %s", e->v2->name.c_str());
      }
      if (print) printf(" -- cap = %.3lf\n", mincap);

      fg->flow += mincap;

      /* Update the residual graph. For each edge on the path, reduce the 
           capacity along that edge by 
           mincap, and then add that capacity to the back edge */

      for (path->first(); !path->endOfList(); path->next()) {
        e = path->get();
        e->cap -= mincap;
        e->backedge->cap += mincap;        /* Note, now we don't need to find
                                              the back edge */
        /* Add flow to the flow graph */
        e->fgedge->cap += mincap;
      }
      if (print) {
        printf("Residual Graph:\n");
        g->print_graph();
        printf("\n");
        printf("Flow Graph:\n");
        fg->print_graph();
        printf("\n");
      }
    }
    delete path;
  }
  printf("Total flow: %.3lf\n", fg->flow);
}