// polyphase synthesis stretcher
// Copyright (C) 2022 Cockos Incorporated
// license: LGPL v3

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

#include "WDL/queue.h"

#define WDL_FFT_REALSIZE 8
#include "WDL/fft.c"

typedef double ReaSample;

class polyphase_stretch
{
  WDL_TypedBuf<ReaSample> m_inbuf;
  WDL_TypedQueue<ReaSample> m_inq, m_outq;

  WDL_TypedBuf<float> m_window; // analysis window followed by output window (fftsize each)
  WDL_TypedBuf<ReaSample> m_workbuf;

  WDL_TypedBuf<ReaSample> m_phase;

  int m_nch, m_fft_size, m_analysis_offset, m_output_offset;
  double m_stretch; // 0.1 for 10x slowdown
  double m_inq_skip;

  bool m_flush;

public:

  polyphase_stretch(int nch)
  {
    m_fft_size = m_analysis_offset = m_output_offset = 0;
    m_nch = nch;
    m_stretch = 1.0;
    m_inq_skip = 0.0;
    m_flush = false;

    Configure();
  }
  ~polyphase_stretch()
  {
  }

  void set_tempo(double tempo) { m_stretch = tempo; }

  void Reset()
  {
    m_inq.Clear();
    m_outq.Clear();
    m_inq_skip = -m_fft_size/2;
    m_workbuf.Resize(0,false);
    m_phase.Resize(0,false);
    m_flush = false;
  }

  // GetBuffer(), BufferDone(), GetSamples(), repeat. Also FlushSamples()/GetSamples() for end of stream
  ReaSample *GetBuffer(int size)
  {
    return m_inbuf.ResizeOK(size*m_nch);
  }
  void BufferDone(int input_filled)
  {
    int sz = input_filled * m_nch;
    if (WDL_NOT_NORMALLY(sz > m_inbuf.GetSize())) sz = m_inbuf.GetSize();
    if (sz<1) return;

    m_inq.Add(m_inbuf.Get(),sz);

    trim_input();
  }

  void FlushSamples()
  {
    m_flush = true;
  }

  int GetSamples(int requested_output, ReaSample *buffer)
  {
    ReaSample *rsbuf = buffer;
    int req = requested_output;

    int avail;
    for (;;)
    {
      avail = m_outq.GetSize()/m_nch - m_fft_size + m_output_offset;
      if (avail >= req) break;
      if (!synthesize()) break;
    }

    if (req > avail) return 0;

    if (rsbuf) memcpy(rsbuf,m_outq.Get(),m_nch * sizeof(*rsbuf) * req);
    m_outq.Advance(m_nch * req);
    m_outq.Compact();

    return req;
  }


private:
  void trim_input()
  {
    const int minsz = m_analysis_offset + m_fft_size;
    if (m_inq_skip >= minsz+1.0)
    {
      int sk = (int)m_inq_skip - minsz;
      const int a = m_inq.GetSize()/m_nch;
      if (sk > a) sk = a;
      if (sk > 0)
      {
        m_inq.Advance(sk*m_nch);
        m_inq.Compact();
        m_inq_skip -= sk;
      }
    }
  }

