/*
 *
 *    soniK digital audio editor
 *    Copyright (C) 2003-2006  Robert Walker <rob@tenfoot.org.uk>
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *    You should have received a copy of the GNU General Public License
 *    along with this program; if not, write to the Free Software
 *    Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 */
#include "spectraldisplay.h"

#include "spectraldisplaysettings.h"
#include "spectraldisplayconfig.h"

#include "wavewidget.h"
#include "data.h"
#include "sonik_util.h"
#include "sonik_sigproc.h"

#include <kcolorbutton.h>
#include <kapplication.h>
#include <kconfig.h>
#include <klocale.h>
#include <kgenericfactory.h>
#include <kdebug.h>

#include <qpainter.h>
#include <qlabel.h>
#include <qcombobox.h>

#include "colourmaps.h"

using Sonik::SpectralDisplay;
using Sonik::WaveWidget;

SpectralDisplay::SpectralDisplay(QObject* parent, const char* name,
                                 const QStringList& args)
  : Sonik::Display("spectral", i18n("Spectral"), parent, name, args),
    mConfigDlg(0)
{
  applyConfig();
}


SpectralDisplay::~SpectralDisplay()
{
}


Sonik::WaveWidget* SpectralDisplay::makeWidget(const Data& data,
                                               uint8_t channel,
                                               QWidget* parent,
                                               const char* name) const
{
  return new SpectralDisplay::Widget(*this, data, channel, parent, name);
}


QWidget* SpectralDisplay::makeConfigPage(QWidget* parent)
{
  mConfigDlg = new SpectralDisplayConfigDlg(parent, "spectral_display_config");

  mConfigDlg->selectColour->setColor(mSelectColour.color());
  mConfigDlg->windowFunction->setCurrentItem(static_cast<int>(mWindowFunction));
  mConfigDlg->windowSize->setCurrentText(QString::number(mWindowSize));
  // TODO: this screws case
  mConfigDlg->colourmap->setCurrentText(mColourmapName);

  return mConfigDlg;
}


void SpectralDisplay::applyConfigPage()
{
  SpectralDisplaySettings::setSelectColour(mConfigDlg->selectColour->color());
  SpectralDisplaySettings::setWindowFunction(
    windowToString(static_cast<WindowFunction>(mConfigDlg->windowFunction->currentItem()))
    );
  SpectralDisplaySettings::setWindowSize(mConfigDlg->windowSize->currentText().toUInt());
  SpectralDisplaySettings::setColourMap(mConfigDlg->colourmap->currentText());
  SpectralDisplaySettings::writeConfig();

  applyConfig();
}


void SpectralDisplay::applyConfig()
{
  mSelectColour   = SpectralDisplaySettings::selectColour();
  mWindowFunction = stringToWindow(SpectralDisplaySettings::windowFunction());
  mWindowSize     = SpectralDisplaySettings::windowSize();
  mColourmapName  = SpectralDisplaySettings::colourMap();
  buildColourmap(findColourmap(mColourmapName));
}


void SpectralDisplay::buildColourmap(const ColourValue* map)
{
  for (size_t i = 0; i < kColourmapSize; ++i)
  {
    mBrushColourmap[i] = QBrush(QColor(map[i].r, map[i].g, map[i].b));
    mPenColourmap[i] = QPen(QColor(map[i].r, map[i].g, map[i].b));
  }
}


SpectralDisplay::Widget::Widget(const SpectralDisplay& display,
                                const Data& data, uint8_t channel,
                                QWidget* parent, const char* name)
  : Sonik::WaveWidget(data, channel, parent, name),
    mDisplay(display),
    mLastWindowSize(0),
    mLastWindowFunction(Sonik::UNKNOWN),
    mFftParms(0)
{
}


SpectralDisplay::Widget::~Widget()
{
  delete mFftParms;

  for (Cache::iterator it = mCache.begin(); it != mCache.end(); ++it)
  {
    CacheData& d = it.data();
    for (CacheData::iterator itD = d.begin(); itD != d.end(); ++itD)
      delete[] (*itD);
  }
  mCache.clear();
}


void SpectralDisplay::Widget::render(QPainter& p, const QRect& r)
{
  // calculate x dimensions
  uint start = scrollPos() + (int)(r.left() / zoom());
  uint end = scrollPos() + (int)(r.right() / zoom());

  if (start > data().length() - 1)
    start = data().length() - 1;

  if (end > data().length() - 1)
    end = data().length() - 1;

  drawSpectrums(p, r, start, end);
  drawSelection(p, r, start, end);
}


