/* -*- c++ -*- */
/* 
 * Copyright 2019 Jay Arrowfield.
 * 
 * This 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, or (at your option)
 * any later version.
 * 
 * This software 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 software; see the file COPYING.  If not, write to
 * the Free Software Foundation, Inc., 51 Franklin Street,
 * Boston, MA 02110-1301, USA.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <gnuradio/io_signature.h>
#include "peak_detect_impl.h"

namespace gr {
  namespace jay {

    peak_detect::sptr
    peak_detect::make(int vlen, float framerate, float peak_height_min, int peak_width_min, int peak_width_max, float alpha_s, float alpha_m)
    {
      return gnuradio::get_initial_sptr
        (new peak_detect_impl(vlen, framerate, peak_height_min, peak_width_min, peak_width_max, alpha_s, alpha_m));
    }

    /*
     * The private constructor
     */
    peak_detect_impl::peak_detect_impl(int vlen, float framerate, float peak_height_min, int peak_width_min, int peak_width_max, float alpha_s, float alpha_m)
		: gr::sync_block("peak_detect"
		, gr::io_signature::make(1, 1, vlen*sizeof(float))
        , gr::io_signature::make(1, 1, vlen*sizeof(float)))
		, m_framerate (framerate)
		, m_peak_height_min (peak_height_min)
		, m_peak_width_min (peak_width_min)
		, m_peak_width_max (peak_width_max)
		, m_alpha_s (alpha_s)
		, m_alpha_m(alpha_m)
		, m_frameCount(0)
    {
		m_pPeakData = new float[vlen];
		m_pX = new float[vlen];
		m_pXs = new float[vlen];
		m_pXm = new float[vlen];
		m_pXd = new float[vlen];
		m_pPeakPos = new int[vlen];
		m_pPeakPos_final = new int[vlen];
		m_pPeakBox_rel = new float[vlen];
		m_pPeakBox_abs = new float[vlen];
		
		memset(m_pPeakData, 0, vlen*sizeof(int));
		memset(m_pX, 0, vlen*sizeof(float));
		memset(m_pXs, 0, vlen*sizeof(float));
		memset(m_pXm, 0, vlen*sizeof(float));
		memset(m_pXd, 0, vlen*sizeof(float));
		memset(m_pPeakPos, 0, vlen*sizeof(int));
		memset(m_pPeakPos_final, 0, vlen*sizeof(int));
		memset(m_pPeakBox_rel, 0, vlen*sizeof(float));
		memset(m_pPeakBox_abs, 0, vlen*sizeof(float));
		
		std::stringstream str;
		str << name() << unique_id();
		m_tag_id = pmt::string_to_symbol(str.str());
	}

    /*
     * Our virtual destructor.
     */
    peak_detect_impl::~peak_detect_impl()
    {
		delete [] m_pPeakBox_abs;
		delete [] m_pPeakBox_rel;
		delete [] m_pPeakPos_final;
		delete [] m_pPeakPos;
		delete [] m_pXd;
		delete [] m_pXm;
		delete [] m_pXs;
		delete [] m_pX;
		delete [] m_pPeakData;
    }

	float peak_detect_impl::mean(float* pSrc, int length)
	{
		float res = 0;
		for (int i=0; i < length; i++)
		{
			res += pSrc[i];
		}
		return res / length;
	}

	void peak_detect_impl::fill(float* pDst, int length, float value)
	{
		for (int i=0; i < length; i++)
		{
			pDst[i] = value;
		}
	}
	int peak_detect_impl::threshold(float *pDst, float *pSrc, int length, float thresh)
	{
		for (int i=0; i < length; i++)
		{
			pDst[i] = (int)(pSrc[i] >= thresh);
		}
		
		return length;
	}
	
	int peak_detect_impl::findPeakPos(int* pDst, float* pSrc, int length)
	{
		int count = 0;
		for (int i=0; i < length; i++)
		{
			if (pSrc[i] > 0)
			{
				pDst[count++] = i;
			}
		}
		return count;
	}
	
    int
    peak_detect_impl::work(int noutput_items,
        gr_vector_const_void_star &input_items,
        gr_vector_void_star &output_items)
    {
		int nitems_per_block = this->output_signature()->sizeof_stream_item(0)/sizeof(float);

		float *iptr = (float *)input_items[0];
		float *optr = (float *)output_items[0];
		
		for (int k = 0; k < noutput_items; k++)
		{
			// Initialization on first frame
			if (m_frameCount == 0)
			{
				float xmean = mean(iptr, nitems_per_block);
				fill(m_pXs, nitems_per_block, xmean);
				fill(m_pXm, nitems_per_block, xmean);
			}
			
			// Create input vector
			for (int i = 0; i < nitems_per_block; i++)
			{
				float x = *(iptr++);
				m_pX[i] = x;
			}
			
			// -----------------------------------------------------------------
			// Update of input statistics
			// -----------------------------------------------------------------
			for (int i = 0; i < nitems_per_block; i++)
			{
				float x = m_pX[i];
				
				// Update of Xs
				m_pXs[i] = (1.f-m_alpha_s)*m_pXs[i] + m_alpha_s*x;
				
				// Update of Xm
				m_pXm[i] = (1.f-m_alpha_m)*m_pXm[i] + m_alpha_m*x;
	
				// Conditional update of Xm
				float last = m_pXm[0];
				for (int j=1; j < nitems_per_block; j++)
				{
					if (m_pXm[j] < (last+2))
					{
						last = (1.f-m_alpha_m)*last + m_alpha_m*m_pXm[j];
					}
					else
					{
						m_pXm[j] = last;	
					}
				}
			}

			// -----------------------------------------------------------------
			// Peak detection
			// -----------------------------------------------------------------
			// Detrend onput data X
			for (int i = 0; i < nitems_per_block; i++)
			{
				// Create peak distances by subtracting Xm
				float Xd = m_pX[i]-m_pXm[i];
				m_pXd[i] = Xd;
//				*(optr++) = Xd;
			}

			// -----------------------------------------------------------------
			// Create peak list
			// -----------------------------------------------------------------
			// Find peak candidates by thresholding
			threshold(m_pPeakData, m_pXd, nitems_per_block, m_peak_height_min);
			int numPeaks = findPeakPos(m_pPeakPos, m_pPeakData, nitems_per_block);
			
			// -----------------------------------------------------------------
			// Search true maximum peak using hill climbing
			// -----------------------------------------------------------------
			int numPeaksLast = 0;
			while(1)
			{
				int numPeaksClimbed = 0;
				for (int i = 0; i < numPeaks; i++)
				{
					int pos = m_pPeakPos[i];
					while(1)
					{
						float y0 = m_pX[pos];
						int pos_p = std::min(nitems_per_block-1, pos+1);
						int pos_n = std::max(0, pos-1);
						if (m_pX[pos_p] > y0)
						{
							m_pPeakData[pos] = 0;
							m_pPeakData[pos_n] = 0;
							pos = pos_p;
						}
						else if (m_pX[pos_n] > y0)
						{
							m_pPeakData[pos] = 0;
							m_pPeakData[pos_p] = 0;
							pos = pos_n;
						}
						else
						{
							m_pPeakData[pos_n] = 0;
							m_pPeakData[pos_p] = 0;
							numPeaksClimbed++;
							break;
						}
					}
				}
				if (numPeaksClimbed == numPeaksLast)
				{
					break;
				}
				numPeaksLast = numPeaksClimbed;
			}
			numPeaks = findPeakPos(m_pPeakPos, m_pPeakData, nitems_per_block);
			
			// -----------------------------------------------------------------
			// Sort peaks
			// -----------------------------------------------------------------
			sort(m_pXd, m_pPeakPos, numPeaks);
			
			// -----------------------------------------------------------------
			// Construct peak boxes
			// -----------------------------------------------------------------
			int numPeaks_final = 0;
			memset(m_pPeakPos_final, 0, nitems_per_block*sizeof(int));
			memcpy(m_pPeakBox_abs, m_pXm, nitems_per_block*sizeof(float));
			memset(m_pPeakBox_rel, 0, nitems_per_block*sizeof(float));
			
			for (int i = 0; i < numPeaks; i++)
			{
				int pbh = 0;
				int pos = m_pPeakPos[i];
				int left = pos;
				int right = pos;
				float nom = m_pXd[pos];
				int left_found = 0;
				int right_found = 0;
				while(!(left_found & right_found))
				{
					// Determine peak width left from center (pos)
					if (m_pXd[left] > nom)
					{
						if (left > 0)
						{
							left--;
						}
						else
						{
							break;
						}
					}
					else
					{
						left_found = 1;
					}
					
					// Determine peak width right from center (pos)
					if (m_pXd[right] > nom)
					{
						if (right < (nitems_per_block-1))
						{
							right++;
						}
						else
						{
							break;
						}
					}
					else
					{
						right_found = 1;
					}
					
					int pbh_left = pos - left;
					int pbh_right = right - pos;

					// Take larger peak width
					pbh = std::max(pbh_left, pbh_right);

					// Ensure peak box is at least pb_min
					pbh = std::max(pbh, m_peak_width_min);
					
					if (pbh > m_peak_width_max)
					{
						pbh = 0;
						break;
					}
				}
				
				if (pbh == 0)
				{
					continue;
				}
				
				// Use peak width for peak box
				int pbox_start = std::max(0, pos-pbh);
				int pbox_end = std::min(nitems_per_block-1, pos+pbh);

				// Find intersection
				float box_height_abs = m_pX[pos];
				float box_height_rel = m_pXd[pos];
				int does_intersect = 0;
				for (int n=pbox_start; n <= pbox_end; n++)
				{
					if (m_pPeakBox_rel[n] > box_height_rel)
					{
						does_intersect = 1;
						break;
					}
				}
				if (!does_intersect)
				{
					for (int n=pbox_start; n <= pbox_end; n++)
					{
						m_pPeakBox_abs[n] = box_height_abs;
						m_pPeakBox_rel[n] = box_height_rel;
					}
					m_pPeakPos_final[numPeaks_final++] = pos;
				}
			}
			
			// -----------------------------------------------------------------
			// Output peak box
			// -----------------------------------------------------------------
			for (int i = 0; i < nitems_per_block; i++)
			{
				*(optr++) = m_pPeakBox_abs[i];
			}

			// -----------------------------------------------------------------
			// Output tags
			// -----------------------------------------------------------------
			for (int i = 0; i < numPeaks_final; i++)
			{
				// Create peak tags
				pmt::pmt_t tag_key = pmt::string_to_symbol("id");
				pmt::pmt_t tag_value = pmt::from_long(i);
				add_item_tag(0, nitems_written(0) + m_pPeakPos_final[i], tag_key, tag_value, m_tag_id);
				tag_key = pmt::string_to_symbol("pos");
				tag_value = pmt::from_long(m_pPeakPos_final[i]);
				add_item_tag(0, nitems_written(0) + m_pPeakPos_final[i], tag_key, tag_value, m_tag_id);
			}

			m_frameCount++;
		}

		// Tell runtime system how many output items we produced.
		return noutput_items;
    }

  } /* namespace jay */
} /* namespace gr */