  bool synthesize()
  {
    int inq_skip = m_inq_skip >= 1.0 ? (int)m_inq_skip : 0;

    if (m_inq.GetSize() < (m_analysis_offset + m_fft_size + inq_skip) * m_nch)
    {
      if (!m_flush) return false;

      // flushing, use last source samples available
      inq_skip = m_inq.GetSize()/m_nch - (m_analysis_offset + m_fft_size);
      if (inq_skip < 0)
      {
        // this implies there were less than m_analysis_offset + m_fft_size ever received, 
        // fill with repeated previous samples or silence
        const int amt = -inq_skip*m_nch;
        const int src_avail = m_inq.GetSize();
        ReaSample *s = m_inq.Add(NULL, amt);
        if (WDL_NOT_NORMALLY(s==NULL)) return false;

        if (src_avail >= 16*m_nch)
          for (int x = 0; x < amt; x += src_avail)
            memcpy(s+x,m_inq.Get(), wdl_min(amt-x,src_avail)*sizeof(*s));
        else
          memset(s,0,amt*sizeof(*s));
        inq_skip = 0;
      }
      if (WDL_NOT_NORMALLY(m_inq.GetSize() < (m_analysis_offset + m_fft_size + inq_skip) * m_nch))
        return false;
    }

    if (WDL_NOT_NORMALLY(m_window.GetSize() != m_fft_size*2)) return false;

    const ReaSample * const rd = m_inq.Get() + (inq_skip + wdl_min((int)m_inq_skip,0)) * m_nch;
    const ReaSample * const minrd = m_inq.Get();
    // rd may be less than minrd, validate before reading

    const int stsize = m_nch * (m_fft_size/2);
    const bool is_init = m_workbuf.GetSize() != m_fft_size*2 || m_phase.GetSize() != stsize;

    ReaSample * const work = m_workbuf.ResizeOK(m_fft_size*2);
    if (WDL_NOT_NORMALLY(!work)) return false;

    ReaSample *phases = m_phase.ResizeOK(stsize);
    if (WDL_NOT_NORMALLY(!phases)) return false;

    ReaSample *wr = m_outq.Add(NULL,m_output_offset * m_nch);
    if (WDL_NOT_NORMALLY(!wr)) return false;
    memset(wr,0,m_output_offset*m_nch*sizeof(*wr));

    wr = m_outq.Get();
    int wrstart = m_fft_size - m_outq.GetSize()/m_nch;
    if (wrstart < 0)
    {
      // queue has more than a fft-length, write at offset
      wr += (-wrstart) * m_nch;
      wrstart = 0;
    }

    const double angadv = m_output_offset / (double) m_analysis_offset;

    const float * const window = m_window.Get();
    const float * const out_window = window + m_fft_size;

    for (int ch = 0; ch < m_nch; ch ++)
    {
      const ReaSample *rdp = rd + ch;
      const int rdo = m_nch * m_analysis_offset;
      for (int y = 0; y < m_fft_size; y ++)
      {
        const ReaSample w = window[y];
        work[y] = rdp >= minrd ? rdp[0] * w : 0.0;
        work[y + m_fft_size] = (rdp+rdo) >= minrd ? rdp[rdo] * w : 0.0;
        rdp+=m_nch;
      }
      WDL_real_fft(work, m_fft_size, 0); // this could probably be a single complex fft
      WDL_real_fft(work + m_fft_size, m_fft_size, 0);

      // work is now the two FFT results

      work[0] = 0.0; // zero DC offset, passthrough nyquist
      ReaSample *p = work + 2;
      const double PI = 3.1415926535897932384;
      for (int y = 1; y < m_fft_size/2; y++)
      {
        const double sre = p[0], sim = p[1];
        const double ere = p[m_fft_size + 0], eim = p[m_fft_size + 1];
        const double ang1 = atan2(sim, sre), ang2 = atan2(eim, ere);

        const double mag = sqrt(sre*sre+sim*sim); // *0.5 + sqrt(ere*ere+eim*eim) * 0.5;
        double ang = is_init ? ang1 : phases[y];
        p[0] = mag*cos(ang);
        p[1] = mag*sin(ang);
        p+=2;

        ang += (ang2-ang1) * angadv;
        if (fabs(ang) >= 2.0*PI) ang = fmod(ang, 2.0*PI);
        phases[y] = ang;
      }

      WDL_real_fft(work, m_fft_size,1);

      ReaSample *wp = wr + ch;
      for (int x = wrstart; x < m_fft_size; x ++)
      {
        wp[0] += work[x]*out_window[x];
        wp += m_nch;
      }
    }

    m_inq_skip += m_stretch * m_output_offset;
    trim_input();
    return true;
  }

  void Configure()
  {
    m_fft_size=32768; // fft size option
    int adiv = 2; // analysis subdivide
    int odiv = 4; // output subdivide
    m_analysis_offset = m_fft_size / adiv;
    m_output_offset = m_fft_size / odiv;
    m_inq_skip = -m_fft_size/2;

    m_workbuf.Resize(0,false);

    float *w = m_window.ResizeOK(m_fft_size * 2);
    if (w)
    {
      make_window(w, m_fft_size, 0 /* analysis window shape */);
      make_window(w + m_fft_size, m_fft_size, 0 /* output window shape */,
          1.0/m_fft_size * sqrt(m_output_offset / (double)m_fft_size));
    }
  }

  static void make_window(float *w, int sz, int mode, double outsc=1.0)
  {
    double v = 0.0, dv=3.1415926535897932384 * 2.0 / sz;
    int x = sz;
    if (mode == 4) dv = 2.0/sz;
    while (x--)
    {
      double tw;
      switch (mode)
      {
        case 4: tw = v < 1.0 ? v : 2.0-v; break;
        case 3: tw = 1.0; break; // rect
        case 2: tw = 0.42 - 0.50 * cos(v) + 0.08 * cos(2.0*v); break; // blackman
        case 1: tw = 0.53836 - cos(v)*0.46164; break; // hamming
        case 0: // blackman-harris
        default:
          tw = 0.35875 - 0.48829 * cos(v) + 0.14128 * cos(2*v) - 0.01168 * cos(3*v); break;
        break;
      }
      *w++ = (float)(tw * outsc);
      v += dv;
    }
  }

};