#include <fftw3.h>
#include <string.h>
#include <math.h>
#include <stdio.h>
#include <vector>
#include <omp.h>

/* USE INSTRUCTIONS:
      Change input_filename below
      Compile
      Run: audio_transform_tandy > transformed_audio.raw

      Audio is expected to be 16-bit stereo. The sampling rate is defined below.

      Why such odd sampling rates like 46305? Frankly, I can't remember.
      Feel free to change it and see if something breaks.

  COPYRIGHT (C) 2012 Joel Yliluoma - http://iki.fi/bisqwit/
  License: Creative Commons Attribution 3.0 Unported
  If you retain my copyright message in this file, it'd be nice.
*/
static const char input_filename[] = "orig_ptal2.raw";

#if 0
static const unsigned rate         = 46305;
static const unsigned slice_ratio  = 1029; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 512;//4096;
#elif 0
static const unsigned rate         = 46305;
static const unsigned slice_ratio  = 245; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 512;//4096;
#elif 0
static const unsigned rate         = 46305;
static const unsigned slice_ratio  = 105; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 512;//4096;
#elif 0
static const unsigned rate         = 46305;
static const unsigned slice_ratio  = 63; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 768;//4096;
#elif 0
static const unsigned rate         = 46305;
static const unsigned slice_ratio  = 45; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 2048;//4096;
#else
static const unsigned rate         = 46300;
static const unsigned slice_ratio  = 30; // Number of times in second the rate is adjusted
static const unsigned fft_length   = 2048;//4096;
#endif
static const unsigned slice_length = rate / slice_ratio;
static const unsigned half_fft     = fft_length / 2;

static void DoFFT(const double *, double* Out, fftw_plan& plan, double Power[half_fft+1])
{
    memset(&Out[0], 0, sizeof(double) * fft_length);
    fftw_execute(plan);

    Power[0] = fabs(Out[0]);
    for(unsigned a=1; a<half_fft; ++a)
    {
        double p1 = Out[a], p2 = Out[fft_length-a];
        Power[a] = sqrt(p1*p1 + p2*p2);
    }
    if(!(fft_length%2))
        Power[fft_length/2] = sqrt(Out[fft_length/2] * Out[fft_length/2]);
}


#define MAX_OUTPUT 0x7fff
#define STEP       0x10000
#define FB_WNOISE  0x14002   /* (16bits) bit16 = bit0(out) ^ bit2 ^ bit15 */
#define FB_PNOISE  0x08000   /* JH 981127 - fixes Do Run Run */
#define NG_PRESET  0x0f35

struct SN76496
{
    int SampleRate;
    unsigned int UpdateStep;
    int VolTable[16];   /* volume table         */
    int Volume[4];      /* volume of voice 0-2 and noise */
    unsigned int RNG;   /* noise generator      */
    int NoiseFB;        /* noise feedback mask */
    int Period[4];
    int Count[4];
    int Output[4];

    unsigned V[4], P[4];

    SN76496()
    {
        /* increase max output basing on gain (0.2 dB per step) */
        double out = MAX_OUTPUT / 3;
        for(int gain=1; gain-- > 0; ) out *= 1.023292992;
        /* build volume table (2dB per step) */
        for (int i = 0;i < 15; ++i)
        {
            /* limit volume to avoid clipping */
            if (out > MAX_OUTPUT / 3) VolTable[i] = MAX_OUTPUT / 3;
            else VolTable[i] = (int)out;
            out /= 1.258925412; /* = 10 ^ (2/20) = 2dB */
        }
        VolTable[15] = 0;
        SampleRate = rate;
        unsigned clock = 3579545;
        UpdateStep = (unsigned int)(((double)STEP * SampleRate * 16) / clock);
        for(unsigned i=0; i<4; ++i)
        {
            Volume[i] = Output[i] = 0;
            Period[i] = Count[i] = UpdateStep;
            //Register[i*2+0] = 0; Register[i*2+1] = 0x0F;
            V[i] = 0;
            P[i] = 1;
        }
        RNG       = NG_PRESET;
        NoiseFB   = FB_WNOISE;
        Output[3] = RNG&1;
    }
};

struct Synth
{
    struct State
    {
        SN76496 chip;
    } state;

