// notes.c - notes and MIDI, as notes.js and basicpitch.js (same arithmetic, same bytes).

#include "bp.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>

#define FPS (BP_SAMPLE_RATE / BP_HOP)
#define LOWEST_MIDI 21
#define ENERGY_TOLERANCE 11

static double js_round(double x) { return floor(x + 0.5); }

static double frame_time(int frame) {
    const int window_frames = 2 * FPS;
    const double offset = (double)BP_HOP / BP_SAMPLE_RATE * (window_frames - (2.0 * BP_SAMPLE_RATE - BP_HOP) / BP_HOP) + 0.0018;
    return (double)frame * BP_HOP / BP_SAMPLE_RATE - offset * floor((double)frame / window_frames);
}

typedef struct { int start, end, key; double amplitude; } span;
typedef struct { span *at; int n, cap; } spans;
typedef struct { float level; int index; } peak;

static int push(spans *s, int start, int end, int key, const float *frames) {
    if (s->n == s->cap) {
        span *grown = realloc(s->at, (s->cap = s->cap ? 2 * s->cap : 256) * sizeof(span));
        if (!grown) return -1;
        s->at = grown;
    }
    double sum = 0;
    for (int r = start; r < end; r++) sum += frames[r * BP_NOTES + key];
    s->at[s->n++] = (span){ start, end, key, sum / (end - start) };
    return 0;
}

// loudest first, then the earlier frame, then the higher key
static int by_level(const void *pa, const void *pb) {
    const peak *a = pa, *b = pb;
    if (a->level != b->level) return a->level > b->level ? -1 : 1;
    const int ra = a->index / BP_NOTES, rb = b->index / BP_NOTES;
    if (ra != rb) return ra - rb;
    return b->index - a->index;
}

static int key_of(double hz) { return (int)trunc(12 * (log2(hz) - log2(440)) + 69 - LOWEST_MIDI); }

int bp_find_notes(const float *frames_in, const float *onsets_in, int rows, const bp_settings *s, bp_note **out) {
    const double onset_threshold = s->onset_threshold, frame_threshold = s->frame_threshold;
    const int n = rows * BP_NOTES, min_frames = (int)js_round(s->min_note_ms / 1000 * FPS);
    const size_t size = n > 0 ? n : 1;   // malloc(0) may be NULL
    float *frames = malloc(size * sizeof(float)), *onsets = malloc(size * sizeof(float)), *rise = calloc(size, sizeof(float));
    float *onset = malloc(size * sizeof(float)), *left = malloc(size * sizeof(float));
    int *starts = malloc(size * sizeof(int)), nstarts = 0, count = -3;
    peak *peaks = malloc(size * sizeof(peak));
    spans found = { 0 };
    *out = NULL;
    if (!frames || !onsets || !rise || !onset || !left || !starts || !peaks) goto done;
    memcpy(frames, frames_in, n * sizeof(float));
    memcpy(onsets, onsets_in, n * sizeof(float));
#define AT(r, k) ((r) * BP_NOTES + (k))

    if (s->max_frequency > 0) {
        const int from = key_of(s->max_frequency);
        for (int r = 0; r < rows; r++)
            for (int k = from < 0 ? 0 : from; k < BP_NOTES; k++) frames[AT(r, k)] = onsets[AT(r, k)] = 0;
    }
    if (s->min_frequency > 0) {
        const int to = key_of(s->min_frequency);
        for (int r = 0; r < rows; r++)
            for (int k = 0; k < to && k < BP_NOTES; k++) frames[AT(r, k)] = onsets[AT(r, k)] = 0;
    }

    memcpy(onset, onsets, n * sizeof(float));
    if (s->boost_onsets) {
        double onset_max = 0, rise_max = 0;
        for (int i = 0; i < n; i++) onset_max = fmax(onset_max, onsets[i]);
        for (int r = 2; r < rows; r++)
            for (int k = 0; k < BP_NOTES; k++) {
                const double v = frames[AT(r, k)], d = fmax(fmin(v - frames[AT(r - 1, k)], v - frames[AT(r - 2, k)]), 0);
                rise[AT(r, k)] = (float)d;
                rise_max = fmax(rise_max, d);
            }
        for (int i = 0; i < n; i++) onset[i] = (float)fmax(onsets[i], rise_max ? onset_max * rise[i] / rise_max : 0);
    }

    for (int r = rows - 1; r >= 0; r--)
        for (int k = BP_NOTES - 1; k >= 0; k--) {
            const float v = onset[AT(r, k)];
            if (v > onset_threshold && (r == 0 || v > onset[AT(r - 1, k)]) && (r == rows - 1 || v > onset[AT(r + 1, k)])) starts[nstarts++] = AT(r, k);
        }

    memcpy(left, frames, n * sizeof(float));   // energy no note has taken yet
#define TAKE(r, k) (left[AT(r, k)] = 0, (k) < BP_NOTES - 1 ? left[AT(r, (k) + 1)] = 0 : 0, (k) > 0 ? left[AT(r, (k) - 1)] = 0 : 0)

    for (int s = 0; s < nstarts; s++) {
        const int start = starts[s] / BP_NOTES, k = starts[s] % BP_NOTES;
        if (start >= rows - 1) continue;
        int r = start + 1, quiet = 0;
        while (r < rows - 1 && quiet < ENERGY_TOLERANCE) { quiet = left[AT(r, k)] < frame_threshold ? quiet + 1 : 0; r++; }
        const int end = r - quiet;
        if (end - start <= min_frames) continue;
        for (int i = start; i < end; i++) TAKE(i, k);
        if (push(&found, start, end, k, frames)) goto done;
    }

    // melodia trick
    int npeaks = 0;
    for (int i = 0; i < n; i++)
        if (left[i] > frame_threshold) peaks[npeaks++] = (peak){ left[i], i };
    qsort(peaks, npeaks, sizeof(peak), by_level);
    for (int p = 0; p < npeaks; p++) {
        if (left[peaks[p].index] != peaks[p].level) continue;   // already taken
        const int mid = peaks[p].index / BP_NOTES, k = peaks[p].index % BP_NOTES;
        left[peaks[p].index] = 0;
        int r = mid + 1, quiet = 0;
        while (r < rows - 1 && quiet < ENERGY_TOLERANCE) { quiet = left[AT(r, k)] < frame_threshold ? quiet + 1 : 0; TAKE(r, k); r++; }
        const int end = r - 1 - quiet;
        r = mid - 1; quiet = 0;
        while (r > 0 && quiet < ENERGY_TOLERANCE) { quiet = left[AT(r, k)] < frame_threshold ? quiet + 1 : 0; TAKE(r, k); r--; }
        const int start = r + 1 + quiet;
        if (end - start > min_frames && push(&found, start, end, k, frames)) goto done;
    }

    if (!(*out = malloc((found.n ? found.n : 1) * sizeof(bp_note)))) goto done;
    for (int i = 0; i < found.n; i++) {
        const span *s = &found.at[i];
        (*out)[i] = (bp_note){ s->key + LOWEST_MIDI, frame_time(s->start), frame_time(s->end) - frame_time(s->start), s->amplitude };
    }
    count = found.n;
done:
    free(frames); free(onsets); free(rise); free(onset); free(left); free(starts); free(peaks); free(found.at);
    return count;
}

