/* Copyright (C) 2025 Joel Yliluoma — Material for https://youtu.be/C4tHxouarGc - https://iki.fi/bisqwit/ */
/* MIT license follows:
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of
the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS
OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT
OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/
#include <algorithm>
#include <string>
#include <vector>
#include <format>
#include <print>
#include <cmath>
#include <map>
#include "fft.hh"
#include <SDL.h>

int main()
{
    SDL_Init(SDL_INIT_VIDEO | SDL_INIT_AUDIO);
    const int iscapture = SDL_TRUE;

    // Specify the type of audio we want to capture 
#ifdef USE_RATE
    const unsigned AudioRate = USE_RATE;
#else
    const unsigned AudioRate = 48000;
#endif
    SDL_AudioSpec wanted_spec{}, actual_spec{};
    wanted_spec.freq     = AudioRate;
    wanted_spec.format   = AUDIO_F32SYS;
    wanted_spec.channels = 1;
    wanted_spec.samples  = 256;

    // Open a dummy playback device. SDL apparently doesn't
    // want to capture, unless a playback device is opened too.
    SDL_AudioDeviceID dummy = SDL_OpenAudioDevice(NULL, !iscapture, &wanted_spec, &actual_spec, SDL_AUDIO_ALLOW_ANY_CHANGE);
    (void)dummy;

    // List the available capture devices:
    const char* name{};
    SDL_GetDefaultAudioInfo(const_cast<char**>(&name), &actual_spec, iscapture);
    for(int j=SDL_GetNumAudioDevices(iscapture), i=0; i<j; ++i)
    {
        const char* n = SDL_GetAudioDeviceName(i, iscapture);
        std::println("Recording device #{} = {}{}", i, n, std::string(n)==name ? " <-- using this one":"");
    }
    // Open the capture device (whichever one the system recommended).
    // Ideally, the callback would just point to SDL_AudioStreamPut
    // with userdata = stream, but stream doesn’t get initialized
    // until after the call to SDL_OpenAudioDevice. A minor inconvenience.
    SDL_AudioStream* stream{};
    wanted_spec.userdata = &stream;
    wanted_spec.callback = [](void* userdata, Uint8* buf, int len)
    {
        SDL_AudioStreamPut(*reinterpret_cast<SDL_AudioStream**>(userdata), buf, len);
    };
    SDL_AudioDeviceID devid = SDL_OpenAudioDevice(name, iscapture, &wanted_spec, &actual_spec, SDL_AUDIO_ALLOW_ANY_CHANGE);

    // Just in case the audio device gave different specs than we requested,
    // create a converter stream so that we get what we wanted. 
    stream = SDL_NewAudioStream(actual_spec.format, actual_spec.channels, actual_spec.freq,
                                wanted_spec.format, wanted_spec.channels, wanted_spec.freq);

    // Now open a graphics window.
    const unsigned VideoWidth = 1920, VideoHeight = 1080, DFTsize = AudioRate*1.0, MaxHz = 7000;
    auto* dft = FindDFT(DFTsize, 1,1, 1,1, 1);
    int p = SDL_WINDOWPOS_UNDEFINED;
    SDL_Window* window = SDL_CreateWindow("Audiodemo", p,p, VideoWidth, VideoHeight, SDL_WINDOW_RESIZABLE);
    SDL_Renderer* renderer = SDL_CreateRenderer(window, -1, 0);
    SDL_Texture* texture = SDL_CreateTexture(renderer, SDL_PIXELFORMAT_ARGB8888, SDL_TEXTUREACCESS_STREAMING, VideoWidth, VideoHeight);
    std::vector<Uint32> hist_vis(VideoWidth*VideoHeight), pixels(hist_vis);

    // Begin recording.
#ifndef BENCHMARK
    SDL_PauseAudioDevice(devid, 0);
#endif

    std::vector<complex> audiodata(DFTsize, complex{});
#ifdef BENCHMARK
    // Warmup DFT:
    { std::println("Warming up"); dft->xform_inplace(&audiodata[0],1); dft->xform_inplace(&audiodata[0],1); std::println("Done"); }
    std::FILE* timing = popen("/usr/bin/time cat", "w");
    unsigned framecount=0;
    std::srand(0);
    std::vector<complex> model(DFTsize);
    for(auto& f: model) f = std::polar(0.7f, std::rand()*1e-8f);
#endif
    for(std::map<int,bool> keys; !keys[SDLK_ESCAPE]; )
    {
#ifdef BENCHMARK
        if(++framecount >= BENCHMARK) { pclose(timing); return 0; }
        auto window = model;
#endif
#ifndef BENCHMARK
        SDL_LockAudioDevice(devid);
        if(int m=sizeof(float), n = SDL_AudioStreamAvailable(stream); n>=m)
        {
            char Buf[65536];
            n = SDL_AudioStreamGet(stream, Buf, std::min(n, int(sizeof(Buf)))/m*m);
            audiodata.insert(audiodata.end(), (float*)Buf, (float*)(Buf+n));
        }
        SDL_UnlockAudioDevice(devid);
        // Delete all but last 1 second of audio:
        audiodata.erase(audiodata.begin(), audiodata.end()-DFTsize);
        // Window the data, so that we can react to quick changes and not smush together 1 second of audio: 
        std::vector<complex> window(DFTsize);
        double sharpness = 180, Ce = std::exp(sharpness/-4), W = DFTsize/2., wsum=0; 
        auto w = [=](std::size_t j) { double p=j/W-0.5; return (std::exp(-sharpness*p*p)-Ce)/(1-Ce); };
        for(std::size_t j=0; j<DFTsize; ++j) wsum += w(j);
        for(std::size_t j=0; j<DFTsize; ++j)
        {
            std::size_t k = j+DFTsize*3/4;
            window[j] = (k<DFTsize ? audiodata[k] : 0) * float(w(j)/wsum);
        }
#endif
        // Run discrete Fourier transform on that data:
        dft->xform_inplace(&window[0],1);
#ifndef BENCHMARK
        // Transform the linear amplitudes into logarithmic volumes:
        std::vector<float> vol(DFTsize);
        for(std::size_t j=0; j<DFTsize; ++j) vol[j] = std::abs(window[j]);
        for(std::size_t j=0; j<DFTsize; ++j) vol[j] = 10*std::log10((vol[j]+1e-30)/4e-10);
        for(std::size_t j=0; j<DFTsize; ++j) vol[j] += 6*std::log((j+!j)/1000.0);
        // Scale and clamp them to 0—1 range.
        float mindB=40, maxdB=75, steepness=4;
        for(std::size_t j=0; j<DFTsize; ++j) vol[j] = (vol[j]-mindB)/(maxdB-mindB);
        for(std::size_t j=0; j<DFTsize; ++j) vol[j] = (std::tanh((vol[j]-.5f)*steepness)+1)*.5f; // Tip: plot this!
        // Now draw the scrolling canvas
        for(unsigned y=0; y<VideoHeight; ++y)
            std::copy_n(&hist_vis[y*VideoWidth+1], VideoWidth-1, &hist_vis[y*VideoWidth]);
        for(unsigned y=0; y<VideoHeight; ++y)
        {
            float p = (VideoHeight-1-y) * 1./VideoHeight; p = (std::pow(2., p)-1);
            unsigned index = MaxHz*p;
            float    f = std::pow(1-vol[index], 1.2f);
            int b = f*768, r = b-256, g = b-256;
            unsigned color = (std::clamp(r,0,255)<<16)+(std::clamp(g,0,255)<<8)+(std::clamp(b,0,255)<<0);
            hist_vis[y*VideoWidth + VideoWidth-1] = color;
        }
        // And draw the realtime amplitude meters
        pixels = hist_vis;
        for(unsigned x=0; x<VideoWidth; ++x)
        {
            float p = x * 1./VideoWidth; p = (std::pow(2., p)-1);
            unsigned index = MaxHz*p;
            unsigned height = vol[index] * VideoHeight;
            for(unsigned y=0; y<height; ++y)
                pixels[(VideoHeight-1-y)*VideoWidth + x] = 0xAA2277;
        }
        // Present the drawings on the screen.
        SDL_UpdateTexture(texture, nullptr, &pixels[0], 4*VideoWidth);
        SDL_RenderCopy(renderer, texture, nullptr, nullptr);
        SDL_RenderPresent(renderer);
        // Handle events
        for(SDL_Event ev; SDL_PollEvent(&ev); ) 
            switch(ev.type)
            {
                case SDL_QUIT: keys[SDLK_ESCAPE] = true; break;
                case SDL_KEYDOWN: keys[ev.key.keysym.sym] = true; break;
                case SDL_KEYUP:   keys[ev.key.keysym.sym] = false; break;
            }
        // Sleep for a token amount.
        SDL_Delay(1);
#endif
    }
}