/*
 * program that processes audio in PCM 16-bit little endian format
 * at a variety of bit rates but for now only mono. it either uses
 * the soundcard line-in/mic device or stdin. it processes the audio
 * and tries to isolate and emphasize the primary notes being played,
 * when then can then be mapped back to the twelve-tone chromatic scale.
 * basically, it currently doesnt isolate the primary notes very well
 * and only really is useful as a novel audio filter that basically makes
 * glitch/idm type of jams.
 *
 * christopher abad
 * aempirei@gmail.com
 * http://www.twentygoto10.com/
 * 2008
 *
 * this code requires FFTW3 and would be best compiled via:
 *
 * gcc -lfftw3 -lm pitch.c -o pitch -O2
 *
 * also, it has an optional text-based spectrum analyzer display.
 *
 */

#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <linux/soundcard.h>
#include <signal.h>
#include <assert.h>
#include <fftw3.h>
#include <math.h>
#include <time.h>

#define MIN(a,b) ((a) < (b) ? (a) : (b))

/*
 * a happy more atomic read function so i dont have to deal with loop
 * code overhead everytime i wanna do a read of a specific length
 */

ssize_t
read2(int fd, void *buf, size_t count)
{
   size_t nleft, nread;
   ssize_t n;

   nleft = count;
   nread = 0;

   do {
      n = read(fd, (char *)buf + nread, nleft);
      if (n == 0)
         return (nread);
      else if (n == -1)
         return (-1);

      nread += n;
      nleft -= n;
   } while (nleft);

   return (nread);

}

ssize_t
write2(int fd, void *buf, size_t count)
{
   size_t nleft, nwrite;
   ssize_t n;

   nleft = count;
   nwrite = 0;

   do {
      n = write(fd, (char *)buf + nwrite, nleft);
      if (n == 0)
         return (nwrite);
      else if (n == -1)
         return (-1);

      nwrite += n;
      nleft -= n;
   } while (nleft);

   return (nwrite);

}

int isdone = 0;

void
sighandler(int signo)
{
   if (signo == SIGINT)
      isdone = 1;
}

#define SHOW_SPECTRUM   0
#define NO_HARMONICS    1
#define TRIM_COEF       0.05
int show_spectrum = SHOW_SPECTRUM;

int bins = 64;
int rate = 8000;
int bits = 16;
int channels = 1;

int *prevbar;

fftw_plan fftplanf;
fftw_plan fftplanb;
fftw_complex *ffti, *ffto, *fftm;

const double bank_weight = 0.75;

double avgfreq = 0.0;

void
fft_abs(fftw_complex * fft, int N)
{
   int i;

   for (i = 0; i < N; i++) {
      fft[i][0] = hypot(fft[i][0], fft[i][1]);
      fft[i][1] = 0.0;
   }
}

void
fft_ramp(fftw_complex *fft, int N)
{
   int i;

   for(i = 0; i < N; i++)
      fft[i][0] *= (double)i * i;
}

double
getfreq(double binno)
{
   return binno * rate / bins;
}

double
decay(double input, double sum, double sum_weight)
{
   return (sum - input) * sum_weight + input;
}

double
fft_real_weight(fftw_complex * fft, int N)
{

   double sum = 0;
   int i;

   /*
    * i skip the first bin because its the relative gain and i dont want that
    */

   for (i = 1; i * 2 < N; i++)
      sum += fft[i][0];

   /*
    * we only want half the weight due to the symmetry of the FFT
    * im not sure of we want to use i < N / 2 in the for loop or what
    */

   return sum;
}

double
fft_find_real_midpoint(fftw_complex * fft, int N)
{

   double d = fft_real_weight(fft, N) / 2.0;
   double lo = 0.0, hi = 0.0, convex;

   int i;

   if (d == 0.0)
      return (0);

   for (i = 1; i + 1 < N; i++) {

      hi = lo + fft[i][0];

      if (d >= lo && d < hi) {
         convex = (d - lo) / (hi - lo);
         convex -= 0.5;         /* some reason the bins hit always in the middle */

         return i + convex;
      }

      lo = hi;
   }

   printf("tring to match %f\n", d);
   printf("made it to [ %f, %f ]\n", lo, hi);
   exit(EXIT_FAILURE);

   return (-1);
}

double
fft_max(fftw_complex * fft, int N)
{
   int i;

   double max = 0.0;

   for (i = 1; i < N; i++)
      if (fft[i][0] > max)
         max = fft[i][0];

   return max;
}

