#include "rawimage.h"

int THUMB_SIZE = 128;
int THUMB_SIZE_BORDER = 138;
int THUMB_SIZE_BORDER_Y = 158;

RawImage::ImgType CV2Type(int cvtype)
{
    switch (cvtype)
    {
    case CV_8U:
        return RawImage::UINT8;
    case CV_16U:
        return RawImage::UINT16;
    case CV_32F:
        return RawImage::FLOAT32;
    case CV_8UC3:
        return RawImage::UINT8C3;
    case CV_8UC4:
        return RawImage::UINT8C4;
    case CV_16UC3:
        return RawImage::UINT16C3;
    case CV_16UC4:
        return RawImage::UINT16C4;
    case CV_32FC3:
        return RawImage::FLOAT32C3;
    default:
        return RawImage::UNKNOWN;
    }
}

int Type2CV(RawImage::ImgType type)
{
    switch (type)
    {
    case RawImage::UINT8:
        return CV_8U;
    case RawImage::UINT16:
        return CV_16U;
    case RawImage::FLOAT32:
        return CV_32F;
    case RawImage::UINT8C3:
        return CV_8UC3;
    case RawImage::UINT8C4:
        return CV_8UC4;
    case RawImage::UINT16C3:
        return CV_16UC3;
    case RawImage::UINT16C4:
        return CV_16UC4;
    case RawImage::FLOAT32C3:
        return CV_32FC3;
    case RawImage::UNKNOWN:
        return CV_8S;
    default:
        return CV_8U;
    }
}

RawImage::RawImage()
{
    m_stats = false;
}

RawImage::RawImage(int w, int h, ImgType type)
{
    m_img.create(h, w, Type2CV(type));
    m_stats = false;
}

RawImage::RawImage(cv::Mat &img)
{
    m_img = img;
    m_stats = false;
    scaleToUnit();
}

RawImage::RawImage(const RawImage &d)
{
    d.m_img.copyTo(m_img);
    m_mean = d.m_mean;
    m_stdDev = d.m_stdDev;
    m_median = d.m_median;
    m_min = d.m_min;
    m_max = d.m_max;
    m_mad = d.m_mad;
    m_stats = d.m_stats;
}

RawImage::RawImage(const QImage &img)
{
    if(img.format() == QImage::Format_RGB32)
    {
        m_img.create(img.height(), img.width(), CV_8UC4);
        for(int i=0; i<img.height(); i++)
            std::memcpy(m_img.ptr(i), img.scanLine(i), img.width()*4);
        cv::cvtColor(m_img, m_img, cv::COLOR_BGRA2RGB);
    }
    else if(img.format() == QImage::Format_ARGB32)
    {
        m_img.create(img.height(), img.width(), CV_8UC4);
        for(int i=0; i<img.height(); i++)
            std::memcpy(m_img.ptr(i), img.scanLine(i), img.width()*4);
        cv::cvtColor(m_img, m_img, cv::COLOR_BGRA2RGBA);
    }
    else if(img.format() == QImage::Format_RGBX64)
    {
        m_img.create(img.height(), img.width(), CV_16UC4);
        for(int i=0; i<img.height(); i++)
            std::memcpy(m_img.ptr(i), img.scanLine(i), img.width()*8);
        cv::cvtColor(m_img, m_img, cv::COLOR_RGBA2RGB);
    }
    else if(img.format() == QImage::Format_RGBA64)
    {
        m_img.create(img.height(), img.width(), CV_16UC4);
        for(int i=0; i<img.height(); i++)
            std::memcpy(m_img.ptr(i), img.scanLine(i), img.width()*8);
    }
    else
    {
        QImage tmp = img.convertToFormat(QImage::Format_RGB888);
        m_img.create(img.height(), img.width(), CV_8UC3);

        for(int i=0; i<tmp.height(); i++)
            std::memcpy(m_img.ptr(i), tmp.scanLine(i), tmp.width()*3);
    }
    m_stats = false;
}

bool RawImage::imageStats(double *mean, double *stdDev, double *median, double *min, double *max, double *mad)
{
    if(!m_stats)calcStats();
    if(mean)*mean = m_mean;
    if(stdDev)*stdDev = m_stdDev;
    if(median)*median = m_median;
    if(min)*min = m_min;
    if(max)*max = m_max;
    if(mad)*mad = m_mad;

    return true;
}

void RawImage::calcStats()
{
    if(m_stats)return;
    m_stats = true;

    cv::Scalar meanS, stdDevS;

    cv::meanStdDev(m_img, meanS, stdDevS);
    cv::minMaxIdx(m_img, &m_min, &m_max);

    cv::Mat img;
    if(m_img.channels() == 1)img = m_img;
    else if (m_img.channels() == 3)cv::cvtColor(m_img, img, cv::COLOR_BGR2GRAY);
    else if (m_img.channels() == 4)cv::cvtColor(m_img, img, cv::COLOR_BGRA2GRAY);

    int histSize = 256;
    if(img.type() == CV_16U || img.type() == CV_32F)histSize = 65536;
    float range[] = {0, (float)histSize};
    if(img.type() == CV_32F)range[1] = 1.0f;
    const float *ranges[] = {range};
    cv::Mat hist;
    cv::calcHist(&img, 1, nullptr, cv::Mat(), hist, 1, &histSize, ranges);

    m_mean = meanS[0];
    m_stdDev = stdDevS[0];
    size_t halfImageSize = size()/2;
    size_t medianSum = 0;
    for(int i=0; i < histSize; i++)
    {
        medianSum += hist.at<float>(0, i);
        if(medianSum >= halfImageSize)
        {
            m_median = i;
            break;
        }
    }
    if(img.type() == CV_32F)m_median /= histSize;
    cv::Mat absDev;
    img.convertTo(absDev, CV_32F, 1, -m_median);
    absDev = cv::abs(absDev);
    cv::Mat madHist;
    medianSum = 0;
    cv::calcHist(&absDev, 1, nullptr, cv::Mat(), madHist, 1, &histSize, ranges);
    for(int i=0; i < histSize; i++)
    {
        medianSum += madHist.at<float>(0, i);
        if(medianSum >= halfImageSize)
        {
            m_mad = i;
            break;
        }
    }
    if(img.type() == CV_32F)m_mad /= histSize;
}

