/**************************************************************************
 * Copyright (C) 2008 Cocha                                               *
 * http://sourceforge.jp/projects/ecodecotool/                            *
 *                                                                        *
 *  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, 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 GNU Make; see the file COPYING.  If not, write to          *
 *  the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA. *
 *                                                                        *
 **************************************************************************/

#include "ConvertWav.h"


CConvertWav::CConvertWav(IUnknown *pUnk, HRESULT *phr)
   : CSimpleTransform(L"ConvertWav", pUnk, CLSID_ConvertWav, phr)
{  

   
   m_rtCurrent = 0;
   m_pConvertBits = NULL;
   m_pConvertChannels = NULL;
   m_pConvertSamples = NULL;
   m_smpld = NULL;

   m_nWait = 0;
   m_llInputSize = 0;

   m_bUseNormalize = false;
   m_dNormalizeScale = 1.0;
}

CConvertWav::~CConvertWav()  
{  

   
   if(m_smpld != NULL)
   {
      delete []m_smpld;
      m_smpld = NULL;
   }

   SAFE_DELETE(m_pConvertBits);
   SAFE_DELETE(m_pConvertChannels);
   SAFE_DELETE(m_pConvertSamples);
}

HRESULT CConvertWav::OnStart(void)
{  

   m_rtCurrent = 0;
   m_llInputSize = 0;

   
   if(m_inFormat.nSamplesPerSec == m_outFormat.nSamplesPerSec)
   {
      if( (m_inFormat.wFormatTag != m_outFormat.wFormatTag) || (m_inFormat.wBitsPerSample != m_outFormat.wBitsPerSample) )
      {
         if(m_pConvertBits == NULL)
            m_pConvertBits = new CConvertBits();

         m_pConvertBits->SetFormat(&m_inFormat, &m_outFormat);
      }

      if(m_inFormat.nChannels != m_outFormat.nChannels)
      {
         if(m_pConvertChannels == NULL)
            m_pConvertChannels = new CConvertChannels();

         m_pConvertChannels->SetFormat(&m_inFormat, &m_outFormat);
      }
   }
   else
   {
      if(m_pConvertSamples == NULL)
         m_pConvertSamples = new CConvertSamples();

      if(m_smpld == NULL)
         m_smpld = new SmplData[131072];

      
      for(int i=0;i<0x20000;i++)
      {
         m_smpld[i].l = 0.0;
         m_smpld[i].r = 0.0;
      }

      m_nSmplSize = 0;

      if(m_bUseNormalize == false)
         m_pConvertSamples->ConvertInit(&m_inFormat, &m_outFormat, 4, 1.0);
      else
         m_pConvertSamples->ConvertInit(&m_inFormat, &m_outFormat, 4, m_dNormalizeScale);

      
      DWORD feature_flag;
      __asm
      {
         mov eax, 1;
         cpuid;
         mov dword ptr [feature_flag], edx;
      }

      if(feature_flag & 0x04000000)
         m_bUseSSE2 = true;
   }

   return S_OK;
}

HRESULT CConvertWav::OnSeek(void)
{  

   m_rtCurrent = 0;

   if(m_pConvertSamples != NULL && m_smpld != NULL)
   {
      
      for(int i=0;i<0x20000;i++)
      {
         m_smpld[i].l = 0.0;
         m_smpld[i].r = 0.0;
      }

      m_nSmplSize = 0;

      if(m_bUseNormalize == false)
         m_pConvertSamples->ConvertInit(&m_inFormat, &m_outFormat, 4, 1.0);
      else
         m_pConvertSamples->ConvertInit(&m_inFormat, &m_outFormat, 4, m_dNormalizeScale);
   }

   return S_OK;
}

