/* -*- 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 bandwidth, float fcenter, float peak_thresh, float peak_hysteresis, float peak_width_min, float peak_width_max, float alpha_s, float alpha_m)
    {
      return gnuradio::get_initial_sptr
        (new peak_detect_impl(vlen, framerate, bandwidth, fcenter, peak_thresh, peak_hysteresis, peak_width_min, peak_width_max, alpha_s, alpha_m));
    }

    /*
     * The private constructor
     */
    peak_detect_impl::peak_detect_impl(int vlen, float framerate, float bandwidth, float fcenter, float peak_thresh, float peak_hysteresis, float peak_width_min, float peak_width_max, float alpha_s, float alpha_m)
		: gr::block("peak_detect"
		, gr::io_signature::make(1, 1, vlen*sizeof(float))
        , gr::io_signature::make(0, 2, vlen*sizeof(float)))
		, m_vlen(vlen)
		, m_framerate (framerate)
		, m_bandwidth (bandwidth)
		, m_fcenter (fcenter)
		, m_peak_thresh (peak_thresh)
		, m_peak_hysteresis (peak_hysteresis)
		, m_peak_width_min ((int)((float)m_vlen/bandwidth*peak_width_min + 0.5f))
		, m_peak_width_max ((int)((float)m_vlen/bandwidth*peak_width_max + 0.5f))
		, 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());
		
		m_msg_port_in = pmt::mp("control");
		m_msg_port_out = pmt::mp("status");
		message_port_register_in(m_msg_port_in);
		message_port_register_out(m_msg_port_out);
		set_msg_handler(m_msg_port_in, boost::bind(&peak_detect_impl::set_msg, this, _1));		
	}

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

	void peak_detect_impl::forecast(int noutput_items, gr_vector_int& ninput_items_required)
	{
		ninput_items_required[0] = noutput_items;
	}
	
	int peak_detect_impl::general_work (int noutput_items, gr_vector_int &ninput_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 = nullptr;
		float *optr_pbox = nullptr;
	
		if (output_items.size() > 0)
		{
			optr_data = (float *)output_items[DATA];
			optr_pbox = (float *)output_items[PBOX];
		}
		
		for (int k = 0; k < ninput_items[DATA]; 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_thresh);
			
			// -----------------------------------------------------------------
			// Search true maximum peak using hill climbing
			// -----------------------------------------------------------------
			std::vector<Peak> peaks;
			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
						{
							float freq_shift_Hz = m_bandwidth * (float)(pos-m_vlen/2)/m_vlen;
							float freq_Hz = m_fcenter + freq_shift_Hz;
							Peak peak(m_peak_id++, pos, freq_Hz);
							peak.height_rel = m_pXds[pos];
							addToListByPos(peaks, peak);
							numPeaksClimbed++;
							break;
						}
					}
				}
				if (numPeaksClimbed == numPeaksLast)
				{
					break;
				}
				numPeaksLast = numPeaksClimbed;
			}
			
			// -----------------------------------------------------------------
			// Sort peaks
			// -----------------------------------------------------------------
			sort(peaks);
			
			// -----------------------------------------------------------------
			// Find peak boxes
			// -----------------------------------------------------------------
			m_peakList.clear();
			for (int i = 0; i < peaks.size(); i++)
			{
				Peak &peak = peaks[i];
				bool success = peak.makeBox(m_peak_width_min, m_peak_width_max, m_pXds, nitems_per_block);
				if (success)
				{
					// Intersection check
					bool intersection = false;
					for (int i = 0; i < m_peakList.size(); i++)
					{
						if (m_peakList[i].isIntersect(peak))
						{
							intersection = true;
							break;
						}
					}

					if (!intersection)
					{
						peak.heightUpdate_rel(m_pXds, nitems_per_block);
						peak.heightUpdate_abs(m_pXs, nitems_per_block);
						m_peakList.push_back(peak);
					}
				}
			}
			
			// -----------------------------------------------------------------
			// Peak management
			// -----------------------------------------------------------------
			temp.clear();
			for (int j=0; j < m_peakHistory.size(); j++)
			{
				Peak &peak = m_peakHistory[j];
				peak.heightUpdate_rel(m_pXds, nitems_per_block);
				peak.heightUpdate_abs(m_pXs, nitems_per_block);
				
				if (peak.height_rel >= (m_peak_thresh-m_peak_hysteresis))
				{
					temp.push_back(peak);
				}
				else
				{
					peak.setState(0);
					// Output tags
					pmt::pmt_t dict = peak.msg();
					message_port_pub(m_msg_port_out, dict);
				}
				
			}
			
			for (int i = 0; i < m_peakList.size(); i++)
			{
				Peak &peak = m_peakList[i];
	
				bool found = false;
				for (int j=0; j < temp.size(); j++)
				{
					if (temp[j].isInRange(peak))
					{
						temp[j].update(peak);
						temp[j].setState(1);
						found = true;
						break;
					}
				}
				
				if (!found)
				{
					peak.setState(1);
					temp.push_back(peak);
				}
			}
			
			for (int j=0; j < temp.size(); j++)
			{
				Peak &peak = temp[j];
				if (peak.isModified)
				{
					pmt::pmt_t dict = peak.msg();
					message_port_pub(m_msg_port_out, dict);
				}
				
			}
			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
			// -----------------------------------------------------------------
			if (optr_data)
			{
				for (int i = 0; i < nitems_per_block; i++)
				{
					*(optr_data++) = m_pXs[i];
				}
			}
			if (optr_pbox)
			{
				for (int i = 0; i < nitems_per_block; i++)
				{
					*(optr_pbox++) = m_pPeakBox_abs[i];
				}
			}

			m_frameCount++;
			consume(0, ninput_items[DATA]);
		}

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

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