#include <stdio.h>
#include <wfdb/wfdb.h>

#define SIGTYPES "signal-types"

struct psig_t {
    char *name;
    int found;
    char *units;
    WFDB_Gain gain;
    int fmt;
    int spf;
    int adcres;
    int adczero;
    int baseline;
};

struct segd_t {
    char *recname;
    char basetime[32];
    long nsamp;
    WFDB_Time start;
    WFDB_Time end;
    WFDB_Time useful_start; /* Sample number of first nonzero sample */
    WFDB_Time useful_end; /* Sample number of last nonzero sample */
    WFDB_Time final_start; /* First sample number to be used */
    WFDB_Time final_end; /* Last sample number to be used */
    WFDB_Time shift;
};

int check_overlaps=0;
int shift_overlaps=0;

void memerr()
{
    fprintf(stderr, "insufficient memory -- exiting\n");
    exit(3);
}

char *copystr(char *oldstring)
{
    char *newstring;

    if (oldstring == NULL || *oldstring == '\0') {
	if ((newstring = calloc(1, 1)) == NULL)
	    memerr();
    }
    else {
	if ((newstring = calloc(1, strlen(oldstring) + 1)) == NULL)
	    memerr();
	strcpy(newstring, oldstring);
    }
    return (newstring);
}

/* Return number of bytes per sample for given format */

int fmtfact(int fmt)
{
  switch(fmt) {
  case 0:
    return 0;
  case 8:
  case 80:
    return 1;
  case 16:
  case 61:
  case 160:
    return 2;
  default:
    fprintf(stderr,"cannot use byte offsets with format %d\n", fmt);
    return 0;
  }
}

/* Find the start and end times for the given segment */

void get_lengths(int nsig, struct segd_t *seg)
{
    WFDB_Sample *samples;
    WFDB_Time t, current;
    int i, found=0;

    seg->start = 0;
    seg->end = seg->nsamp;

    if (!check_overlaps) {
      seg->useful_start = seg->start;
      seg->useful_end = seg->end;
      return;
    }

    samples = malloc(nsig * sizeof(WFDB_Sample));
    if (!samples)
	memerr();

    t = 0;

    /* Find the earliest nonzero sample in any signal */
    while (!found && getvec(samples) > 0) {
	for (i=0; i<nsig; i++)
	    if (samples[i] != 0 && samples[i] != -32768)
		found=1;
	t++;
    }

    if (t > 5)
	seg->useful_start = t - 2;
    else
	seg->useful_start = 0;

    /* Search backwards from the end to find the latest nonzero sample */
    current = seg->nsamp;

    do {
	current -= 1024;
	if (current < 0)
	    current = 0;
	isigsettime(current);
	seg->useful_end = -1;
	for (t=current; t<seg->nsamp && t<(current+1024); t++) {
	    if (0 > getvec(samples)) {
		exit(2);
	    }
	    for (i=0; i<nsig; i++)
		if (samples[i] != 0 && samples[i] != WFDB_INVALID_SAMPLE)
		    seg->useful_end = t + 1;
	}
    } while (seg->useful_end == -1 && current > 0);

    if (seg->useful_end == -1)
	seg->useful_end = 0;

    if (seg->useful_end > seg->nsamp - 5)
	seg->useful_end = seg->nsamp;

    free(samples);
}

char *usage = "usage: %s [-cs] [-o output-record] input-record ...\n"
" -c: check for empty space at start and end of segments, and remove\n"
"     it if necessary\n"
" -s: when adjacent segments overlap, shift the \"later\" segment\n"
"     forward in time (default is to drop the overlapping portion)\n"
" Note: input records must be in chronological order!\n";