void CConvertWav::ConvertToDouble(BYTE *pbData, int nDataLength)
{  

   int i, j, nPos;

   int nInSmplSize = nDataLength / m_inFormat.nBlockAlign;


   if(m_inFormat.wFormatTag == WAVE_FORMAT_PCM)
   {
      if(m_inFormat.wBitsPerSample == 16)
      {
         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;

            m_smpld[nPos].l = short((pbData[j+1] << 8) | pbData[j]);
            j += 2;

            if(m_inFormat.nChannels == 2)
            {
               m_smpld[nPos].r = short((pbData[j+1] << 8) | pbData[j]);
               j += 2;
            }
         }
      }
      else if(m_inFormat.wBitsPerSample == 8)
      {
         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;

            m_smpld[nPos].l = pbData[j] - 128.0;
            j++;

            if(m_inFormat.nChannels == 2)
            {
               m_smpld[nPos].r = pbData[j] - 128.0;
               j++;
            }
         }
      }
      else if(m_inFormat.wBitsPerSample == 24)
      {
         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;

            if(pbData[j+2] & 0x80)
               m_smpld[nPos].l = int(0xff000000 | (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j]);
            else
               m_smpld[nPos].l = (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j];

            j += 3;

            if(m_inFormat.nChannels == 2)
            {
               if(pbData[j+2] & 0x80)
                  m_smpld[nPos].r = int(0xff000000 | (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j]);
               else
                  m_smpld[nPos].r = (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j];
   
               j += 3;
            }
         }
      }
      else if(m_inFormat.wBitsPerSample == 32)
      {
         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;

            m_smpld[nPos].l = int((pbData[j+3] << 24) | (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j]);
            j += 4;

            if(m_inFormat.nChannels == 2)
            {
               m_smpld[nPos].r = int((pbData[j+3] << 24) | (pbData[j+2] << 16) | (pbData[j+1] << 8) | pbData[j]);
               j += 4;
            }
         }
      }
   }
   else 
   {
      if(m_inFormat.wBitsPerSample == 32)
      {
         float fData;

         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;
            ::CopyMemory(&fData, &pbData[j], 4);
            m_smpld[nPos].l = fData;
            j += 4;

            if(m_inFormat.nChannels == 2)
            {
               ::CopyMemory(&fData, &pbData[j], 4);
               m_smpld[nPos].r = fData;
               j += 4;
            }
         }
      }
      else if(m_inFormat.wBitsPerSample == 64)
      {
         j = 0;

         for(i=0;i<nInSmplSize;i++)
         {
            nPos = (m_nSmplSize + i) & 0x1ffff;
            ::CopyMemory(&m_smpld[nPos].l, &pbData[j], 8);
            j += 8;

            if(m_inFormat.nChannels == 2)
            {
               ::CopyMemory(&m_smpld[nPos].r, &pbData[j], 8);
               j += 8;
            }
         }
      }
   }

   m_nSmplSize = (m_nSmplSize + nInSmplSize) & 0x1ffff;
}

void CConvertWav::Normalize(BYTE *pbInBuffer, int nInBufferLength, WAVEFORMATEX *pwf)
{  

   int i;

   if(pwf->wFormatTag == WAVE_FORMAT_PCM)
   {
      if(pwf->wBitsPerSample == 16)
      {
         short sData;
         int nData;

         for(i=0;i<nInBufferLength;i+=2)
         {
            ::CopyMemory(&sData, &pbInBuffer[i], 2);

            nData = (int)(sData * m_dNormalizeScale);

            if(nData < -32768)
               nData = -32768;

            if(32767 < nData)
               nData = 32767;

            pbInBuffer[i] = BYTE(nData & 0x000000ff);
            pbInBuffer[i+1] = BYTE((nData & 0x0000ff00) >> 8);
         }
      }
      else if(pwf->wBitsPerSample == 8)
      {
         int nData;

         for(i=0;i<nInBufferLength;i++)
         {
            nData = (int)((pbInBuffer[i] - 128.0) * m_dNormalizeScale) + 128;

            if(nData < 0)
               nData = 0;

            if(255 < nData)
               nData = 255;

            pbInBuffer[i] = (BYTE)nData;
         }
      }
      else if(pwf->wBitsPerSample == 24)
      {
         int nData;

         for(i=0;i<nInBufferLength;i+=3)
         {
            if(pbInBuffer[i+2] & 0x80)
            {
               nData = (int)( ( ((pbInBuffer[i+2] << 16) | (pbInBuffer[i+1] << 8) | pbInBuffer[i]) - 0x01000000) * m_dNormalizeScale );

               if(nData < -8388608)
                  nData = -8388608;
            }
            else
            {
               nData = (int)( ((pbInBuffer[i+2] << 16) | (pbInBuffer[i+1] << 8) | pbInBuffer[i]) * m_dNormalizeScale );

               if(8388607 < nData)
                  nData = 8388607;
            }

            pbInBuffer[i] = BYTE(nData & 0x000000ff);
            pbInBuffer[i+1] = BYTE((nData & 0x0000ff00) >> 8);
            pbInBuffer[i+2] = BYTE((nData & 0x00ff0000) >> 16);
         }
      }
      else 
      {
         int nData;
         double dData;

         for(i=0;i<nInBufferLength;i+=4)
         {
            ::CopyMemory(&nData, &pbInBuffer[i], 4);

            dData = nData * m_dNormalizeScale;

            if(dData < -2147483648.0)
               dData = -2147483648.0;

            if(2147483647.0 < dData)
               dData = -2147483647.0;

            nData = (int)dData;
            ::CopyMemory(&pbInBuffer[i], &nData, 4);
         }
      }
   }
   else 
   {
      if(pwf->wBitsPerSample == 32)
      {
         float fData;

         for(i=0;i<nInBufferLength;i+=4)
         {
            ::CopyMemory(&fData, &pbInBuffer[i], 4);

            fData = (float)(fData * m_dNormalizeScale);

            if(fData < -1.0f)
               fData = -1.0f;

            if(1.0f < fData)
               fData = 1.0f;

            ::CopyMemory(&pbInBuffer[i], &fData, 4);
         }
      }
      else 
      {
         double dData;

         for(i=0;i<nInBufferLength;i+=8)
         {
            ::CopyMemory(&dData, &pbInBuffer[i], 8);

            dData = dData * m_dNormalizeScale;

            if(dData < -1.0)
               dData = -1.0;

            if(1.0 < dData)
               dData = 1.0;

            ::CopyMemory(&pbInBuffer[i], &dData, 8);
         }
      }
   }
}

