/* -*- 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(2, 2, 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_peakList(0)
		, m_peakHistory(0)
		, m_peak_id(0)
    {
		m_pX = new float[vlen];
		m_pXs = new float[vlen];
		m_pXm = new float[vlen];
		m_pXd = new float[vlen];
		m_pXds = new float[vlen];
		m_pPeakBox_rel = new float[vlen];
		m_pPeakBox_abs = new float[vlen];
		
		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_pXds, 0, vlen*sizeof(float));
		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_pXds;
		delete [] m_pXd;
		delete [] m_pXm;
		delete [] m_pXs;
		delete [] m_pX;
    }

	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;
		}
	}

	std::vector<int> peak_detect_impl::findPeakPos(float* pSrc, int length, float thresh)
	{
		std::vector<int> result;
		for (int i=0; i < length; i++)
		{
			if (pSrc[i] > thresh)
			{
				result.push_back(i);
			}
		}
		return result;
	}

    int
    peak_detect_impl::work(int noutput_items,
        gr_vector_const_void_star &input_items,
        gr_vector_void_star &output_items)
    {
		const int DATA = 0;
		const int PBOX = 1;
		
		int nitems_per_block = this->output_signature()->sizeof_stream_item(0)/sizeof(float);

		float *iptr = (float *)input_items[DATA];
		float *optr_data = (float *)output_items[DATA];
		float *optr_pbox = (float *)output_items[PBOX];
		
		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);
				m_peak_id = 0;
			}
			
			// 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];
				float Xds = m_pXs[i]-m_pXm[i];
				m_pXd[i] = Xd;
				m_pXds[i] = Xds;
			}

			// -----------------------------------------------------------------
			// Create peak list
			// -----------------------------------------------------------------
			// Find peak candidates by thresholding
			std::vector<int> peakPos = findPeakPos(m_pXds, nitems_per_block, m_peak_height_min);
			
			// -----------------------------------------------------------------
			// Search true maximum peak using hill climbing
			// -----------------------------------------------------------------
			std::vector<int> newPeakPos;
			int numPeaksLast = 0;
			while(1)
			{
				int numPeaksClimbed = 0;
				for (int i = 0; i < peakPos.size(); i++)
				{
					int pos = peakPos[i];
					while(1)
					{
						float y0 = m_pXs[pos];
						int pos_p = std::min(nitems_per_block-1, pos+1);
						int pos_n = std::max(0, pos-1);
						if (m_pXs[pos_p] > y0)
						{
							pos = pos_p;
						}
						else if (m_pXs[pos_n] > y0)
						{
							pos = pos_n;
						}
						else
						{
							addToList(newPeakPos, pos);
							numPeaksClimbed++;
							break;
						}
					}
				}
				if (numPeaksClimbed == numPeaksLast)
				{
					break;
				}
				numPeaksLast = numPeaksClimbed;
			}
			peakPos = newPeakPos;
			
			// -----------------------------------------------------------------
			// Sort peaks
			// -----------------------------------------------------------------
			sort<float>(m_pXds, peakPos);
			
			// -----------------------------------------------------------------
			// Find peak boxes
			// -----------------------------------------------------------------
			m_peakList.clear();
			for (int i = 0; i < peakPos.size(); i++)
			{
				int pbh = 0;
				int pos = peakPos[i];
				int left = pos;
				int right = pos;
				float nom = 10;
				int left_found = 0;
				int right_found = 0;
				while(!(left_found & right_found))
				{
					// Determine peak width left from center (pos)
					if (m_pXds[left] > nom)
					{
						if (left > 0)
						{
							left--;
						}
						else
						{
							break;
						}
					}
					else
					{
						left_found = 1;
					}
					
					// Determine peak width right from center (pos)
					if (m_pXds[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;
				}
				
				Peak peak(m_peak_id, pos, pbh, m_pXds[pos], m_pXs[pos]);
				
				// Intersection check
				bool intersection = false;
				for (int i = 0; i < m_peakList.size(); i++)
				{
					if (m_peakList[i].isInRange(peak))
					{
						intersection = true;
					}
				}
				
				if (!intersection)
				{
					m_peakList.push_back(peak);
					m_peak_id++;
				}
			}
			
			// -----------------------------------------------------------------
			// Peak management
			// -----------------------------------------------------------------
			temp.clear();
			for (int i = 0; i < m_peakList.size(); i++)
			{
				Peak &peak = m_peakList[i];
		
				bool found = false;
				for (int j=0; j < m_peakHistory.size(); j++)
				{
					if (m_peakHistory[j].isInRange(peak))
					{
						temp.push_back(m_peakHistory[j].update(peak));
						found = true;
						break;
					}
				}
				
				if (!found)
				{
					peak.isValid = 1;
					temp.push_back(peak);
					fprintf(stderr, "New Peak #%u at %d\n", (int)peak.id, (int)peak.pos);
				}
			}
			
			for (int j=0; j < m_peakHistory.size(); j++)
			{
				Peak &peak = m_peakHistory[j];
				if (!peak.isValid)
				{
					fprintf(stderr, "Lost Peak #%u at %d\n", (int)peak.id, (int)peak.pos);
				}
			}
			m_peakHistory.clear();
			m_peakHistory = temp;
			
			// -----------------------------------------------------------------
			// Create peak box data
			// -----------------------------------------------------------------
			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 < m_peakHistory.size(); i++)
			{
				Peak &peak = m_peakHistory[i];
				// Use peak width for peak box
				int pbox_start = std::max(0, peak.left());
				int pbox_end = std::min(nitems_per_block-1, peak.right());
				
				for (int j=pbox_start; j <= pbox_end; j++)
				{
					m_pPeakBox_abs[j] = peak.height_abs;
					m_pPeakBox_rel[j] = peak.height_rel;
				}
			}

			// -----------------------------------------------------------------
			// Output peak box data
			// -----------------------------------------------------------------
			for (int i = 0; i < nitems_per_block; i++)
			{
				*(optr_data++) = m_pXs[i];
				*(optr_pbox++) = m_pPeakBox_abs[i];
			}

			// -----------------------------------------------------------------
			// Output tags
			// -----------------------------------------------------------------
			for (int i = 0; i < m_peakHistory.size(); i++)
			{
				// Create peak tags
				Peak &peak = m_peakHistory[i];
				pmt::pmt_t tag_key = pmt::string_to_symbol(std::string("pos_id_") + std::to_string(peak.id));
				pmt::pmt_t tag_value = pmt::from_long(peak.pos);
				add_item_tag(DATA, nitems_written(DATA) + k, tag_key, tag_value, m_tag_id);
				peak.isValid = false;
			}

			m_frameCount++;
		}

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

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

