// A base class for ultra-fast audio recognition.

// [Part of the Ultrafa.st framework]

// Idea: We do NOT spend time on a full FFT/DCT in the early stages

// of the recognition.

// Instead we stay in the time domain, turn the sample data into pixels

// and then convert that verybig*1px image into an "integral" image.

// Then we use the integral image access functions to probe for

// various frequencies and wavelets we come up with during the

// recognition of whatever we are currently listening to (that's

// what the higher-level algorithms based on this class do).

// Note we want a full 16 bit range for each pixel's value to make

// this truly hi-fi, so we actually reserve a whole 8 bytes for each 

// cell in the (1D) table (could make that 6 but that's annoying to

// handle).

// Stefan Reich, Gaz.AI, Sep 3 2021

//

// [Insert very liberal license here]

sclass AudioRecognizer {

  IAudioSample mainSample;

  double defaultInputSampleRate() { ret 44100; }

  // It works like this: There is a general interface for accessing an "integrated" audio clip - IAudioSample.

  interface IAudioSample {

    int channels(); // 1 for mono, 2 for left+right, 3 for center+left+right... or whatever channel model you prefer

    double length();     // in samples according to sampleRate

    double sampleRate(); // in hertz

    // Query the integral.

    // Result is in the range -32768*(end-start) to 32767*(end-start)...

    // unless you applied too much gain (there is no clipping).

    // channel is between 0 and channels()-1 from here on out

    double sampleSum(int channel, double start, double end);

    // Here the range is -1 to 1 just to spice things up

    default double getPixel(int channel, double start, double end) {

      ret doubleRatio(sampleSum(channel, start, end), (end-start)*32768);

    }

    // RENDERING FUNCTIONS (visualize audio as BufferedImage)

    // render audio as black-and-white (grayscale) stripes

    // h = height per channel

    default BufferedImage stripes(int h default 50) {

      int w = iceil(length());

      int channels = channels();

      ret imageFromFunction(w, h*channels, (x, y) -> {

        int channel = y/h;

        double value = sampleSum(channel, x, x+1);

        // lose lower 8 bits and shift to 0 to 255

        int digital = ifloor(value/256)+128;

        ret rgbIntFullAlpha(digital, digital, digital);

      });

    }

    // render audio as graph

    // h = height per channel

    default BufferedImage graph(int h default 100) {

      int w = iceil(length());

      ret mergeBufferedImagesVertically(

        countIteratorToList(channels(), c ->

          simpleGraph(w, h, x -> sampleSum(c, x, x+1), -32768, 32767)));

    }

    // render audio as stripes + graph (best way to look at it)

    default BufferedImage render(int h default 100) {

      ret mergeBufferedImagesVertically(stripes(h/2), graph(h));

    }

    // END OF RENDERING FUNCTIONS

    // find maximum amplitude, going pixel-by-pixel

    // (remember: This clip may already have been temporally

    // scaled with speedUp(), so a "pixel" may represent the average

    // of multiple audio samples.)

    default double maxAmplitude() {

      int n = iceil(length()), channels = channels();

      double max = 0;

      for i to n:

        for c to channels: 

          max = max(max, abs(sampleSum(c, i, i+1)));

      ret min(32767, max);

    }

    // There are various non-destructive virtual transformations

    // which you can do on the audio clip (gain, speed-up and time-shift).

    // All transformations are affine in time and amplitude and thus

    // preserve the "integral image" property.

    default IAudioSample gain(double factor) {

      ret factor == 1 ? this : new Gain(factor, this);

    }

    // gain to maximum volume possible without clipping

    // (even though clipping isn't even a thing in integral audio wonderland,

    // so we just define "clipping" as exceeding the 32767 value we are used to from real audio.)

    default IAudioSample normalize() {

      ret gain(doubleRatio(32767, maxAmplitude()));

    }

    // resample with a factor

    public default IAudioSample speedUp(double factor) {

      ret factor == 1 ? this : new SpeedUp(factor, this);

    }

    // resample to a target frequency

    public default IAudioSample sampleAt(double freq) {

      ret speedUp(sampleRate()/freq);

    }

    public default IAudioSample timeShift aka shift(double shift) {

      ret shift == 0 ? this : new TimeShift(shift, this);

    }

    // For debug-printing. Valued from 0 to 1 this time because why not. First channel only

    default L<Double> firstPixels(int n default 20) {

      double[] pixels = new[n];

      for i to n:

        pixels[i] = sampleSum(0, i, i+1)/32768;

      ret wrapDoubleArrayAsList(pixels);

    }

  } // end of IAudioSample

  // The core integral 1D image.

  sclass AudioSample implements IAudioSample {

    int channels;

    double sampleRate;

    int length;

    // Here they are: the partial sums of the 16 bit audio samples.

    // Channels are stored interleaved

    long[] data;

    public double sampleRate() { ret sampleRate; }

    public int channels() { ret channels; }

    public double length() { ret length; }

    // result is in the range -32768*(end-start) to 32767*(end-start)

    public double sampleSum(int channel, double start, double end) {

      // We could do linear interpolation here if we weren't so basic.

      int a = ifloor(start), b = ifloor(end);

      ret getEntry(channel, b-1)-getEntry(channel, a-1);

    }

    // Get an entry of the sum table - allow for out-of-bounds

    // requests (those just default to silence).

    long getEntry(int channel, int i) {

      if (i < 0) ret 0;

      i = min(i, length-1);

      ret data[i*channels+channel];

    }

    // perform the integration of the raw audio data

    *(L<short[]> samples, int *channels, double *sampleRate) {

      length = lengthLevel2_shortArrays(samples);

      data = new long[length*channels];

      long[] sums = new[channels];

      int iSample = 0, iChunk = 0, iInArray = 0;

      short[] chunk = null;

      for i to length:

        for c to channels: {

          if (chunk == null || iInArray >= chunk.length) {

            chunk = samples.get(iChunk++);

            iInArray = 0;

          }

          data[iSample++] = (sums[c] += chunk[iInArray++]);

        }

    }

  }

  // implementation of gain modifier

  srecord noeq Gain(double factor, IAudioSample original) implements IAudioSample {

    public double sampleRate() { ret original.sampleRate(); }

    public int channels() { ret original.channels(); }

    public double length() { ret original.length(); }

    public double sampleSum(int channel, double start, double end) {

      ret original.sampleSum(channel, start, end)*factor;

    }

    // coalesce consecutive gains

    public IAudioSample gain(double factor) {

      ret original.gain(this.factor*factor);

    }

  }

  // Implementation of the time-shift modifier.

  // moves the input <shift> samples to the left (cuts off beginning).

  // Shift can be fractional - we're in integral image (audio) wonderland after all

  // where a traditional pixel has no meaning.

  srecord noeq TimeShift(double shift, IAudioSample original) implements IAudioSample {

    public double sampleRate() { ret original.sampleRate(); }

    public int channels() { ret original.channels(); }

    public double length() { ret original.length()-shift; }

    public double sampleSum(int channel, double start, double end) {

      ret original.sampleSum(channel, start+shift, end+shift);

    }

    // coalesce consecutive time-shifts

    public IAudioSample timeShift(double shift) {

      ret original.timeShift(this.shift+shift);

    }

  }

  // Implementation of the speed-up modifier which transforms every frequency f to f*factor.

  // This is for convenience, you could also just call sampleSum() directly with larger intervals.

  sclass SpeedUp implements IAudioSample {

    double factor, invFactor;

    IAudioSample original;

    *(double *factor, IAudioSample *original) {

      if (factor < 1) fail("Can't slow down. " + factor);

      invFactor = 1/factor;

    }

    public double sampleRate() { ret original.sampleRate()*invFactor; }

    public int channels() { ret original.channels(); }

    public double length() { ret original.length()*invFactor; }

    public double sampleSum(int channel, double start, double end) {

      ret original.sampleSum(channel, start*factor, end*factor)*invFactor;

    }

    // coalesce consecutive speed-ups

    public IAudioSample speedUp(double factor) {

      ret original.speedUp(this.factor*factor);

    }

  }

  // Constructors from various types of PCM data (including rendered-on-the-spot)

  *() {}

  *(short[] samples, int channels) {

    this(ll(samples), channels);

  }

  *(L<short[]> samples, int channels) {

    mainSample = new AudioSample(samples, channels, defaultInputSampleRate());

  }

  *(double seconds, VF1<double[]> soundSource, int channels) {

    this(soundSourceToShortArrays(seconds, soundSource, channels), channels);

  }

  // in-place modifiers for mainSample (convenience functions)

  void applyGain(double factor) { mainSample = mainSample.gain(factor); }

  void normalize                { mainSample = mainSample.normalize(); }

  void speedUp(double factor)   { mainSample = mainSample.speedUp(factor); }

  // Here come the actual analysis functions.

  // This looks at a number of periods of a given frequency starting at a certain time in the audio

  // and returns an intensity value.

  // No phase adjustment here, so you have to call this twice to get meaningful (complex) results.

  srecord noeq SumOfVibrations(IAudioSample sample, int channel, double start, double freq, int periods) {

    double period, end;

    double rawSum() {

      period = sample.sampleRate()/freq;

      double sum = 0, t = start;

      for p to periods: {

        // Subtract an expected trough from an expected neighboring peak and add to overall sum.

        // Nota bene: Trough and peak have the same area (=length), so this is basically a Haar-like feature!

        // By the use of which we automatically get around nasty complications like DC offsets in the input data.

        sum += sample.sampleSum(channel, t, t+period/2)

             - sample.sampleSum(channel, t+period/2, t+period);

        t += period;

      }

      end = t;

      ret sum;

    }

    // alternate calculation adjusted for duration

    double sumDividedByDuration() {

      ret rawSum()/(end-start);

    }

  }

  // divided by duration

  Complex complexSumOfVibrations(IAudioSample sample, int channel, double start, double freq, int periods) {

    double duration = sample.sampleRate()/freq*periods;

    ret div(complexSumOfVibrations_raw(sample, channel, start, freq, periods), duration);

  }

  // Not divided by duration - this seems like the best frequency detector at this point.

  // As in a proper FFT/DCT, we return a complex value to represent phase.

  // Call abs() to get the desired intensity value.

  Complex complexSumOfVibrations_raw(IAudioSample sample, int channel, double start, double freq, int periods) {

    SumOfVibrations sum = new(sample, channel, start, freq, periods);

    double re = sum.rawSum();

    sum.start += sum.period/4; // 90° phase shift to catch the other half of the circle

    double im = sum.rawSum();

    ret Complex(re, im);

  }

}