/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA 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 3 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 MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */


#include <stdexcept>
#include <string>
#include <vector>

#include <fftw3.h>

#include <mia/core/msgstream.hh>

NS_MIA_BEGIN


template <typename T>
class TCSTPlan
{
public:
       typedef typename T::value_type value_type;

       TCSTPlan(fftwf_r2r_kind forward, std::vector<int> size);

       virtual ~TCSTPlan();

       template <typename D>
       void execute(const D& in_data, D& out_data) const;

       const std::vector<int>& get_size() const;

private:
       mutable value_type *m_in;
       mutable value_type *m_out;

       std::vector<int> m_size;
       float m_scale;
       size_t m_n;

       fftwf_plan m_forward_plan;
       fftwf_plan m_backward_plan;

       virtual void do_execute(value_type *buffer) const = 0;

       struct Scale {
              Scale(float scale): m_scale(scale)
              {
              }
              value_type operator ()(const value_type& x) const
              {
                     return m_scale * x;
              }
              float m_scale;
       };
};

template <typename T>
const std::vector<int>& TCSTPlan<T>::get_size() const
{
       return m_size;
}

template <typename T>
TCSTPlan<T>::TCSTPlan(fftwf_r2r_kind forward, std::vector<int> size):
       m_size(size),
       m_scale(1.0)
{
       std::string msg;
       size_t rank = m_size.size();
       fftwf_r2r_kind backward;

       switch (forward) {
       case FFTW_RODFT00:
              for (size_t i = 0; i < rank; ++i)
                     m_scale *= 0.5 / ( size[i] + 1.0);

              break;

       case FFTW_REDFT00:
              for (size_t i = 0; i < rank; ++i)
                     m_scale *= 0.5 / ( size[i] - 1.0);

              break;

       default:
              for (size_t i = 0; i < rank; ++i)
                     m_scale *= 0.5 / size[i];

              break;
       }

       switch (forward) {
       case FFTW_RODFT00:
       case FFTW_REDFT00:
       case FFTW_RODFT11:
       case FFTW_REDFT11:
              backward = forward;
              break;

       case FFTW_RODFT10:
              backward = FFTW_RODFT01;
              break;

       case FFTW_REDFT10:
              backward = FFTW_REDFT01;
              break;

       case FFTW_RODFT01:
              backward = FFTW_RODFT10;
              break;

       case FFTW_REDFT01:
              backward = FFTW_REDFT10;
              break;

       default:
              throw std::invalid_argument("TCSTPlan: unknown transformtion requested");
       }

       std::vector<fftwf_r2r_kind> fw_kind(rank, forward);
       std::vector<fftwf_r2r_kind> bw_kind(rank, backward);
       const int howmany = sizeof(value_type) / sizeof(float);
       cvdebug() << "howmany = " << howmany << "\n";
       m_n = 1;

       for (size_t i = 0; i < rank; ++i)
              m_n *= m_size[i];

       if (NULL == (m_in = (value_type *) fftwf_malloc(sizeof(T) * m_n))) {
              msg = "unable to allocate FFTW data";
              goto in_fail;
       }

       if ( NULL == (m_out  = (value_type *) fftwf_malloc(sizeof(T) * m_n))) {
              msg = "unable to allocate FFTW data";
              goto out_fail;
       }

       if (0 == (m_forward_plan = fftwf_plan_many_r2r(rank, &size[0], howmany,
                                  m_in, NULL, howmany, 1,
                                  m_out, NULL, howmany, 1,
                                  &fw_kind[0], FFTW_ESTIMATE))) {
              msg = "unable to create FFTW forward plan";
              goto plan_fw_fail;
       }

       if (0 == (m_backward_plan = fftwf_plan_many_r2r(rank, &size[0], howmany,
                                   m_out, NULL, howmany, 1,
                                   m_in,  NULL, howmany, 1,
                                   &bw_kind[0], FFTW_ESTIMATE))) {
              msg = "unable to create FFTW backward plan";
              goto  plan_bw_fail;
       }

       return;
plan_bw_fail:
       fftwf_destroy_plan(m_forward_plan);
plan_fw_fail:
       fftwf_free(m_out);
out_fail:
       fftwf_free(m_in);
in_fail:
       throw std::runtime_error(msg);
}

template <typename T>
TCSTPlan<T>::~TCSTPlan()
{
       fftwf_destroy_plan(m_backward_plan);
       fftwf_destroy_plan(m_forward_plan);
       fftwf_free(m_out);
       fftwf_free(m_in);
}

template <typename T>
template <typename D>
void TCSTPlan<T>::execute(const D& in_data, D& out_data) const
{
       assert(m_n == in_data.size());
       assert(m_n == out_data.size());
       copy(in_data.begin(), in_data.end(), m_in);
       fftwf_execute( m_forward_plan);
       do_execute(m_out);
       fftwf_execute(m_backward_plan);
       transform( m_in,  m_in + m_n, out_data.begin(), Scale(m_scale) );
}

NS_MIA_END