HRESULT CConvertWav::Convert(BYTE *pbInBuffer, int nInBufferLength, BYTE *pbOutBuffer, int *pnOutBufferLength)
{  

   HRESULT hr;

   if(m_inFormat.nSamplesPerSec == m_outFormat.nSamplesPerSec)
   {  

      if(m_inFormat.wFormatTag == m_outFormat.wFormatTag && m_inFormat.nChannels == m_outFormat.nChannels && m_inFormat.wBitsPerSample == m_outFormat.wBitsPerSample)
      {  

         
         ::CopyMemory(pbOutBuffer, pbInBuffer, nInBufferLength);
         *pnOutBufferLength = nInBufferLength;
      }
      else if( (m_inFormat.wFormatTag != m_outFormat.wFormatTag) || (m_inFormat.wBitsPerSample != m_outFormat.wBitsPerSample) )
      {  

         if(m_inFormat.nChannels == m_outFormat.nChannels)
         {  
            hr = m_pConvertBits->Convert(pbInBuffer, nInBufferLength, pbOutBuffer, pnOutBufferLength);
            if(FAILED(hr)) return S_FALSE;
         }
         else if(m_inFormat.nChannels == 1 && m_outFormat.nChannels == 2)
         {  

            hr = m_pConvertBits->Convert(pbInBuffer, nInBufferLength, pbOutBuffer, pnOutBufferLength);
            if(FAILED(hr)) return S_FALSE;

            int nLength = *pnOutBufferLength;
            hr = m_pConvertChannels->Convert(pbOutBuffer, nLength, pnOutBufferLength);
            if(FAILED(hr)) return S_FALSE;
         }
         else if(m_inFormat.nChannels == 2 && m_outFormat.nChannels == 1)
         {  

            hr = m_pConvertChannels->Convert(pbInBuffer, nInBufferLength, pnOutBufferLength);
            if(FAILED(hr)) return S_FALSE;

            int nLength = *pnOutBufferLength;
            hr = m_pConvertBits->Convert(pbInBuffer, nLength, pbOutBuffer, pnOutBufferLength);
            if(FAILED(hr)) return S_FALSE;
         }
      }
      else if(m_inFormat.nChannels != m_outFormat.nChannels)
      {  
         hr = m_pConvertChannels->Convert(pbInBuffer, nInBufferLength, pbOutBuffer, pnOutBufferLength);

      }

      if(m_bUseNormalize == true) 
         Normalize(pbOutBuffer, *pnOutBufferLength, &m_outFormat);
   }
   else  
   {
      
      ConvertToDouble(pbInBuffer, nInBufferLength);

      
      if(m_bUseSSE2 == true)
         m_pConvertSamples->ConvertSSE2(m_smpld, nInBufferLength / m_inFormat.nBlockAlign, pbOutBuffer, pnOutBufferLength);
      else
         m_pConvertSamples->Convert(m_smpld, nInBufferLength / m_inFormat.nBlockAlign, pbOutBuffer, pnOutBufferLength);
   }
 
   return S_OK;
}

