#ifndef RAWIMAGE_H
#define RAWIMAGE_H

#include <vector>
#include <algorithm>
#include <stdint.h>
#include <math.h>
#include <memory.h>

class Peak
{
    uint32_t m_v;
    uint32_t m_x,m_y;
public:
    Peak() : m_v(0), m_x(0), m_y(0) {}
    Peak(uint32_t v, uint32_t x, uint32_t y) : m_v(v), m_x(x), m_y(y) {}
    uint32_t v() const { return m_v; }
    uint32_t x() const { return m_x; }
    uint32_t y() const { return m_y; }
    double distance(const Peak &d) const
    {
        uint32_t dx = m_x-d.m_x;
        uint32_t dy = m_y-d.m_y;
        return dx*dx + dy*dy;
    }
    bool operator <(const Peak &d) const
    {
        return m_v > d.m_v;
    }
};

class RawImageAbs
{
protected:
    uint32_t m_width,m_height;
public:
    virtual ~RawImageAbs(){}
    virtual bool imageStats(uint64_t *mean, double *stdDev, uint64_t *median, uint64_t *min, uint64_t *max) const = 0;
    virtual void rect(int &x, int &y, int w, int h, std::vector<double> &r) const = 0;
    virtual int findPeaks(uint64_t background, double distance, std::vector<Peak> &peaks) const = 0;
    virtual RawImageAbs* medianFilter() const = 0;
    virtual void quarter() = 0;
    uint32_t width() const
    {
        return m_width;
    }
    uint32_t height() const
    {
        return m_height;
    }
};

template<typename T>
class RawImage : public RawImageAbs
{
    std::vector<T> m_img;
    bool checkPixel(T c, uint32_t x, uint32_t y) const
    {
        T d = pixel(x, y);
        return c>=d;
    }
public:
    RawImage()
    {
        m_width = m_height = 0;
    }
    RawImage(int w, int h)
    {
        m_width = w;
        m_height = h;
        m_img.resize(w*h);
    }
    RawImage(int w, int h, std::vector<T> &img)
    {
        m_width = w;
        m_height = h;
        img.resize(w*h);
        m_img = std::move(img);
        img.clear();
    }
    T* data()
    {
        return m_img.data();
    }
    std::vector<T> dataArray() const
    {
        return m_img;
    }
    bool imageStats(uint64_t *mean, double *stdDev, uint64_t *median, uint64_t *min, uint64_t *max) const
    {
        if(m_img.size()==0)return false;

        uint64_t sum = 0;
        uint64_t sqrSum = 0;
        uint64_t tMin = UINT64_MAX;
        uint64_t tMax = 0;
        uint32_t histogram[65536];
        memset(histogram, 0, sizeof(histogram));
        const int shift = sizeof(T)>2 ? sizeof(T)*8 - 16 : 0;

        for(T i : m_img)
        {
            sum += i;
            sqrSum += i*i;
            tMin = tMin>i ? i : tMin;
            tMax = tMax<i ? i : tMax;
            histogram[i>>shift]++;
        }

        if(mean)*mean = sum/m_img.size();
        if(min)*min = tMin;
        if(max)*max = tMax;
        if(stdDev)*stdDev = sqrt((sqrSum - (double)sum*sum/m_img.size()) / (double)(m_img.size()-1));


        if(median)
        {
            size_t medianSum = 0;
            for(int i=0;i<65536;i++)
            {
                medianSum += histogram[i];
                if(medianSum>=m_img.size()/2)
                {
                    *median = i;
                    break;
                }
            }
        }

        return true;
    }
    T pixel(uint32_t x, uint32_t y) const
    {
        if(x>=m_width)x=0;
        if(y>=m_height)y=0;
        return m_img[y*m_width+x];
    }
    void rect(int &x, int &y, int w, int h, std::vector<double> &r) const
    {
        r.resize(w*h);
        x -= w/2;
        y -= h/2;
        if(x<0)x = 0;
        if(y<0)y = 0;
        if(x+w >= m_width)x = m_width-w;
        if(y+h >= m_height)y = m_height-h;
        uint32_t d = 0;
        for(int i=y;i<y+h;i++)
            for(int o=x;o<x+w;o++)
                r[d++] = pixel(o, i);
    }
    int findPeaks(uint64_t background, double distance, std::vector<Peak> &peaks) const
    {
        std::vector<Peak> tmpPeaks;
        const int r = 1;
        for(uint32_t i=r; i<m_height-r; i++)
        {
            for(uint32_t o=r; o<m_width-r; o++)
            {
                T c = pixel(o, i);
                if(c>background && checkPixel(c, o-r, i-r) && checkPixel(c, o, i-r) &&
                        checkPixel(c, o+r, i-r) && checkPixel(c, o-r, i) && checkPixel(c, o+r, i)
                        && checkPixel(c, o-r, i+r) && checkPixel(c, o, i+r) && checkPixel(c, o+r, i+r))
                {
                    tmpPeaks.push_back(Peak(c, o, i));
                }
            }
        }

        int num = tmpPeaks.size();
        std::sort(tmpPeaks.begin(), tmpPeaks.end());
        tmpPeaks.resize(std::min(num, 10000));
        uint32_t d = distance*distance;
        for(const Peak &p : tmpPeaks)
        {
            bool pass = true;
            for(const Peak &k : peaks)
            {
                if(p.distance(k) < d)
                {
                    pass = false;
                    break;
                }
            }
            if(pass)
            {
                peaks.push_back(p);
            }
        }
        return num;
    }
    RawImageAbs* medianFilter() const
    {
        RawImage<T> *ret = new RawImage<T>;
        std::vector<T> tmp;
        tmp.resize(m_width*m_height);
        #pragma omp parallel for
        for(uint32_t y=0;y<m_height;y++)
        {
            for(uint32_t x=0;x<m_width;x++)
            {
                T array[9] = { pixel(x-1, y-1), pixel(x, y-1), pixel(x+1, y-1),
                             pixel(x-1, y), pixel(x, y), pixel(x+1, y),
                             pixel(x-1, y+1), pixel(x, y+1), pixel(x+1, y+1)};
                std::nth_element(std::begin(array), array+4, std::end(array));
                tmp[y*m_width+x] = array[4];
            }
        }
        ret->m_width = m_width;
        ret->m_height = m_height;
        ret->m_img = std::move(tmp);
        return ret;
    }
    void quarter()
    {
        std::vector<T> tmp;
        tmp.resize(m_width*m_height);
        uint32_t d = 0;
        for(uint32_t i=0;i<m_height;i+=2)
        {
            for(uint32_t o=0;o<m_width;o+=2)
            {
                tmp[d++] = m_img[i*m_width+o];
            }
        }
        m_img = tmp;
        m_width /= 2;
        m_height /= 2;
    }
};

#endif // RAWIMAGE_H