    void GenerateTo(short* buffer, unsigned length)
    {
        /* Uses:  Count, Period, RNG, NoiseFB, Output
        */
        SN76496* R = &state.chip;
        for(int i=0; i<4; ++i)
        {
            unsigned p = R->P[i];
            if(i == 3)
            {
                if(p == 3) p = R->P[2] * 2; else p = 1 << (5 + p);
            }

            R->Period[i] = R->UpdateStep * p;
            R->Volume[i] = R->VolTable[ R->V[i] ];

            if(R->Volume[i] == 0)
                if (R->Count[i] <= (int)length*STEP) R->Count[i] += length*STEP;
        }
        unsigned count=length;
        while (count)
        {
            int vol[4], left;
            unsigned int out;

            /* vol[] keeps track of how long each square wave stays */
            /* in the 1 position during the sample period. */
            vol[0] = vol[1] = vol[2] = vol[3] = 0;

            for (unsigned i = 0;i < 3;i++)
            {
                if (R->Output[i]) vol[i] += R->Count[i];
                R->Count[i] -= STEP;
                /* Period[i] is the half period of the square wave. Here, in each */
                /* loop I add Period[i] twice, so that at the end of the loop the */
                /* square wave is in the same status (0 or 1) it was at the start. */
                /* vol[i] is also incremented by Period[i], since the wave has been 1 */
                /* exactly half of the time, regardless of the initial position. */
                /* If we exit the loop in the middle, Output[i] has to be inverted */
                /* and vol[i] incremented only if the exit status of the square */
                /* wave is 1. */
                while (R->Count[i] <= 0)
                {
                    R->Count[i] += R->Period[i];
                    if (R->Count[i] > 0)
                    {
                        R->Output[i] ^= 1;
                        if (R->Output[i]) vol[i] += R->Period[i];
                        break;
                    }
                    R->Count[i] += R->Period[i];
                    vol[i] += R->Period[i];
                }
                if (R->Output[i]) vol[i] -= R->Count[i];
            }

            left = STEP;
            do
            {
                int nextevent;


                if (R->Count[3] < left) nextevent = R->Count[3];
                else nextevent = left;

                if (R->Output[3]) vol[3] += R->Count[3];
                R->Count[3] -= nextevent;
                if (R->Count[3] <= 0)
                {
                    if (R->RNG & 1) R->RNG ^= R->NoiseFB;
                    R->RNG >>= 1;
                    R->Output[3] = R->RNG & 1;
                    R->Count[3] += R->Period[3];
                    if (R->Output[3]) vol[3] += R->Period[3];
                }
                if (R->Output[3]) vol[3] -= R->Count[3];

                left -= nextevent;
            } while (left > 0);

            out = vol[0] * R->Volume[0] + vol[1] * R->Volume[1] +
                  vol[2] * R->Volume[2] + vol[3] * R->Volume[3];

            if (out > MAX_OUTPUT * STEP) out = MAX_OUTPUT * STEP;

            *(buffer++) = (short)(out / STEP);

            count--;
        }
    }
};


struct FFTcontext
{
    double In[fft_length];
    double Out[fft_length];
    fftw_plan plan;

    FFTcontext()
    {
        memset(&Out, 0, sizeof(Out));
        plan = fftw_plan_r2r_1d
            (fft_length,
             &In[0],
             &Out[0],
             FFTW_R2HC,
     		 FFTW_MEASURE | FFTW_PATIENT | FFTW_EXHAUSTIVE);
        memset(&In,  0, sizeof(In));
        memset(&Out, 0, sizeof(Out));
    }
};