HRESULT CConvertWav::OnTransform(IMediaSample *pInSample, IMediaSample *pOutSample)
{  

   if(0 <= m_nWait && m_nWait <= 3)
   {  

      
      int nWait = int((m_llInputSize + pInSample->GetActualDataLength()) * 10 / m_inFormat.nAvgBytesPerSec) - int(m_llInputSize * 10 / m_inFormat.nAvgBytesPerSec);

      if(0< nWait)
         ::Sleep(nWait * (4 - m_nWait));
   }

   BYTE *pbInBuffer = NULL;
   BYTE *pbOutBuffer = NULL;

   
   pInSample->GetPointer(&pbInBuffer);

   
   pOutSample->GetPointer(&pbOutBuffer);

   
   m_llInputSize += pInSample->GetActualDataLength();

   
   int nOutBufferLength = 0;
   Convert(pbInBuffer, (int)pInSample->GetActualDataLength(), pbOutBuffer, &nOutBufferLength);

   
   pOutSample->SetSyncPoint(TRUE);
   pOutSample->SetDiscontinuity(FALSE);
   pOutSample->SetPreroll(FALSE);
   pOutSample->SetActualDataLength(nOutBufferLength);

   
   REFERENCE_TIME rtStop;

   rtStop = m_rtCurrent + (REFERENCE_TIME)((LONGLONG)nOutBufferLength * ONE_SECOND / m_outFormat.nAvgBytesPerSec);
   pOutSample->SetTime(&m_rtCurrent, &rtStop);
   m_rtCurrent = rtStop;

   return S_OK;
}

HRESULT CConvertWav::OnStop(bool bEos)
{
   if(bEos == true)
   {  
      

      if(m_inFormat.nSamplesPerSec != m_outFormat.nSamplesPerSec)
      {
         HRESULT hr;
         int nLength = 0;
         int nOutBufferLength = 0;

         
         BYTE *pbInBuffer = NULL;
         BYTE *pbOutBuffer = NULL;
         IMediaSample *pOutSample = NULL;

         hr = m_pOutput->GetDeliveryBuffer(&pOutSample, NULL, NULL, 0);
         if FAILED(hr)
            return S_FALSE;

         pOutSample->GetPointer(&pbOutBuffer);
         pOutSample->SetSyncPoint(TRUE);
         pOutSample->SetDiscontinuity(FALSE);
         pOutSample->SetPreroll(FALSE);
         pOutSample->SetActualDataLength(0);

         for(int i=0;i<1000;i++)
         {  
            m_smpld[(m_nSmplSize + i) & 0x1ffff].l = 0.0;
            m_smpld[(m_nSmplSize + i) & 0x1ffff].r = 0.0;
         }

         m_pConvertSamples->ConvertEnd(m_smpld, m_nSmplSize, pbOutBuffer, &nLength);
         nOutBufferLength = nOutBufferLength + nLength;

         
         pOutSample->SetActualDataLength(nOutBufferLength);

         
         if(nOutBufferLength > 0)
            hr = m_pOutput->Deliver(pOutSample);

         
         SAFE_RELEASE(pOutSample);
      }
   }

   return S_OK;
}

HRESULT CConvertWav::GetInFormat(WAVEFORMATEX *pwf)
{
   CheckPointer(pwf, E_POINTER);
   ::CopyMemory(pwf, &m_inFormat, sizeof(m_inFormat));
   return S_OK;
}

HRESULT CConvertWav::GetOutFormat(WAVEFORMATEX *pwf)
{
   CheckPointer(pwf, E_POINTER);
   ::CopyMemory(pwf, &m_outFormat, sizeof(m_outFormat));
   return S_OK;
}

