// cli.c - basicpitch [-m model_dir] [-onset 0.5] [-frame 0.3] [-min-ms 128] [-min-hz 0]
//                     [-max-hz 3000] [-no-boost] [-tempo 120] in.wav out.mid
// The WAV must be 22050 Hz (ffmpeg -i song.mp3 -ar 22050 song.wav).

#include "bp.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static uint32_t u32(const unsigned char *p) { return p[0] | p[1] << 8 | p[2] << 16 | (uint32_t)p[3] << 24; }
static uint16_t u16(const unsigned char *p) { return p[0] | p[1] << 8; }

// PCM or float WAV, as mono floats
static float *read_wav(const char *path, int *samples, int *rate) {
    FILE *f = fopen(path, "rb");
    if (!f) return NULL;
    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    unsigned char *wav = malloc(size);
    float *mono = NULL;
    if (!wav || fread(wav, 1, size, f) != (size_t)size || size < 12 || memcmp(wav, "RIFF", 4) || memcmp(wav + 8, "WAVE", 4)) goto done;
    int format = 0, channels = 0, bits = 0;
    for (long at = 12; at + 8 <= size;) {
        const uint32_t len = u32(wav + at + 4);
        const unsigned char *chunk = wav + at + 8;
        if (!memcmp(wav + at, "fmt ", 4) && len >= 16) {
            format = u16(chunk); channels = u16(chunk + 2); *rate = u32(chunk + 4); bits = u16(chunk + 14);
            if (format == 0xfffe && len >= 26) format = u16(chunk + 24);   // WAVE_FORMAT_EXTENSIBLE
        } else if (!memcmp(wav + at, "data", 4) && channels) {
            const int bytes = bits / 8, frames = (int)((len < (uint32_t)(size - at - 8) ? len : size - at - 8) / (bytes * channels));
            if (!(format == 1 && bytes >= 1 && bytes <= 4) && !(format == 3 && bits == 32)) goto done;
            if (!(mono = calloc(frames ? frames : 1, sizeof(float)))) goto done;
            for (int i = 0; i < frames; i++) {
                double sum = 0;
                for (int c = 0; c < channels; c++) {
                    const unsigned char *s = chunk + ((size_t)i * channels + c) * bytes;
                    float v;
                    if (format == 3) memcpy(&v, s, 4);
                    else if (bytes == 1) v = (s[0] - 128) / 128.0f;
                    else {
                        int32_t x = 0;
                        for (int b = 0; b < bytes; b++) x |= (int32_t)s[b] << (8 * (4 - bytes + b));
                        v = x / 2147483648.0f;
                    }
                    sum += v;
                }
                mono[i] = (float)(sum / channels);
            }
            *samples = frames;
            break;
        }
        at += 8 + len + (len & 1);
    }
done:
    free(wav);
    fclose(f);
    return mono;
}

int main(int argc, char **argv) {
    const char *dir = "basicpitch", *in = NULL, *out = NULL;
    bp_settings settings = BP_SETTINGS;
    double tempo = 120;
    for (int i = 1; i < argc; i++) {
        const int value = i + 1 < argc;
        if (value && !strcmp(argv[i], "-m")) dir = argv[++i];
        else if (value && !strcmp(argv[i], "-onset")) settings.onset_threshold = atof(argv[++i]);
        else if (value && !strcmp(argv[i], "-frame")) settings.frame_threshold = atof(argv[++i]);
        else if (value && !strcmp(argv[i], "-min-ms")) settings.min_note_ms = atof(argv[++i]);
        else if (value && !strcmp(argv[i], "-min-hz")) settings.min_frequency = atof(argv[++i]);
        else if (value && !strcmp(argv[i], "-max-hz")) settings.max_frequency = atof(argv[++i]);
        else if (value && !strcmp(argv[i], "-tempo")) tempo = atof(argv[++i]);
        else if (!strcmp(argv[i], "-no-boost")) settings.boost_onsets = 0;
        else if (!in) in = argv[i];
        else out = argv[i];
    }
    if (!in || !out || tempo <= 0) {
        fprintf(stderr, "usage: %s [-m model_dir] [-onset 0.5] [-frame 0.3] [-min-ms 128] [-min-hz 0] [-max-hz 3000] [-no-boost] [-tempo 120] in.wav out.mid\n", argv[0]);
        return 2;
    }

    char plan[4096], weights[4096];
    snprintf(plan, sizeof plan, "%s/plan.bin", dir);
    snprintf(weights, sizeof weights, "%s/weights.bin", dir);
    bp_model *model = bp_load_files(plan, weights);
    if (!model) { fprintf(stderr, "can't load %s and %s\n", plan, weights); return 1; }

    int samples = 0, rate = 0;
    float *audio = read_wav(in, &samples, &rate);
    if (!audio) { fprintf(stderr, "can't read %s (PCM or float WAV)\n", in); return 1; }
    if (rate != BP_SAMPLE_RATE) { fprintf(stderr, "%s is %d Hz, it must be %d Hz\n", in, rate, BP_SAMPLE_RATE); return 1; }

    const int rows = bp_frame_count(samples);
    float *frames = malloc((size_t)rows * BP_NOTES * sizeof(float)), *onsets = malloc((size_t)rows * BP_NOTES * sizeof(float));
    bp_note *notes;
    unsigned char *midi;
    int count, size;
    if (!frames || !onsets || bp_transcribe(model, audio, samples, frames, onsets, NULL) < 0 ||
        (count = bp_find_notes(frames, onsets, rows, &settings, &notes)) < 0 || (size = bp_midi(notes, count, tempo, &midi)) < 0) {
        fprintf(stderr, "out of memory\n");
        return 1;
    }
    FILE *f = fopen(out, "wb");
    if (!f || fwrite(midi, 1, size, f) != (size_t)size || fclose(f)) { fprintf(stderr, "can't write %s\n", out); return 1; }
    printf("%s: %d notes in %.1f s of audio\n", out, count, (double)samples / BP_SAMPLE_RATE);
    free(midi); free(notes); free(frames); free(onsets); free(audio);
    bp_unload(model);
    return 0;
}