void RawImage::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_img.cols)x = m_img.cols-w;
    if(y+h >= m_img.rows)y = m_img.rows-h;
    cv::Mat roiImg(m_img, cv::Rect(x, y, w, h));
    cv::Mat doubleMat;
    roiImg.convertTo(doubleMat, CV_64F);
    r = std::vector<double>(doubleMat.begin<double>(), doubleMat.end<double>());
}

int RawImage::findPeaks(double background, double distance, std::vector<Peak> &peaks) const
{
    std::vector<std::vector<cv::Point>> contours;

    cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(distance, distance));

    cv::Mat mask, dilate, locMax, result;
    cv::dilate(m_img, dilate, kernel);
    cv::compare(m_img, dilate, locMax, cv::CMP_GE);
    cv::compare(m_img, cv::Scalar(background), mask, cv::CMP_GT);
    cv::bitwise_and(locMax, mask, result);

    cv::findContours(result, contours, cv::noArray(), cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
    peaks.reserve(contours.size());
    for(auto contour : contours)
    {
        peaks.push_back(Peak(1, contour[0].x, contour[0].y));
    }

    return peaks.size();
}

RawImage* RawImage::medianFilter() const
{
    RawImage *ret = new RawImage();
    cv::medianBlur(m_img, ret->m_img, 3);
    return ret;
}

void RawImage::quarter()
{

}

uint32_t RawImage::width() const
{
    return m_img.cols;
}

uint32_t RawImage::height() const
{
    return m_img.rows;
}

uint32_t RawImage::size() const
{
    return width()*height();
}

RawImage::ImgType RawImage::type() const
{
    return CV2Type(m_img.type());
}

int RawImage::dataType() const
{
    return m_img.type();
}

uint32_t RawImage::norm() const
{
    switch(m_img.type())
    {
    case CV_8U:
    case CV_8UC3:
    case CV_8UC4:
        return UINT8_MAX;
    case CV_16U:
    case CV_16UC3:
        return UINT16_MAX;
    default:
        return 1;
    }
}

void* RawImage::data()
{
    return m_img.ptr();
}

const void *RawImage::data() const
{
    return m_img.ptr();
}

void RawImage::convertToThumbnail()
{
    m_thumbAspect = (float)width() / height();
    switch(CV_MAT_DEPTH(m_img.type()))
    {
    case CV_8U:
        m_img.convertTo(m_img, CV_16U, 255);
        break;
    case CV_32F:
        m_img.convertTo(m_img, CV_16U, 65535);
        break;
    case CV_16U:
        break;
    default:
        break;
    }

    if(m_img.channels() == 1)
        cv::cvtColor(m_img, m_img, cv::COLOR_GRAY2RGB);
    if(m_img.channels() == 4)
        cv::cvtColor(m_img, m_img, cv::COLOR_RGBA2RGB);
    cv::Size dsize(THUMB_SIZE, THUMB_SIZE);
    cv::resize(m_img, m_img, dsize, 0, 0, cv::INTER_NEAREST);
}

float RawImage::thumbAspect() const
{
    return m_thumbAspect;
}

const cv::Mat& RawImage::mat() const
{
    return m_img;
}

bool RawImage::pixel(int x, int y, QVector3D &rgb) const
{
    if(x < 0 || y < 0 || x >= (int)width() || y >= (int)height())return false;

    switch(m_img.type())
    {
    case CV_8U:
    {
        uint8_t v = m_img.at<uint8_t>(y, x);
        rgb = QVector3D(v, v, v);
        break;
    }
    case CV_16U:
    {
        uint16_t v = m_img.at<uint16_t>(y, x);
        rgb = QVector3D(v, v, v);
        break;
    }
    case CV_32F:
    {
        float v = m_img.at<float>(y, x);
        rgb = QVector3D(v, v, v);
        break;
    }
    case CV_8UC3:
    {
        cv::Vec3b v = m_img.at<cv::Vec3b>(y, x);
        rgb = QVector3D(v[0], v[1], v[2]);
        break;
    }
    case CV_8UC4:
    {
        cv::Vec4b v = m_img.at<cv::Vec4b>(y, x);
        rgb = QVector3D(v[0], v[1], v[2]);
        break;
    }
    case CV_16UC3:
    {
        cv::Vec3w v = m_img.at<cv::Vec3w>(y, x);
        rgb = QVector3D(v[0], v[1], v[2]);
        break;
    }
    case CV_32FC3:
    {
        cv::Vec3f v = m_img.at<cv::Vec3f>(y, x);
        rgb = QVector3D(v[0], v[1], v[2]);
        break;
    }
    default:
        rgb = QVector3D(0, 0, 0);
        break;
    }
    return true;
}

void RawImage::scaleToUnit()
{
    if(CV_MAT_DEPTH(m_img.type()) == CV_32F)
    {
        double min, max;
        cv::minMaxIdx(m_img, &min, &max);
        if(min < 0 || max > 1)
        {
            float scale = 1.0 / (max - min);
            float zero = min * scale;
            m_img = m_img * scale - zero;
        }
    }
}