typedef struct { int tick, on, pitch, velocity, order; } event;

// by tick, offs before ons, then as made (JS's sort is stable)
static int by_tick(const void *pa, const void *pb) {
    const event *a = pa, *b = pb;
    return a->tick != b->tick ? a->tick - b->tick : a->on != b->on ? a->on - b->on : a->order - b->order;
}

int bp_midi(const bp_note *notes, int count, double tempo, unsigned char **bytes) {
    const double ticks_per_second = 8 * tempo;   // 480 a beat
    const unsigned us = (unsigned)floor(60e6 / tempo);
    event *events = malloc((count ? 2 * count : 1) * sizeof(event));
    unsigned char *out = malloc(22 + 7 + (size_t)count * 2 * 8 + 4);   // events take at most 8 bytes
    *bytes = NULL;
    if (!events || !out) { free(events); free(out); return -3; }
    for (int i = 0; i < count; i++) {
        const int on = (int)js_round(notes[i].start * ticks_per_second);
        const int off = (int)fmax(on + 1, js_round((notes[i].start + notes[i].duration) * ticks_per_second));
        const int velocity = (int)fmin(127, fmax(1, js_round(notes[i].amplitude * 127)));
        events[2 * i] = (event){ on, 1, notes[i].pitch, velocity, 2 * i };
        events[2 * i + 1] = (event){ off, 0, notes[i].pitch, 0, 2 * i + 1 };
    }
    qsort(events, 2 * count, sizeof(event), by_tick);
    const unsigned char head[] = { 'M', 'T', 'h', 'd', 0, 0, 0, 6, 0, 0, 0, 1, 0x01, 0xe0, 'M', 'T', 'r', 'k', 0, 0, 0, 0,   // format 0, 480 ticks a beat
                                   0x00, 0xff, 0x51, 0x03, (us >> 16) & 255, (us >> 8) & 255, us & 255 };              // the tempo
    size_t n = sizeof head;
    memcpy(out, head, n);
    for (int i = 0, last = 0; i < 2 * count; i++) {
        unsigned delta = events[i].tick - last, shift = 28;
        last = events[i].tick;
        while (shift && !(delta >> shift)) shift -= 7;   // variable length, 7 bits a byte
        for (; shift; shift -= 7) out[n++] = 0x80 | ((delta >> shift) & 0x7f);
        out[n++] = delta & 0x7f;
        out[n++] = events[i].on ? 0x90 : 0x80;
        out[n++] = events[i].pitch;
        out[n++] = events[i].velocity;
    }
    out[n++] = 0x00; out[n++] = 0xff; out[n++] = 0x2f; out[n++] = 0x00;
    const size_t track = n - 22;
    out[18] = track >> 24; out[19] = track >> 16; out[20] = track >> 8; out[21] = track;
    free(events);
    *bytes = out;
    return (int)n;
}

int bp_audio_to_midi(bp_model *model, const float *audio, int samples, const bp_settings *settings, double tempo, unsigned char **bytes) {
    const bp_settings defaults = BP_SETTINGS;
    const int rows = bp_frame_count(samples);
    float *frames = malloc(((size_t)rows * BP_NOTES + 1) * sizeof(float)), *onsets = malloc(((size_t)rows * BP_NOTES + 1) * sizeof(float));
    bp_note *notes = NULL;
    int result = frames && onsets ? bp_transcribe(model, audio, samples, frames, onsets, NULL) : -3;
    *bytes = NULL;
    if (result >= 0) result = bp_find_notes(frames, onsets, rows, settings ? settings : &defaults, &notes);
    if (result >= 0) result = bp_midi(notes, result, tempo, bytes);
    free(frames); free(onsets); free(notes);
    return result;
}