void SpectralDisplay::Widget::drawSpectrums(QPainter& p, const QRect& r, uint start, uint end)
{
  // window parameters
  uint windowSize   = QMAX(mDisplay.windowSize(), (uint)(2.0/zoom()));
  uint windowSize_2 = windowSize >> 1;
  uint fftSize      = windowSize_2 + 1;

  uint wndCentre = start - (start % windowSize_2);

  int xStep = (int)(windowSize_2 * zoom());
  int xPos = Sonik::timeToScreenM(wndCentre, scrollPos(), zoom()) - xStep/2;

  bool scaleDown = ((uint)rect().height() <= fftSize);

  if (scaleDown)
    p.setBrush(Qt::NoBrush);
  else
    p.setPen(Qt::NoPen);

  double pointsPerPixel = (double)fftSize / rect().height();
  // TODO: brensham instead of floats
  double yStep = (float)rect().height() / fftSize;

  for (; xPos <= r.right() && wndCentre < end + windowSize_2; xPos += xStep)
  {
    int wndStart = wndCentre - windowSize_2;

    // get the fft
    float* spect = getSpectrum(windowSize, wndStart);
    assert(spect != 0);

    if (scaleDown)
      drawSliceScaleDown(p, xPos, r.top(), xStep, r.height(),
                         spect, fftSize, pointsPerPixel);
    else
      drawSliceScaleUp(p, xPos, r.top(), xStep, r.height(),
                       spect, fftSize, yStep);


    wndCentre += windowSize_2;
  }
}


void SpectralDisplay::Widget::drawSliceScaleDown(QPainter& p,
                                                 int x, int y, int w, int h,
                                                 float* spect, size_t spectSize, double scale)
{
  for (int yPos = y; h > 0; ++yPos, --h)
  {
    int i = (int)(yPos * scale);
    int c = (int)(255 * spect[spectSize - i]);
    // TODO: linear interpolate
    if (c < 0)
      c = 0;
    else if (c > 255)
      c = 255;

    p.setPen(mDisplay.penColourmap()[c]);
    p.drawLine(x, yPos, x+w, yPos);
  }
}


void SpectralDisplay::Widget::drawSliceScaleUp(QPainter& p,
                                               int x, int y, int w, int h,
                                               float* spect, size_t spectSize, double step)
{
  double yPos = 0.0;
  double yNext = step;
  int i = spectSize-1;
  for (; i >= 0 && (int)yPos <= y + h - 1; --i)
  {
    int y0 = (int)yPos;

    int yH;
    if (i == 0)
      yH = y + h - y0; // fudge to ensure last line is drawn
    else
      yH = ((int)yNext) - y0;

    yPos = yNext;
    yNext += step;

    if (y0+yH < y ||
        yH <= 0)
      continue;

    int c = (int)(255 * spect[i]);
    if (c < 0)
      c = 0;
    else if (c > 255)
      c = 255;

    p.setBrush(mDisplay.brushColourmap()[c]);
    p.drawRect(x, y0, w, yH);
  }
}


void SpectralDisplay::Widget::drawSelection(QPainter& p, const QRect& r, uint start, uint end)
{
  if (selLength() > 0)
  {
    uint selEnd = selStart() + selLength() - 1;

    if (selStart() <= (int)end && selEnd >= start)
    {
      int selLeft = Sonik::timeToScreenL(selStart(), scrollPos(), zoom());
      if (selLeft < r.left())
        selLeft = r.left();

      int selRight = Sonik::timeToScreenU(selEnd, scrollPos(), zoom());
      if (selRight > r.right())
        selRight = r.right();

      p.setRasterOp(Qt::NotROP);
      p.fillRect(selLeft, r.top(),
                 selRight - selLeft + 1, r.height(),
                 mDisplay.selectColour());
    }
  }
  else
  {
    int cursorPos = Sonik::timeToScreenM(selStart(), scrollPos(), zoom());

    if (cursorPos >= r.left() && cursorPos <= r.right())
    {
      p.setPen(Qt::SolidLine);
      p.setRasterOp(Qt::NotROP);
      p.drawLine(cursorPos, r.top(), cursorPos, r.bottom());
    }
  }
}


void SpectralDisplay::Widget::clearCache(CacheClearOp op, off_t start, size_t length)
{
  for (Cache::iterator cacheEntry = mCache.begin();
       cacheEntry != mCache.end();
       ++cacheEntry)
  {
    CacheData& c = cacheEntry.data();
    size_t ws_2 = cacheEntry.key().ws >> 1;
    if (op == WaveWidget::LengthChanged)
    {
      size_t blocks = ((length + ws_2) / ws_2) + 1;

      // clear blocks past new end
      for (size_t i = blocks; i < c.size(); ++i)
        delete[] c[i];

      c.resize(blocks, 0);
    }
    else if (op == WaveWidget::DataChanged)
    {
      size_t i = start/ws_2;
      size_t e = i + length / ws_2 + 2;
      for ( ; i < c.size() && i < e; ++i)
      {
        delete[] c[i];
        c[i] = 0;
      }
    }
  }
}


