#include "MultibandRasterRenderer.h"
#include "RasterDataProvider.h"
#include <QDebug>
#include <algorithm>
#include <cmath>

MultibandRasterRenderer::MultibandRasterRenderer(RasterDataProvider* provider)
    : RasterRendererBase(provider)
{
}

void MultibandRasterRenderer::stretchTo8bit(const std::vector<float>& in, std::vector<unsigned char>& out, float minVal, float maxVal)
{
    size_t n = in.size();
    out.resize(n);
    if (maxVal <= minVal)
    {
        std::fill(out.begin(), out.end(), 0);
        return;
    }
    float denom = 1.0f / (maxVal - minVal);
    for (size_t i = 0; i < n; ++i)
    {
        float v = (in[i] - minVal) * denom;
        v = std::min(std::max(v, 0.0f), 1.0f);
        out[i] = static_cast<unsigned char>(v * 255.0f);
    }
}

void MultibandRasterRenderer::computeHistogram(const std::vector<float>& in, std::vector<long>& hist, float& minVal, float& maxVal)
{
    if (in.empty())
    {
        minVal = 0.0f;
        maxVal = 1.0f;
        hist.clear();
        return;
    }

    minVal = in[0];
    maxVal = in[0];
    for (size_t i = 0; i < in.size(); ++i)
    {
        if (in[i] < minVal) minVal = in[i];
        if (in[i] > maxVal) maxVal = in[i];
    }

    if (maxVal <= minVal)
    {
        hist.assign(256, 0);
        hist[0] = static_cast<long>(in.size());
        minVal = 0.0f;
        maxVal = 1.0f;
        return;
    }

    const int numBins = 65536;
    hist.assign(numBins, 0);

    for (size_t i = 0; i < in.size(); ++i)
    {
        ++hist[static_cast<int>(in[i])];
    }
}

void MultibandRasterRenderer::findStretchValuesByRatio(const std::vector<long>& hist, float ratio, float& minVal, float& maxVal)
{
    if (hist.empty()) return;

    long totalPixels = 0;
    for (size_t i = 0; i < hist.size(); ++i)
    {
        totalPixels += hist[i];
    }
    if (totalPixels == 0) return;

    long targetCount = static_cast<long>(totalPixels * ratio);

    long cumulative = 0;
    for (size_t i = 0; i < hist.size(); ++i)
    {
        cumulative += hist[i];
        if (cumulative >= targetCount)
        {
            minVal = static_cast<float>(i);
            break;
        }
    }

    cumulative = 0;
    for (size_t i = hist.size(); i > 0; --i)
    {
        cumulative += hist[i - 1];
        if (cumulative >= targetCount)
        {
            maxVal = static_cast<float>(i - 1);
            break;
        }
    }
}

void MultibandRasterRenderer::stretchTo8bitByHistogram(const std::vector<float>& in, std::vector<unsigned char>& out, float minVal, float maxVal)
{
    size_t n = in.size();
    out.resize(n);
    if (maxVal <= minVal)
    {
        std::fill(out.begin(), out.end(), 0);
        return;
    }
    float denom = 1.0f / (maxVal - minVal);
    for (size_t i = 0; i < n; ++i)
    {
        float v = (in[i] - minVal) * denom;
        v = std::min(std::max(v, 0.0f), 1.0f);
        out[i] = static_cast<unsigned char>(v * 255.0f);
    }
}

QImage MultibandRasterRenderer::render()
{
    if (!m_provider) return QImage();
    int bands = m_provider->bandCount();
    int w = m_provider->width();
    int h = m_provider->height();
    if (w <= 0 || h <= 0) return QImage();

    // Find nearest bands for requested wavelengths if wavelengths available
    std::vector<double> wavelengths = m_provider->bandWavelengths();

    auto chooseBandIndexForWave = [&](double wave)->int {
        if (wavelengths.empty()) {
            // fallback: select R,G,B as first three bands
            if (bands >= 3) return (wave==m_params.rWave?0:(wave==m_params.gWave?1:2));
            if (bands >= 1) return 0;
            return -1;
        }
        int best = -1; double bestDiff = 1e12;
        for (int i = 0; i < (int)wavelengths.size(); ++i) {
            if (wavelengths[i] < 0) continue;
            double d = std::abs(wavelengths[i] - wave);
            if (d < bestDiff) { bestDiff = d; best = i; }
        }
        if (best >= 0) return best;
        // fallback
        return std::min(2, bands-1);
    };

    int rIdx = chooseBandIndexForWave(m_params.rWave);
    int gIdx = chooseBandIndexForWave(m_params.gWave);
    int bIdx = chooseBandIndexForWave(m_params.bWave);

    std::vector<float> rbuf, gbuf, bbuf;
    if (rIdx >= 0) m_provider->readBandAsFloat(rIdx, rbuf);
    if (gIdx >= 0) m_provider->readBandAsFloat(gIdx, gbuf);
    if (bIdx >= 0) m_provider->readBandAsFloat(bIdx, bbuf);

    std::vector<unsigned char> r8, g8, b8;

    // Compute histogram-based stretch for each channel
    float ratio = m_stretchRatio;

    // R channel stretch
    if (!rbuf.empty()) {
        std::vector<long> rHist;
        float rMin, rMax;
        computeHistogram(rbuf, rHist, rMin, rMax);
        float rStretchMin = rMin, rStretchMax = rMax;
        findStretchValuesByRatio(rHist, ratio, rStretchMin, rStretchMax);
        stretchTo8bitByHistogram(rbuf, r8, rStretchMin, rStretchMax);
    }

    // G channel stretch
    if (!gbuf.empty()) {
        std::vector<long> gHist;
        float gMin, gMax;
        computeHistogram(gbuf, gHist, gMin, gMax);
        float gStretchMin = gMin, gStretchMax = gMax;
        findStretchValuesByRatio(gHist, ratio, gStretchMin, gStretchMax);
        stretchTo8bitByHistogram(gbuf, g8, gStretchMin, gStretchMax);
    }

    // B channel stretch
    if (!bbuf.empty()) {
        std::vector<long> bHist;
        float bMin, bMax;
        computeHistogram(bbuf, bHist, bMin, bMax);
        float bStretchMin = bMin, bStretchMax = bMax;
        findStretchValuesByRatio(bHist, ratio, bStretchMin, bStretchMax);
        stretchTo8bitByHistogram(bbuf, b8, bStretchMin, bStretchMax);
    }

    QImage out(w, h, QImage::Format_RGB888);
    for (int y = 0; y < h; ++y) {
        unsigned char* scan = out.scanLine(y);
        for (int x = 0; x < w; ++x) {
            int idx = y * w + x;
            unsigned char rc = (r8.size() > (size_t)idx) ? r8[idx] : 0;
            unsigned char gc = (g8.size() > (size_t)idx) ? g8[idx] : 0;
            unsigned char bc = (b8.size() > (size_t)idx) ? b8[idx] : 0;
            scan[x*3 + 0] = rc;
            scan[x*3 + 1] = gc;
            scan[x*3 + 2] = bc;
        }
    }
    return out;
}