HRESULT CConvertWav::SetOutFormat(WAVEFORMATEX *pwf)
{
   CheckPointer(pwf, E_POINTER);

   
   if(m_pOutput->IsConnected() == TRUE)
      return E_FAIL;

   
   if(pwf->wFormatTag != WAVE_FORMAT_PCM && pwf->wFormatTag != WAVE_FORMAT_IEEE_FLOAT)
      return E_FAIL;

   if(pwf->nBlockAlign != (pwf->nChannels * pwf->wBitsPerSample / 8))
      return E_FAIL;

   if(pwf->nAvgBytesPerSec != pwf->nBlockAlign * pwf->nSamplesPerSec)
      return E_FAIL;

   
   ::CopyMemory(&m_outFormat, pwf, sizeof(m_outFormat));
   return S_OK;
}

HRESULT CConvertWav::GetTransformedBytes(LONGLONG *llTransformedBytes)
{
   CAutoLock lock(&m_Lock);

   *llTransformedBytes = m_llInputSize;

   return S_OK;
}

HRESULT CConvertWav::SetWait(int nWait)
{
   CAutoLock lock(&m_Lock);

   m_nWait = nWait;

   return S_OK;
}

HRESULT CConvertWav::SetNormalize(double dNormalize)
{  
   m_bUseNormalize = true;
   m_dNormalizeScale = dNormalize;

   return S_OK;
}

HRESULT CConvertWav::OnConnectInPin(const CMediaType *pmtIn)
{  
   

   if(pmtIn->subtype != MEDIASUBTYPE_PCM && pmtIn->subtype != MEDIASUBTYPE_IEEE_FLOAT)
      return S_FALSE;

   if(pmtIn->formattype != FORMAT_WaveFormatEx)
      return S_FALSE;

   WAVEFORMATEX *pwf = (WAVEFORMATEX *)pmtIn->Format();

   
   int i;
   int nAcceptSamplingRate[] = {8000, 11025, 12000, 16000, 22050, 24000, 32000, 44100, 48000, 88200, 96000, 176400, 192000};

   for(i=0;i<13;i++)
   {
      if(pwf->nSamplesPerSec == nAcceptSamplingRate[i])
         break;
   }

   if(i == 13)
      return S_FALSE;

   
   if(pwf->wFormatTag == WAVE_FORMAT_PCM)
   {
      if(pwf->wBitsPerSample != 8 && pwf->wBitsPerSample != 16 && pwf->wBitsPerSample != 24 && pwf->wBitsPerSample != 32)
         return S_FALSE;
   }
   else if(pwf->wFormatTag == WAVE_FORMAT_IEEE_FLOAT)
   {
      if(pwf->wBitsPerSample != 32 && pwf->wBitsPerSample != 64)
         return S_FALSE;
   }
   else
      return S_FALSE;

   
   if(pwf->nChannels != 1 && pwf->nChannels != 2)
      return S_FALSE;

   
   ::CopyMemory(&m_inFormat, pwf, sizeof(m_inFormat));

   
   ::CopyMemory(&m_outFormat, &m_inFormat, sizeof(m_outFormat));

   return S_OK;
}

HRESULT CConvertWav::OnConnectOutPin(const CMediaType *pmtIn, int nInBufferSize, CMediaType *pmtOut, int *pnOutBufferSize)
{  

   
   pmtOut->SetType(&MEDIATYPE_Audio);
   pmtOut->SetSubtype(&MEDIASUBTYPE_PCM);
   pmtOut->SetTemporalCompression(FALSE);
   pmtOut->SetSampleSize(0);
   pmtOut->SetFormatType(&FORMAT_WaveFormatEx);
   pmtOut->SetFormat((BYTE*)&m_outFormat, sizeof(m_outFormat));

   
   *pnOutBufferSize = (int)((double)nInBufferSize / (double)m_inFormat.nAvgBytesPerSec * (double)m_outFormat.nAvgBytesPerSec * 1.1);

   return S_OK;
}

STDMETHODIMP CConvertWav::OnQueryInterface(REFIID riid, void ** ppv)
{
   if(riid == IID_IConvertWavInterface)
      return GetInterface((IConvertWavInterface *)this, ppv);

   return CTransformFilter::NonDelegatingQueryInterface(riid, ppv);
}

CUnknown * WINAPI CConvertWav::CreateInstance(LPUNKNOWN punk, HRESULT *phr)
{  

   ASSERT(phr);
    
   CConvertWav *pNewObject = new CConvertWav(punk, phr);
   if (pNewObject == NULL)
   {
      if(phr != NULL)
         *phr = E_OUTOFMEMORY;
   }

   return dynamic_cast<CUnknown *>(pNewObject);
}