void
fft_normalize(fftw_complex * fft, int N)
{

   int i;

   double max = fft_max(fft, N);

   fft[0][0] = 0.0;
   fft[0][1] = 0.0;

   if (max != 0.0) {
      for (i = 0; i < N; i++)
         fft[i][0] /= max;
      for ( ; i < N; i++)
         fft[i][0] = 0.0;
   }
}

void
display_spectrum(fftw_complex * fft, int N)
{

   int i, j;

   for (i = 1; 2 * i < N; i++) {

      int left = fft[i][0] * 80;
      int skip;

      printf("%2i ", i);

      skip = MIN(left, prevbar[i]);

      if (skip > 0)
         printf("\33[%iC", skip);

      for (j = skip; j < left; j++)
         putchar('#');

      for (; j < prevbar[i]; j++)
         putchar(' ');

      prevbar[i] = left;

      putchar('\n');
   }
}

void
display_maxima(const double *m, int N) {

   int i;

   for(i = 0; i < N; i++)
      if(m[i] > TRIM_COEF)
         printf("%4i/%-4.0f ", i, getfreq(i));


   puts("\33[K");
}

double get_maxargmax(const double *m, int N) {

   int i;

   int argmax = -1;
   double max = 0.0;

   for(i = 0; i < N; i++) {
      if(m[i] > max) {
         argmax = i;
         max = m[i];
      }
   }

   return argmax;
}


double get_maxmax(const double *m, int N) {

   int i;

   int argmax = -1;
   double max = 0.0;

   for(i = 0; i < N; i++) {
      if(m[i] > max) {
         argmax = i;
         max = m[i];
      }
   }

   return max;
}

double *
get_maxima(fftw_complex * fft, int N)
{

   const int window = 2;
   double max;

   int i, j;
   int arg;
   
   double *m;

   m = malloc(sizeof(*m) * N);

   memset(m, 0, sizeof(*m) * N);

   for (i = 1; i < N; i++) {

      arg = i;
      max = fft[i][0];

      for (j = -window; j <= window; j++) {
         if (i + j >= 1 && i + j < N) {
            if (fft[i + j][0] > max) {
               arg = i + j;
               max = fft[i + j][0];
            }
         }

      }

      if (arg == i)
         m[i] = max;
   }


#if NO_HARMONICS
   for(i = 1; i < N; i++) {
      if(m[i] > 0.0) {
         for(j = 2; i * j < N; j++) {
            if(m[i * j]) {
               m[i] = 0.0;
               break;
            }
         }
      }
   }
#endif

   max = get_maxmax(m, N);

   if(max > 0.0)
      for(i = 1; i < N; i++)
         m[i] /= max;

   return m;
}

void
pitch_detect(short *buffer, int outfd)
{
   int i;
   double *maxima;
   // double freq;

   for (i = 0; i < bins; i++)
      ffti[i][0] = (double)buffer[i];

   fftw_execute(fftplanf);

   /*
    * convert to norm/magnitudes only
    */

   fft_abs(ffto, bins);

   //for(i = 0; i < bins; i++) printf("%3i  %f\n", i, ffto[i][0]);
   //exit(0);

   // fft_ramp(ffto, bins);

   fft_normalize(ffto, bins);

   fputs("\33[H", stdout);

   if (show_spectrum)
      display_spectrum(ffto, bins);

   putchar('\n');

   maxima = get_maxima(ffto, bins);


   /*

   freq = getfreq(fft_find_real_midpoint(ffto, bins));
   avgfreq = decay(freq, avgfreq, bank_weight);

   printf("freq: %5.0f avg-freq: %5.0f\33[K\n", freq, avgfreq);

   */

   if(outfd >= 0) {

      int m = 0;

      for (i = 0; i < bins; i++)
         ffti[i][0] = (double)buffer[i];

      fftw_execute(fftplanf);

//      fft_abs(ffto, bins);

      for(i = 0; i < bins; i++) {
         if(maxima[i] <= TRIM_COEF) {
            ffto[i][0] = 0.0;
            ffto[i][1] = 0.0;
            m++;
         } 
      }

      m = bins - m;

      printf("maxima = %2i / %i\33[K\n", m, bins);

      fftw_execute(fftplanb);

      for(i = 0; i < bins; i++) {
         short v;
         v = (int)rint(fftm[i][0] / bins);
         write(outfd, &v, sizeof(v));
      }
   }

   display_maxima(maxima, bins);
}