static void TryMatchSpectrum(const double RefPower[half_fft+1])
{
    static Synth synth;
    static std::vector<FFTcontext> ctx ( omp_get_num_procs() );

    for(auto& c : ctx)
        for(unsigned a=0; a<fft_length - slice_length; ++a)
            c.In[a] = c.In[a + slice_length];

    Synth::State base_state = synth.state;
    Synth::State cur_state = synth.state;

    double best_error = 1e99;
    Synth::State best_state;

    for(bool improved=true; improved; )
    {
        improved = false;
        for(unsigned index=0; index<4; ++index)
        {
            #pragma omp parallel for schedule(static) \
                default(none) \
                shared(improved),firstprivate(index),shared(RefPower), \
                shared(stderr),shared(best_error),shared(best_state), \
                shared(cur_state),shared(ctx)
            for(unsigned vol0=0;    vol0<=15;       ++vol0)
            {
                FFTcontext& c = ctx[omp_get_thread_num() ];

                for(unsigned period0=1; period0<=0x3FF; ++period0)
                for(unsigned vol3=0; vol3<=15; ++vol3)
                for(unsigned mode3=0; mode3<1; ++mode3)
                {
                    if((index != 2) && (vol3||mode3)) break;
                    if(mode3 == 0 && vol3 != 0) continue;
                    if(vol3 == 15 && mode3 != 0) break;
                    if(index == 3 && (period0-1) >= 3) break;

                    Synth s = { cur_state };
                    SN76496* R = &s.state.chip;

                    /* Tweak synthesizers */
                    if(index == 2 && mode3 != 0)
                    {
                        R->V[3] = vol3;
                        R->P[3] = 3;
                    }
                    R->P[index] = index==3 ? period0-1 : period0;
                    R->V[index] = vol0;

                    /* Synthesize audio   */
                    short SynthBuf[slice_length];
                    s.GenerateTo(SynthBuf, slice_length);

                    /* Feed audio to FFTW */
                    for(unsigned a=0; a<slice_length; ++a)
                        c.In[fft_length - slice_length + a] =
                        0.5 + (int(SynthBuf[a])) / 65536.0;

                    double TestPower[half_fft+1];
                    DoFFT(c.In, c.Out, c.plan, TestPower);

                    /* Determine error */
                    double error = 0;
                    for(unsigned a=1; a<half_fft; ++a)
                    {
                        double significance = 1.0;// / log2(a+a);
                        /*if(a >= 512 && a <= 1024) significance *= 2.0;
                        if(a >= 2048) significance *= 0.5;
                        if(a >= 4096) significance *= 0.5;*/
                        error += fabs( (TestPower[a]) - (RefPower[a]) );
                    }

                    #pragma omp critical(test_error)
                  {
                    #pragma omp flush(best_error)
                    if(error < best_error)
                    {
                        /* Save this synthesization candidate */
                        best_error = error;
                        best_state = s.state;
                        improved = true;
                        fprintf(stderr, "Error %10g with %4u:%2u %4u:%2u %4u:%2u %4u:%2u\n", error,
                            R->P[0], R->V[0],
                            R->P[1], R->V[1],
                            R->P[2], R->V[2],
                            R->P[3], R->V[3]);
                        #pragma omp flush(improved,best_error,best_state)
                    }
                  }

                    if(vol0 == 15 && period0 == 1) break;
                    //if(index == 3 && period0 == 4+1) break;
                }
            }
            for(unsigned i=0; i<4; ++i)
            {
                cur_state.chip.P[i] = best_state.chip.P[i];
                cur_state.chip.V[i] = best_state.chip.V[i];
            }
        }
    }

    /* Render the audio and update the synthesizer */
    short SynthBuf[slice_length];
    synth.state = base_state;
    for(unsigned i=0; i<4; ++i)
    {
        synth.state.chip.P[i] = best_state.chip.P[i];
        synth.state.chip.V[i] = best_state.chip.V[i];
    }
    synth.GenerateTo(SynthBuf, slice_length);
    fwrite(&SynthBuf, 2, slice_length, stdout);
    fflush(stdout);
    for(auto& c: ctx)
        for(unsigned a=0; a<slice_length; ++a)
            c.In[fft_length - slice_length + a] =
            0.5 + (int(SynthBuf[a])) / 65536.0;

	{FILE *wfp = fopen("fftw.wisdom", "wb");
	if(wfp){fftw_export_wisdom_to_file(wfp);
	fclose(wfp);}}
}

int main()
{
	{FILE *wfp = fopen("fftw.wisdom", "rb");
	if(wfp){fftw_import_wisdom_from_file(wfp);
	fclose(wfp);}}

    static double In[fft_length]  = { 0 };
    static double Out[fft_length] = { 0 };

	fftw_plan plan = fftw_plan_r2r_1d
		(fft_length,
		 &In[0],
		 &Out[0],
		 FFTW_R2HC,
		 FFTW_MEASURE | FFTW_PATIENT | FFTW_EXHAUSTIVE);

	{FILE *wfp = fopen("fftw.wisdom", "wb");
	if(wfp){fftw_export_wisdom_to_file(wfp);
	fclose(wfp);}}

    FILE* fp = fopen(input_filename, "rb");
    setbuf(fp, NULL); // So that strace can identify read buffer size

    for(unsigned sliceno=0;;)
    {
        fprintf(stderr, "slice %u... rate=%u,slice_ratio=%u,slice_length=%u,fft_length=%u\n",
            sliceno++, rate, slice_ratio, slice_length, fft_length);
        short buffer[slice_length * 2];
        if(fread(&buffer, 4, slice_length, fp) <= 0) break;

        for(unsigned a=0; a<fft_length - slice_length; ++a)
            In[a] = In[a + slice_length];

        for(unsigned a=0; a<slice_length; ++a)
            In[fft_length - slice_length + a] =
            0.5 +
               (int(buffer[a*2+0])
              + int(buffer[a*2+1])) / 131072.0;

        static double Power[half_fft+1];
        DoFFT(In, Out, plan, Power);

        TryMatchSpectrum(Power);
    }

    fclose(fp);
}