int main(int argc, char **argv)
{
  static char buf[256], *p, *q, *sigtypes, *recordname, *filename;
  double *dv, *g, vf;
  FILE *sfile;
  int i, j, nsig, nosig, nseg, s, *smap, ss, t, tsig = 0;
  struct psig_t *psig = NULL;
  struct segd_t *seg;
  unsigned long seglen;
  WFDB_Sample *vin, *vout;
  WFDB_Siginfo *si, *so;
  WFDB_Time tt;
  char *outrecname = "out";

  if ((seg = calloc(sizeof(struct segd_t), argc)) == NULL)
    memerr();
  nseg = 0;

  for (i=1; i<argc; i++) {
    if (argv[i][0] == '-') {
      for (j=1; argv[i][j]; j++) {
	switch(argv[i][j]) {
	case 's':
	  shift_overlaps = 1;
	  break;
	case 'c':
	  check_overlaps = 1;
	  break;
	case 'o':
	  if (argv[i][j+1]) {
	    outrecname = &(argv[i][j+1]);
	  }
	  else if (++i < argc) {
	    outrecname = argv[i];
	  }
	  if (i < argc)
	    j = strlen(argv[i]) - 1;
	  else {
	    fprintf(stderr,"%s: -o requires an argument\n",argv[0]);
	    fprintf(stderr,usage,argv[0]);
	    return 1;
	  }
	  break;
	default:
	  fprintf(stderr,"%s: unrecognized option: %c\n",argv[0],argv[i][j]);
	  fprintf(stderr,usage,argv[0]);
	  return 1;
	}
      }
    }
    else {
      seg[nseg++].recname = argv[i];
    }
  }

  if (nseg < 1) {
    fprintf(stderr,usage,argv[0]);
    return 1;
  }

  /* Initialize psig (list of all known signal types). */
  if ((sigtypes = wfdbfile(SIGTYPES, NULL)) == NULL) {
    fprintf(stderr, "%s: can't find %s\n", argv[0], SIGTYPES);
    exit(1);
  }
  if ((sfile = fopen(sigtypes, "r")) == NULL) {
    fprintf(stderr, "%s: can't open %s\n", argv[0], SIGTYPES);
    exit(1);
  }
  while (fgets(buf, sizeof(buf), sfile)) {
    while (buf[strlen(buf)-1] == '\n' || buf[strlen(buf)-1] == '\r')
      buf[strlen(buf)-1] = 0;

    if (tsig++ == 0) {
      if ((psig = calloc(1, sizeof(struct psig_t))) == NULL ||
	  (g = calloc(1, sizeof(double))) == NULL)
	memerr();
    }
    else {
      if ((psig = realloc(psig, sizeof(struct psig_t) * tsig)) == NULL ||
	  (g = realloc(g, sizeof(double) * tsig)) == NULL)
	memerr();
    }
    psig[tsig - 1].name = copystr(buf);
    psig[tsig - 1].found = 0;
  }

  /* Pass 1: collect a list of all of the available signals. */
  for (i = 0; i < nseg; i++) {
    /* Strip off '.hea' if present in record name arguments. */
    if (p = strstr(seg[i].recname, ".hea"))
      *p = '\0';
    /* Determine the number of signals in the record. */
    if ((nsig = isigopen(seg[i].recname, NULL, 0)) <= 0)
      exit(2);
    /* Allocate storage for the signal information structures. */
    if ((si = malloc(nsig * sizeof(WFDB_Siginfo))) == NULL) {
      fprintf(stderr, "%s: insufficient memory\n", argv[0]);
      exit(3);
    }
    /* Get the signal information. */
    if (isigopen(seg[i].recname, si, check_overlaps?nsig:-nsig) != nsig)
      exit(2);
    strcpy(seg[i].basetime, mstimstr((WFDB_Time) 0));
    seg[i].nsamp = si[0].nsamp;
    get_lengths(nsig, &(seg[i]));
    seg[i].shift = 0;

    for (s = 0; s < nsig; s++) {	/* for each input signal */
      /* First, make sure each input has a unique name (if necessary,
	 append a '+' to the name to make it unique). */
      for (ss = 0; ss < s; ss++) {
	if (strcmp(si[s].desc, si[ss].desc) == 0) {
	  if ((p = calloc(1, strlen(si[s].desc) + 2)) == NULL)
	    memerr();
	  sprintf(p, "%s+", si[s].desc);
	  si[s].desc = p;
	}
      }
      /* Next, check each unique name against the list of all
	 signal names, and flag the matching name if it is found. */
      for (t = 0; t < tsig; t++)
	if (strcmp(si[s].desc, psig[t].name) == 0)
	  break;
      /* If the signal name was not in the list, add it to the end of
	 the list. */
      if (t >= tsig) {
	psig = realloc(psig, sizeof(struct psig_t) * ++tsig);
	g = realloc(g, sizeof(double) * tsig);
	if (psig == NULL || g == NULL)
	  memerr();
	psig[t].name = copystr(si[s].desc);
      }
      if (psig[t].found != 1) {
	psig[t].found = 1;
	psig[t].units = copystr(si[s].units);
	psig[t].gain = g[t] = si[s].gain;
	psig[t].fmt = si[s].fmt;
	psig[t].spf = si[s].spf;
	psig[t].adcres = si[s].adcres;
	psig[t].adczero = si[s].adczero;
	psig[t].baseline = si[s].baseline;
      }
      else {
	if (strcmp(psig[t].units, si[s].units))
	  fprintf(stderr, "warning: units of %s change\n",
		  si[s].desc);
	if (psig[t].gain != si[s].gain) {
	  if (si[s].gain > psig[t].gain) {
	    psig[t].gain = si[s].gain;
	    psig[t].baseline = si[s].baseline;
	  }
	  if (si[s].gain < g[t])
	    g[t] = si[s].gain;
	  fprintf(stderr, "warning: gain of %s changes\n",
		  si[s].desc);
	}
	if (psig[t].fmt != si[s].fmt)
	  fprintf(stderr, "warning: format of %s changes\n",
		  si[s].desc);
	if (psig[t].spf != si[s].spf)
	  fprintf(stderr, "warning: spf of %s changes\n", si[s].desc);
	if (psig[t].adcres != si[s].adcres)
	  fprintf(stderr, "warning: adcres of %s changes\n",
		  si[s].desc);
	if (psig[t].adczero != si[s].adczero)
	  fprintf(stderr, "warning: adczero of %s changes\n",
		  si[s].desc);
	if (psig[t].baseline != si[s].baseline)
	  fprintf(stderr, "warning: baseline of %s changes\n",
		  si[s].desc);
      }
    }
    wfdbquit();
  }

  /* Determine the final starting and ending points for the segments */
  
  if ((si = calloc(sizeof(WFDB_Siginfo), tsig)) == NULL)
    memerr();
  
  if ((nsig = isigopen(seg[0].recname, si, -tsig)) < 1)
    exit(2);
  for (i = 1; i < nseg; i++) {
    tt = strtim(seg[i].basetime);
    if (tt < 0) tt = -tt;
    seg[i].start += tt;
    seg[i].end += tt;
    seg[i].useful_start += tt;
    seg[i].useful_end += tt;
    
    if (seg[i].useful_start < seg[i-1].useful_end) {
	if (shift_overlaps) {
	    fprintf(stderr,
		    "%s: cannot resolve overlap between %s and %s, shifting forward by %s\n",
		    argv[0], seg[i-1].recname, seg[i].recname,
		    mstimstr(seg[i-1].useful_end - seg[i].useful_start));
	    seg[i].shift = seg[i-1].useful_end - seg[i].useful_start;
	    seg[i].start += seg[i].shift;
	    seg[i].end += seg[i].shift;
	    seg[i].useful_start += seg[i].shift;
	    seg[i].useful_end += seg[i].shift;
	    seg[i].final_start = seg[i].useful_start;
	    seg[i-1].final_end = seg[i-1].useful_end;
	}
	else {
	    fprintf(stderr,
		    "%s: cannot resolve overlap between %s and %s, discarding %s of earlier record\n",
		    argv[0], seg[i-1].recname, seg[i].recname,
		    mstimstr(seg[i-1].useful_end - seg[i].useful_start));

	    seg[i].final_start = seg[i-1].final_end = seg[i].useful_start;
	}
    }
    else if (seg[i].start < seg[i-1].useful_end)
      seg[i].final_start = seg[i-1].final_end = seg[i-1].useful_end;
    else if (seg[i].start <= seg[i-1].end)
      seg[i].final_start = seg[i-1].final_end = seg[i].start;
    else {
      // insert empty segment
      nseg++;
      seg = realloc(seg, nseg*sizeof(struct segd_t));
      if (!seg)
	memerr();
      for (j=(nseg-1); j>i; j--)
	memcpy(&(seg[j]), &(seg[j-1]), sizeof(struct segd_t));

      seg[i].recname = "~";
      seg[i].final_start = seg[i].useful_start = seg[i].start = seg[i-1].end;
      seg[i].final_end = seg[i].useful_end = seg[i].end = seg[i+1].start;
      seg[i+1].final_start = seg[i+1].start;
      seg[i-1].final_end = seg[i-1].end;
      i++;
    }
  }
  seg[0].final_start = 0;
  seg[nseg-1].final_end = seg[nseg-1].end;

  for (i=0; i<nseg; i++) {
    printf("%16s %s ", seg[i].recname, seg[i].basetime);
    printf("%10d (s+%4d) ", seg[i].final_start, seg[i].final_start - seg[i].start);
    printf("%10d (s+%9d) ", seg[i].final_end, seg[i].final_end - seg[i].start);
    printf("shift %d\n", seg[i].shift);
  }


  /* Adjust output signal info */

  for (s = t = 0; s < tsig; s++)
    if (psig[s].found)
      t++;
  if ((so = calloc(sizeof(WFDB_Siginfo), t)) == NULL)
    memerr();

  /* snprintf(buf, sizeof(buf), "%s_00000.dat", argv[1]); */
  
  for (s = t = 0; s < tsig; s++)
    if (psig[s].found) {
      so[t].fname = "~";/*copystr(buf);*/
      so[t].desc = copystr(psig[s].name);
      so[t].units = copystr(psig[s].units);
      so[t].gain = psig[s].gain;
      so[t].group = 0;
      so[t].fmt = 0;
      so[t].spf = psig[s].spf;
      so[t].bsize = 0;
      so[t].adcres = psig[s].adcres;
      if (psig[s].gain != g[s]) {	/* there are gain changes */
	double dg = psig[s].gain / g[s];
	      
	if (dg < 0.0)
	  dg = -dg;
	while (dg > 1.0) {
	  so[t].adcres++;
	  dg /= 2.0;
	}
	if (so[t].adcres > 16) {
	  /* The highest precision supported by the current version
	     of the WFDB library is 16 bits.  If we are here, the
	     gain of signal t has varied so much that we cannot
	     represent the smallest quantization steps in 16 bits. */
	  fprintf(stderr, "warning: there may be loss of precision"
		  " in the %s signal\n", so[t].desc);
	  while (--so[t].adcres > 16)
	    so[t].gain /= 2.0;
	}
      }
      so[t].adczero = psig[s].adczero;
      so[t].baseline = psig[s].baseline;
      so[t].nsamp = 0;
      so[t].cksum = 0;
      so[t].initval = so[t].adczero - (1 << (so[t].adcres-1));
      t++;
    }
  nosig = t;

  snprintf(buf, sizeof(buf), "%s_layout", outrecname);
  if (0>setheader(buf, so, nosig))
    exit(4);

  FILE *heafile, *subheafile;
  char newbasetime[256];
  char *info, **infostrs;
  int ninfostrs;
  int *groupcount;

  snprintf(buf, sizeof(buf), "%s.hea", outrecname);
  heafile = fopen(buf, "wb");
  if (!heafile) {
    perror(buf);
    exit(4);
  }

  strcpy(newbasetime, mstimstr(0));
  newbasetime[strlen(newbasetime)-1] = 0;
  
  fprintf(heafile,"%s/%d %d %lg %ld %s\r\n", outrecname, nseg+1, nosig,
	  (double) sampfreq(NULL), (long) seg[nseg-1].final_end, newbasetime+1);
  fprintf(heafile,"%s_layout 0\r\n",outrecname);

  wfdbquit();

  for (i=0; i<nseg; i++) {
    fprintf(heafile,"%s %ld\r\n", seg[i].recname,
	    (long) (seg[i].final_end - seg[i].final_start));

    if (strcmp(seg[i].recname, "~")) {
      if ((nsig = isigopen(seg[i].recname, si, -tsig)) < 1)
	exit(2);

      if ((seg[i].final_start + seg[i].shift - seg[i].start) > 0) {
	strcpy(newbasetime,
	       mstimstr(seg[i].start - seg[i].shift - seg[i].final_start));
      }
      else {
	strcpy(newbasetime, seg[i].basetime);
      }
      newbasetime[strlen(newbasetime)-1] = 0;

      infostrs = NULL;
      ninfostrs = 0;

      if (info = getinfo(seg[i].recname)) {
	infostrs = malloc(sizeof(char*));
	if (!infostrs)
	  memerr();
	ninfostrs = 1;
	infostrs[0] = copystr(info);
	while (info = getinfo(NULL)) {
	  infostrs = realloc(infostrs, (++ninfostrs)*sizeof(char*));
	  if (!infostrs)
	    memerr();
	  infostrs[ninfostrs-1] = copystr(info);
	}
      }

      groupcount = calloc(nsig, sizeof(int));
      if (!groupcount)
	memerr();

      for (j=0; j<nsig; j++) {
	si[j].nsamp = seg[i].final_end - seg[i].final_start;
	groupcount[si[j].group]++;
      }

      snprintf(buf, sizeof(buf), "%s.hea", seg[i].recname);
      subheafile = fopen(buf, "wb");
      if (!subheafile) {
	perror(buf);
	exit(4);
      }

      fprintf(subheafile,"%s %d %lg %ld %s\r\n", seg[i].recname, nsig,
	      (double) sampfreq(NULL),
	      (long) seg[i].final_end - seg[i].final_start,
	      newbasetime+1);

      for (j=0; j<nsig; j++) {
	fprintf(subheafile, "%s %d", si[j].fname, si[j].fmt);
	if (si[j].spf > 1)
	  fprintf(subheafile, "x%d", si[j].spf);
	if (seg[i].final_start > seg[i].start)
	  fprintf(subheafile, "+%d",
		  fmtfact(si[j].fmt)
		  * (seg[i].final_start - seg[i].start)
		  * groupcount[si[j].group]);
	fprintf(subheafile, " %lg", (double) si[j].gain);
	if (si[j].baseline != si[j].adczero)
	  fprintf(subheafile, "(%d)", si[j].baseline);
	fprintf(subheafile, "/%s %d %d %d %d %d %s\r\n", si[j].units,
		si[j].adcres, si[j].adczero, si[j].initval, si[j].cksum,
		si[j].bsize, si[j].desc);
      }

      if (seg[i].shift)
	fprintf(subheafile, "# Shifted by %d\r\n", seg[i].shift);
      if (seg[i].final_start > seg[i].start)
	fprintf(subheafile, "# %d samples removed from start\r\n",
		seg[i].final_start - seg[i].start);
      if (seg[i].end > seg[i].final_end)
	fprintf(subheafile, "# %d samples removed from end\r\n",
		seg[i].end - seg[i].final_end);

      seg[i].basetime[strlen(seg[i].basetime)-1] = 0;
      fprintf(subheafile, "# %s %d %lg %ld %s\r\n# -----\r\n",
	      seg[i].recname, nsig, (double) sampfreq(NULL),
	      (long) seg[i].end - seg[i].start, seg[i].basetime+1);

      for (j=0; j<ninfostrs; j++) {
	fprintf(subheafile, "#%s\r\n", infostrs[j]);
      }

      fclose(subheafile);
      wfdbquit();

      if (infostrs)
	free(infostrs);
      free(groupcount);
    }
  }

  fclose(heafile);
  wfdbquit();
  return 0;
}