float* SpectralDisplay::Widget::getSpectrum(size_t ws, off_t pos)
{
  float* spect = 0;

  // map pos to positive val
  assert(pos >= -((off_t)ws >> 1));

  // Find any existing cache element
  CacheKey k(mDisplay.windowFunction(), ws);
  size_t posIndex = ((pos + (ws >> 1)) / (ws >> 1));
  CacheData& c = mCache[k];

  if (c.size() <= posIndex)
    c.resize(posIndex + 1, 0);

  spect = c[posIndex];

  if (spect == 0)
  {
    size_t fftSize = (ws >> 1) + 1;
    Sonik::SampleBuffer wnd(ws);
    Sonik::SampleBuffer buf(ws);
    Sonik::ComplexSampleBuffer fftData(fftSize);
    Sonik::SampleBuffer fftMag(fftSize);

    // get data & apply window
    if (ws != mLastWindowSize ||
        mDisplay.windowFunction() != mLastWindowFunction ||
        mWindow.size() == 0)
    {
      if (mWindow.capacity() < ws)
        mWindow.reset(ws);
      else
        mWindow.resize(ws);
      Sonik::makeWindow(mWindow, mDisplay.windowFunction());
    }
    data().data(channel(), pos, ws, buf);
    Sonik::mul(buf, mWindow);

    // compute FFT
    if (ws != mLastWindowSize || !mFftParms)
    {
      delete mFftParms;
      mFftParms = createFFTParms(ws);
    }
    Sonik::fft(buf.data(), fftData.data(), ws, mFftParms);
    fftData.abs(fftMag);

    // take log and scale to [0-1]
    // TODO: log optional, custom range
    Sonik::scale(fftMag, 1.0f/ws);
    Sonik::log10(fftMag);
    Sonik::scale(fftMag, 20.0f);
    Sonik::offset(fftMag, 200.0f);
    Sonik::scale(fftMag, 1.0f/200.0f);

    spect = fftMag.release();
    c[posIndex] = spect;

    mLastWindowSize = ws;
    mLastWindowFunction = mDisplay.windowFunction();
  }


  return spect;
}


bool SpectralDisplay::Widget::CacheKey::operator==(const CacheKey& r) const
{
  return (wf == r.wf && ws == r.ws);
}


bool SpectralDisplay::Widget::CacheKey::operator<(const CacheKey& r) const
{
  if (wf != r.wf)
    return wf < r.wf;
  if (ws != r.ws)
    return ws < r.ws;
  return false;
}


// void SpectralDisplay::drawWindows()
// {
//   QPen pens[] = { QPen(Qt::red), QPen(Qt::blue),
//                   QPen(Qt::green), QPen(Qt::magenta) };
//     if (widget->zoom() < 1.0)  // several samples per pixel
//     {
//       Sample* data = fftMag.data();
//       uint samplesLeft = windowSize;
//       int minPos, maxPos;
//       Sample oMin, oMax;
//       Sonik::minmax(data, 1, minPos, oMin, maxPos, oMax);

//       uint xPos = Sonik::timeToScreenL(wndStart,
//                                        widget->scrollPos(), widget->zoom());
//       uint xMax = Sonik::timeToScreenL(QMIN(wndEnd, end),
//                                        widget->scrollPos(), widget->zoom());
//       for (; samplesLeft > 0 && xPos <= xMax; xPos++)
//       {
//         Sample min, max;
//         size_t l = QMIN(samplesLeft, samplesPerPixel);
//         Sonik::minmax(data, l, minPos, min, maxPos, max);
//         data += l;
//         samplesLeft -= l;

//         int yTop = vCentre - QMAX(max, oMin) * maxDisp;
//         int yBottom = vCentre - QMIN(min, oMax) * maxDisp;

//         if (yTop != yBottom)
//           p.drawLine(xPos, yTop, xPos, yBottom);
//         else
//           p.drawPoint(xPos, yTop);

//         oMin = min;
//         oMax = max;
//       }
//     }
//     else               // several pixels per sample
//     {
//       Sample* data = fftMag.data();
//       xStep = (int)widget->zoom();
//       uint xMin = Sonik::timeToScreenL(wndStart,
//                                        widget->scrollPos(), widget->zoom());
//       xPos = xMin - (xMin % xStep); // round to nearest sample

//       yPos = vCentre - *data++ * maxDisp;

//       p.moveTo(xPos - xStep, yPos);
//       for (size_t i = 0; i < windowSize-2; i++)
//       {
//         yPos = vCentre - *data++ * maxDisp;
//         p.lineTo(xPos, yPos);
//         xPos += xStep;
//       }

//       // tail sample
//       yPos = vCentre - *data * maxDisp;
//       p.lineTo(xPos, yPos);
//     }
// }


K_EXPORT_COMPONENT_FACTORY(libsonik_displayspectral,
                           KGenericFactory<SpectralDisplay>(
                             "sonik-display-spectral")
                           );

