/*******************************************************************************/
/* aln2wig: converts psl and shrimp files to wig format                        */
/*                                                                             */
/* Autor: Ralph Krimmel                                                        */
/*******************************************************************************/

#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include <stdlib.h>
#include <unistd.h>
#include <getopt.h>


typedef int bool;
typedef char** stringArray;
#define TRUE   (1)
#define FALSE  (0)
#define FILEBUFFER 1048576
#define SHRIMP_CHROM_LENGTH 250000000
#define START_TOKEN_SIZE 128
#define SIZE_TOKEN_SIZE 128
#define TOKEN_SIZE 4096


/* chomp */
void chomp(const char *s)
{
  char *p;
  while (NULL != s && NULL != (p = strrchr(s, '\n')))
    {
      *p = '\0';
    }
}


/*returns the tokens of a string split by the separator in an array of char **/
int strtoken(char ***target, char *in, char sep) {
  int ret = 0;
  int i = 0;
  int si = 0;
  char **out = *target;

  while (in[i] != '\0') /* && i < TOKEN_SIZE)*/
    {
      if (in[i] != sep)
	{
	  out[ret][si++] = in[i];
	}
      else
	{
	  out[ret][si] = '\0';
	  si = 0;
	  ret++; }
      i++;
    }
  out[ret][si] = '\0';
  if (i > 0)
    {
      ret++;
    }
  return ret;
}

/*creates and writes the wigfile*/
void createWig(char *name, int *coveragefield, long int fieldlength, bool span)
{
  int i;
  int oldindex;
  int oldcoverage;
  if (span)
    {
      i=0;
      oldindex=0;
      oldcoverage=coveragefield[0];

      do
	{
	  if(oldcoverage==coveragefield[i])
	    {
	      i++;
	    }
	  else
	    {	if (coveragefield[i] != 0)
		{
		  printf("variableStep chrom=%s span=%i\n",name,i-oldindex);
		  printf("%i %i\n",i+1,coveragefield[i]);
		}
	      oldcoverage=coveragefield[i];
	      oldindex=i;
	      i++;
	    }

	} while(i<fieldlength);
    }
  else
    {
      printf("variableStep chrom=%s\n",name);
      for (i=0;i< fieldlength; i++)
	{
	  if(coveragefield[i]!=0)
	    {
	      printf("%i %i\n",i+1,coveragefield[i]);
	    }
	}
    }
}

/*checks if the input is still sorted*/
int stillSorted(char ***run,char *chkstr,int strcount)
{
  int i;
  char **target = *run;
  for (i=0; i < strcount; i++)
    {
      if (strcmp(target[i], chkstr) == 0)
	{
	  fprintf(stderr, "Found %s!. This should never happen on sorted input data.\nPlease sort the input.\nYou can do this by using the GNU sort tool: sort -k 14,14 <filename>", *target);
	  return FALSE;
	}
    }
  return TRUE;
}

int detectTokenCount(char *input, char tok)
{
  int i;
  int count=0;
  for (i=0;i < strlen(input); i++)
    {
      if (input[i] == tok)
	{
	  count ++;
	}
    }
  return count + 1;
}



