#include "port.h"

#include <stdio.h>
#ifndef UNIX
#include <io.h>
#endif
#include <fcntl.h>

void write_ppm( FILE* f, int width, int height, unsigned char* data );
void write_pgm( FILE* f, int width, int height, unsigned char* data );
unsigned char* read_pgm( char* filename, int* width, int* height );

#define MAX_LABELS   80000

#define I(x,y)   (image[(y)*(width)+(x)])
#define N(x,y)   (nimage[(y)*(width)+(x)])

void uf_union( int x, int y, unsigned int parent[] )
{
     while ( parent[x] )
	  x = parent[x];
     while ( parent[y] )
	  y = parent[y];
     if ( x != y ) {
	  if (y<x) parent[x] = y;
	  else parent[y] = x;
     }	  
}

int next_label = 1;

int uf_find( int x, unsigned int parent[], unsigned int label[] )
{
     while ( parent[x] )
	  x = parent[x];
     if ( label[x] == 0 )
	  label[x] = next_label++;
     return label[x];
}

main( int argc, char **argv )
{
     int width, height;
     unsigned char* image;
     unsigned int* nimage;
     unsigned int parent[MAX_LABELS];
     unsigned int labels[MAX_LABELS];
     unsigned char* oimage;
     int next_region = 1;
     int i, j;
     int x, y, k;
     int output_type = 0;

#ifndef UNIX
	 _setmode( _fileno( stdin ), _O_BINARY );
	 _setmode( _fileno( stdout ), _O_BINARY );
#endif

     if ( argc != 3 && argc != 2 )
     {
	  fprintf( stderr, "usage: %s [g|c] inputfile\n", argv[0] );
	  exit( 1 );
     }

     if ( argc == 3 )
     {
	  switch ( argv[1][0] )
	  {
	     case 'g':
	     case 'G':
	       output_type = 1;
	       break;

	     case 'c':
	     case 'C':
	       output_type = 2;
	       break;

	     default:
	       fprintf( stderr, "[g|c] argument must be g or c.\n" );
	       exit( 1 );
	  }
	  image = read_pgm( argv[2], &width, &height );
     }
     else
	  image = read_pgm( argv[1], &width, &height );
     
     nimage = (unsigned int*)malloc( width * height * sizeof( unsigned int ) );
     memset( nimage, 0, width * height * sizeof( unsigned int ) );

     memset( parent, 0, sizeof(parent) );
     memset( labels, 0, sizeof(labels) );
     
     for ( y = 0; y < height; ++y )
     {
	  for ( x = 0; x < width; ++x )
	  {
	       if ( I(x,y) == 0 ) continue;

	       k = 0;
	       if ( x > 0 && I(x-1,y) == I(x,y) )
		    k = N(x-1,y);
	       if ( y > 0 && I(x,y-1) == I(x,y) && N(x,y-1) < k )
		    k = N(x,y-1);
	       if ( k == 0 )
	       {
		    k = next_region;
		    next_region++;
	       }

	       if ( k >= MAX_LABELS )
	       {
		    fprintf( stderr, "maximum number of labels reached.  increase MAX_LABELS and recompile.\n" );
		    exit( 1 );
	       }
	       
	       N(x,y) = k;

	       if ( x > 0 && I(x-1,y) == I(x,y) && N(x-1,y) != k )
		    uf_union( k, N(x-1,y), parent );
	       if ( y > 0 && I(x,y-1) == I(x,y) && N(x,y-1) != k )
		    uf_union( k, N(x,y-1), parent );
	  }
     }

     for ( i = 0; i < width*height; ++i )
	  nimage[i] = uf_find( nimage[i], parent, labels )-1;

     --next_label;
     
     fprintf( stderr, "%d regions\n", next_label );

     if ( output_type == 0 )
     {
	  if ( next_label <= 255 )
	  {
	       fprintf( stderr, "outputting grayscale image\n" );
	       output_type = 1;
	  }
	  else
	  {
	       fprintf( stderr, "outputting color image\n" );
	       output_type = 2;
	  }
     }
     
     if ( output_type == 2 )
     {
	  oimage = (unsigned char*)malloc( width * height * 3 );
	  for ( i = 0; i < width*height; ++i )
	  {
	       oimage[i*3] = nimage[i] % 256;
	       oimage[i*3+1] = ( nimage[i] / 256 ) % 256;
	       oimage[i*3+2] = ( nimage[i] / 65536 );
	  }
	  write_ppm( stdout, width, height, oimage );
     }
     else
     {
	  for ( i = 0; i < width*height; ++i )
	       image[i] = nimage[i];
	  write_pgm( stdout, width, height, image );
     }
}

unsigned char* read_pgm( char* filename, int* width, int* height )
{
     char buffer[80];
     int phase = 0;
     int type, maxval;
     int j;
     unsigned char* data;
     FILE* f;

     if ( strcmp( filename, "-" ) == 0 )
     {
	  f = stdin;
     }
     else
     {
	  f = fopen( filename, "rb" );
	  if ( f == NULL )
	       return NULL;
     }
     
     while ( phase < 4 )
     {
	  fgets( buffer, 80, f );
	  if ( buffer[0] == '#' ) continue;
	  switch( phase )
	  {
	     case 0: j = sscanf( buffer, "P%d %d %d %d\n", &type, width, height, &maxval ); break;
	     case 1: j = sscanf( buffer, "%d %d %d\n", width, height, &maxval ); break;
	     case 2: j = sscanf( buffer, "%d %d\n", height, &maxval ); break;
	     case 3: j = sscanf( buffer, "%d\n", &maxval ); break;
	     case 4: j = 0; break;
	  }
	  phase += j;
     }

     if ( type != 5 )
     {
	  fprintf( stderr, "bad input image format.\n" );
	  return NULL;
     }

     data = (unsigned char*)malloc ( *width * *height );
     fread( data, *width * *height, 1, f );

     fclose( f );
     
     return data;
}

void write_ppm( FILE* f, int width, int height, unsigned char* data )
{
     fprintf( f, "P6\n%d %d\n%d\n", width, height, 255 );
     fwrite( data, width*height*3, 1, f );
     fclose( f );
}

void write_pgm( FILE* f, int width, int height, unsigned char* data )
{
     fprintf( f, "P5\n%d %d\n%d\n", width, height, 255 );
     fwrite( data, width*height, 1, f );
     fclose( f );
}