void
usage(const char *p)
{

   printf("\nusage: %s [-h] [-r rate] [-b bins] [-o outfile] [-s] [-v]\n\n", p);
   puts("  -h          display help");
   puts("  -r rate     set sample rate (default is 8000)");
   puts("  -o outfile  save vocoded audio out to a file (pcm_s16le)");
   puts("  -s          input audio from stdin (pcm_s16le)");
   puts("  -b bins     set FFT bin count (default 64)");
   puts("  -v          show verbose (spectrum)");
   putchar('\n');
}

int
main(int argc, char **argv)
{
   int fd = -1;                 /* sound device file descriptor */
   int outfd = -1;
   int bsize;
   void *buffer;
   int status;
   ssize_t n;
   unsigned long frames;
   char *outfile = NULL;
   int c;

   while ((c = getopt(argc, argv, "hso:r:b:v")) != -1) {

      switch (c) {

      case 'v':

         show_spectrum = 1;
         break;

      case 'b':

         bins = atoi(optarg);
         assert(bins > 0);
         break;

      case 's':

         fd = 0;
         break;

      case 'r':

         rate = atoi(optarg);
         assert(rate >= 4000);
         break;

      case 'o':

         outfile = optarg;
         break;

      case '?':

         /* fall-thru */

      case 'h':

         usage(*argv);
         exit(EXIT_FAILURE);

      }

   }

   /* open sound device */
   if (fd == -1)
      fd = open("/dev/dsp", O_RDWR);
   if (fd == -1) {
      perror("open of /dev/dsp failed");
      exit(EXIT_FAILURE);
   }

   /* set sampling parameters */

   if (fd > 0) {
      status = ioctl(fd, SOUND_PCM_WRITE_BITS, &bits);
      if (status == -1)
         perror("SOUND_PCM_WRITE_BITS ioctl failed");

      status = ioctl(fd, SOUND_PCM_WRITE_CHANNELS, &channels);
      if (status == -1)
         perror("SOUND_PCM_WRITE_CHANNELS ioctl failed");

      status = ioctl(fd, SOUND_PCM_WRITE_RATE, &rate);
      if (status == -1)
         perror("SOUND_PCM_WRITE_WRITE ioctl failed");
   }

   printf("%10s: %ihz\n", "rate", rate);
   printf("%10s: %i\n", "bits", bits);
   printf("%10s: %i\n", "channels", channels);
   printf("%10s: %i\n", "bins", bins);
   printf("%10s: %ihz\n", "nyquist", rate / 2);
   printf("%10s: %.1fhz\n", "resolution", (double)rate / bins);
   printf("%10s: %01.2fsec\n", "sample", (double)bins / rate);

   assert(bits == 16);
   assert(channels == 1);

   prevbar = malloc(bins * sizeof(*prevbar));

   assert(prevbar != NULL);

   memset(prevbar, 0, bins * sizeof(*prevbar));

   bsize = bins * bits * channels / 8;

   buffer = malloc(bsize);
   if (buffer == NULL) {
      perror("malloc()");
      return (-1);
   }

   ffti = fftw_malloc(sizeof(fftw_complex) * bins);
   ffto = fftw_malloc(sizeof(fftw_complex) * bins);
   fftm = fftw_malloc(sizeof(fftw_complex) * bins);

   fftplanf = fftw_plan_dft_1d(bins, ffti, ffto, FFTW_FORWARD, FFTW_MEASURE);
   fftplanb = fftw_plan_dft_1d(bins, ffto, fftm, FFTW_BACKWARD, FFTW_MEASURE);

   // setvbuf(stdout, NULL, _IONBF, 0);

   assert(signal(SIGINT, sighandler) != SIG_ERR);

   fputs("\33[2J", stdout);

   frames = 0;

   if(outfile != NULL) {
      outfd = open(outfile, O_WRONLY | O_TRUNC | O_CREAT, 0644);
      assert(outfd >= 0);
   }

   while (!isdone) {            /* loop until Control-C */

      n = read2(fd, buffer, bsize);     /* record some sound */
      if (n == 0)
         break;

      pitch_detect(buffer, outfd);

      frames += bins;

      printf("time elapsed: %.2f\33[K\n", (double)frames / rate);

   }

   puts("done...");

   close(outfd);
   close(fd);

   free(buffer);

   fftw_destroy_plan(fftplanf);
   fftw_destroy_plan(fftplanb);

   fftw_free(ffti);
   fftw_free(ffto);

   exit(EXIT_SUCCESS);
}