int main (int argc, char **argv)
{
  int c;
  char *filename=NULL;
  char *trackname=NULL;
  /*Our filehandle*/
  FILE *handle;

  char usage[]="USAGE: aln2wig -f <filename>\nInput file can be a psl file or a shrimp file\nOutput goes to STDOUT\n\nOptions:\n\t-f\t<filename>\n\t-s\tUse span notation\n\t-t\tName of the track\n";

  char buffer [FILEBUFFER];

  /* line buffer of 1024 bytes */
  char line [ 1024 ];

  /* we need information about the length of our target sequence from the first line*/
  char firstline [TOKEN_SIZE];

  /*the list of completed chromosomes. we need this to check if the input is sorted or not*/
  char **completedChromosomes=NULL;

  int completedChromosomeCount=0;

  /* the name of the actual chromosome */
  char name[TOKEN_SIZE];

  /*we write the wigfile on every chromosome change and free the memory after it. To determine the change, we need to save to old name. */
  char oldname[TOKEN_SIZE];

  /*The tokens of the psl file we split by \t*/
  char **tokens;

  /* we save the chromosomeLength to create a coverageArray big enough to hold all positions. */
  long int chromosomeLength;

  /* the coverage array */
  int *coverage;

  bool span=FALSE;

  int i,j;


  /*****************/
  /* PSL variables */
  /*****************/

  /*number of tokens in the file. We need this to determine if input is in psl or shrimp format*/
  int tokenCounts;

  /*The tokens of the psl file we split by ",". Namely,this is the block size*/
  char **sizeTokens;

  /* The tokens of the psl file we split by ",". This is the absolute starting Positions of the block. */
  char **startTokens;

  /*the string we split into the  startTokens*/
  char *blockStarts;

  /*the string we split into the sizeTokens*/
  char *blockSizes;

  /* the number of blocks we have in each line */
  int blockCounts;


  /*********+++********/
  /* shrimp variables */
  /************+++*****/

  long int oldChromosomeLength;
  int contigStart=0;
  int contigEnd=0;

  opterr = 0;

  if ( argc == 1 )
    {
      printf ("%s",usage);
      exit(2);
    }

  while ((c = getopt (argc, argv, "t:sf:")) != -1)
    switch (c)
      {
      case 'f':
	filename = optarg;
	break;
      case 't':
	trackname = optarg;
	break;
      case 's':
	span=TRUE;
	break;
      case '?':
	if (optopt == 'c')
	  fprintf (stderr, "Option -%c requires an argument.\n", optopt);
	else if (isprint (optopt))
	  {
	    fprintf (stderr, "Unknown option `-%c'.\n", optopt);
	    fprintf (stderr, "%s",usage);
	  }
	else
	  {
	    fprintf (stderr,"Unknown option character `\\x%x'.\n",optopt);
	    fprintf (stderr, "%s",usage);
	    return 1;
	  }
      default:
	abort ();
      }

  handle = fopen ( filename, "r" );

  if ( handle != NULL )
    {
      /*print track name*/
      completedChromosomes = malloc(sizeof(char*));
      printf("track name=%s type=wiggle_0\n",trackname==NULL? filename : trackname);


      /*we want buffered input */
      setbuf(handle, buffer);

      /*read in the first line*/
      if (fgets (firstline, sizeof firstline, handle) == NULL){
	perror ( filename );
	printf ("Could not read first line.");
	exit(2);
      }
      chomp(firstline);

      tokenCounts = detectTokenCount(firstline,'\t');
      if (tokenCounts == 21)
	{
	  /*psl detected*/
	  tokens = malloc(21*sizeof(char*));
	  for (i = 0; i < 21; i++) {
	    tokens[i] = malloc(TOKEN_SIZE);
	    if (tokens[i] == NULL) exit(2);
	  }

	  strtoken(&tokens,firstline,'\t');
	  chromosomeLength=atoi(tokens[14]);

	  /* The 14th token is the name of the actual chromosome */
	  strncpy(oldname,tokens[13],TOKEN_SIZE-1);
	  oldname[TOKEN_SIZE-1] = '\0';

	  coverage =  (int *) calloc(chromosomeLength,sizeof(int));
	  if (coverage==NULL) exit (2);

	  /*set the file pointer back to start because we also need to parse the first line too */
	  fseek(handle, 0, SEEK_SET);

	  while (fgets ( line, sizeof line, handle ) != NULL )
	    {
	      chomp(line);
	      strtoken(&tokens,line,'\t');
	      strncpy(name, tokens[13], TOKEN_SIZE -1);
	      name[TOKEN_SIZE-1] = '\0';
	      /*if this happens the chromosome changes*/
	      if (strcmp(name,oldname)!=0)
		{
		  /*create the wigfile for the parsed chromosome*/
		  createWig(oldname,coverage,chromosomeLength,span);
		  completedChromosomeCount++;

		  /*remember the parsed chromosome to check if the input is sorted*/
		  completedChromosomes = realloc(completedChromosomes,completedChromosomeCount*sizeof(char*));
		  completedChromosomes[completedChromosomeCount-1] = malloc(TOKEN_SIZE*sizeof(char));
		  if (completedChromosomes[completedChromosomeCount-1] == NULL )
		    {
		      printf("Couldnt allocate Memory, exiting\n");
		      exit(1);
		    }
		  /*add the old chromosome to the list of completed chromosomes */
		  strcpy(completedChromosomes[completedChromosomeCount-1],oldname);

		  free(coverage);
		  /*The 15th token is the target sequence size*/
		  chromosomeLength=atoi(tokens[14]);
		  coverage = (int *) calloc(chromosomeLength,sizeof(int));
		  if (coverage==NULL) exit (2);
		  strcpy(oldname, name);

		  if (!stillSorted(&completedChromosomes,name,completedChromosomeCount))
		    {
		      exit (1);
		    }
		}

	      blockCounts = atoi(tokens[17])+1;
	      blockSizes =  tokens[18];
	      blockStarts = tokens[20];
	      startTokens = malloc(blockCounts*sizeof(char*));
	      if( startTokens == NULL) exit(2);
	      for (i = 0; i < blockCounts; i++)
		{
		  startTokens[i] = malloc(START_TOKEN_SIZE);
		  if (startTokens[i] == NULL) exit(2);
		}

	      sizeTokens = malloc(blockCounts*sizeof(char*));
	      if (sizeTokens == NULL) exit(2);
	      for (i = 0; i < blockCounts; i++)
		{
		  sizeTokens[i] = malloc(SIZE_TOKEN_SIZE);
		  if (sizeTokens[i] == NULL) exit(2);
		}

	      strtoken((char***) &startTokens,blockStarts,',');
	      strtoken((char***) &sizeTokens,blockSizes,',');


	      for (i = 0; i < blockCounts-1; i++)
		{
		  for (j = atoi(startTokens[i]); j < atoi(startTokens[i])+atoi(sizeTokens[i]); j++)
		    {
		      coverage[j]++;
		    }
		}

	      /*Cleaning up memory needed for parsing this line */

	      /*freeing startTokens array */
	      for (i=0; i< blockCounts; i++)
		{
		  free(startTokens[i]);
		  startTokens[i]=NULL;
		}
	      free(startTokens);
	      startTokens=NULL;

	      /*freeing sizeTokens array*/
	      for (i=0; i< blockCounts; i++)
		{
		  free(sizeTokens[i]);
		  sizeTokens[i]=NULL;
		}
	      free(sizeTokens);
	      sizeTokens=NULL;
	    }
	  createWig(oldname,coverage,chromosomeLength,span);
	  /*All done, clean up*/
	  for(i=0;i<21;i++)
	    {
	      free(tokens[i]);
	      tokens[i]=NULL;
	    }
	  free(tokens);
	  free(coverage);

	  for(i=0;i<completedChromosomeCount;i++)
	    {
	      free(completedChromosomes[i]);
	      completedChromosomes[i]=NULL;
	    }
	  free(completedChromosomes);
	}
      else
	{
	  /*Shrimp detected */
	  tokens = malloc(10*sizeof(char*));
	  oldChromosomeLength=SHRIMP_CHROM_LENGTH;
	  coverage =  (int *) calloc(oldChromosomeLength,sizeof(int));
	  for (i = 0; i < 11; i++) {
	    tokens[i] = malloc(TOKEN_SIZE);
	  }
	  /*Get the name of the first chromosome*/
	  strtoken(&tokens,firstline,'\t');
	  strncpy(oldname,tokens[1],TOKEN_SIZE -1);

	  /*set the file pointer back to start because we also need to parse the first line too */
	  fseek(handle, 0, SEEK_SET);

	  while (fgets ( line, sizeof line, handle ) != NULL )
	    {
	      strtoken(&tokens,line,'\t');
	      strncpy(name,tokens[1],TOKEN_SIZE-1);
	      chromosomeLength=atoi(tokens[4]);
	      if (chromosomeLength > oldChromosomeLength)
		{
		  oldChromosomeLength += (chromosomeLength-oldChromosomeLength+(SHRIMP_CHROM_LENGTH/5));
		  coverage = realloc(coverage,oldChromosomeLength*sizeof(int));
		}
	      if (strcmp(name,oldname)!=0)
		{
		  /*create the wigfile for the parsed chromosome*/
		  createWig(oldname,coverage,oldChromosomeLength,span);
		  completedChromosomeCount++;
		  /*remember the parsed chromosome to check if the input is sorted*/
		  completedChromosomes = realloc(completedChromosomes,completedChromosomeCount*sizeof(char*));
		  completedChromosomes[completedChromosomeCount-1] = malloc(TOKEN_SIZE*sizeof(char));
		  if (completedChromosomes[completedChromosomeCount-1] == NULL )
		    {
		      printf("Couldnt allocate Memory, exiting\n");
		      exit(1);
		    }
		  strcpy(completedChromosomes[completedChromosomeCount-1], oldname);

		  /*free(coverage); *
		    chromosomeLength=atoi(tokens[4]);
		    coverage = (int *) calloc(oldChromosomeLength,sizeof(int)); */
		  memset ( coverage, 0x00, oldChromosomeLength);
		  strncpy(oldname, name, TOKEN_SIZE-1);
		  oldname[TOKEN_SIZE-1] = '\0';
		  if (coverage==NULL) exit (2);
		  if (!stillSorted(&completedChromosomes,name,completedChromosomeCount))
		    {
		      exit (1);
		    }
		}
	      contigStart=atoi(tokens[3]);
	      contigEnd=atoi(tokens[4]);
	      for (i=contigStart;i<contigEnd+1;i++)
		{
		  coverage[i]++;
		}
	    }
	  createWig(oldname,coverage,oldChromosomeLength,span);
	  /*All done, clean up*/
	  for(i=0;i<10;i++)
	    {
	      free(tokens[i]);
	      tokens[i]=NULL;
	    }
	  free(tokens);
	  free(coverage);

	  for(i=0;i<completedChromosomeCount;i++)
	    {
	      free(completedChromosomes[i]);
	      completedChromosomes[i]=NULL;
	    }
	  free(completedChromosomes);
	}

      fclose ( handle );
    }
  else
    {
      perror ( filename ); /* why didn't the file open? */
      printf("%s",usage);
    }
  return 0;
}
