#include "effects/pixelsort.hpp"

#include <algorithm>
#include <array>
#include <cstring>
#include <vector>

#include "core/parallel.hpp"
#include "effects/sampling.hpp"
#include "effects/source_fill.hpp"

namespace effector {

const std::vector<ParamSpec>& PixelSortEffect::params() const {
    static const std::vector<ParamSpec> specs = {
        ParamSpec::makeDouble("sort_low", "Lower threshold", 0.25, 0.0, 1.0,
                              "Pixels darker than this break a run and stay where they are."),
        ParamSpec::makeDouble("sort_high", "Upper threshold", 0.8, 0.0, 1.0,
                              "Pixels brighter than this break a run and stay where they are."),
        ParamSpec::makeChoice("sort_key", "Sort by", "luma", {"luma", "hue", "saturation"},
                              "What orders the pixels inside a run."),
        ParamSpec::makeChoice("sort_direction", "Direction", "horizontal",
                              {"horizontal", "vertical"}, "Whether rows or columns are sorted."),
        ParamSpec::makeBool("sort_reverse", "Reverse", false,
                            "Sort from high to low instead of from low to high."),
        ParamSpec::makeDouble("sort_line_chance", "Line chance", 1.0, 0.0, 1.0,
                              "Chance of each row (or column) being sorted at all."),
        ParamSpec::makeInt("sort_max_run", "Longest run", 0, 0, 8192,
                           "Each run is cut after a random length up to this many pixels, which "
                           "frays the edges of the sorted streaks (0 = no limit)."),
    };
    return specs;
}

namespace {

int luma(const std::uint8_t* px) { return (px[0] * 54 + px[1] * 183 + px[2] * 19) >> 8; }

// Hue on a 0..1535 scale (six sectors of 256), in integers.
int hue(const std::uint8_t* px) {
    const int r = px[0], g = px[1], b = px[2];
    const int hi = std::max({r, g, b});
    const int lo = std::min({r, g, b});
    const int range = hi - lo;
    if (range == 0) return 0;
    if (hi == r) return g >= b ? (g - b) * 255 / range : 1535 - (b - g) * 255 / range;
    if (hi == g) return 512 + (b - r) * 255 / range;
    return 1024 + (r - g) * 255 / range;
}

int saturation(const std::uint8_t* px) {
    return std::max({px[0], px[1], px[2]}) - std::min({px[0], px[1], px[2]});
}

// The image with rows and columns swapped, so vertical sorting is horizontal
// sorting of the transpose.
Image transposed(const Image& src) {
    Image out(src.height(), src.width());
    parallelFor(static_cast<std::size_t>(src.width()), [&](std::size_t begin, std::size_t end) {
        for (std::size_t x = begin; x < end; ++x) {
            std::uint8_t* row = out.row(static_cast<int>(x));
            for (int y = 0; y < src.height(); ++y) {
                std::memcpy(row + static_cast<std::size_t>(y) * 3, src.row(y) + x * 3, 3);
            }
        }
    });
    return out;
}

struct LinePlan {
    bool sorted = false;
    std::uint64_t seed = 0;  // for the run lengths, when they are limited
};

}  // namespace

Image PixelSortEffect::render(const std::vector<Image>& sources, int width, int height, Rng& rng,
                              const ParamValues& values) const {
    if (width < 1 || height < 1) throw EffectError("Canvas size must be at least 1x1 pixels.");

    const int low = static_cast<int>(std::clamp(paramDouble(values, "sort_low", 0.25), 0.0, 1.0) * 255.0 + 0.5);
    const int high = static_cast<int>(std::clamp(paramDouble(values, "sort_high", 0.8), 0.0, 1.0) * 255.0 + 0.5);
    const std::string keyName = paramChoice(values, "sort_key", "luma");
    const bool vertical = paramChoice(values, "sort_direction", "horizontal") == "vertical";
    const bool reverse = paramBool(values, "sort_reverse", false);
    const double lineChance = std::clamp(paramDouble(values, "sort_line_chance", 1.0), 0.0, 1.0);
    const int maxRun = static_cast<int>(std::max(0L, paramInt(values, "sort_max_run", 0)));

    Image canvas = canvasFromSources(sources, width, height, rng);
    if (vertical) canvas = transposed(canvas);
    const int lines = canvas.height();
    const int span = canvas.width();

    // Which lines are sorted, and a seed for each line's run lengths, are drawn
    // serially; the lines then sort in parallel, each from its own generator.
    std::vector<LinePlan> plan(static_cast<std::size_t>(lines));
    for (auto& line : plan) {
        line.sorted = chance(rng.next(), lineChance);
        if (maxRun > 0) line.seed = rng.next();
    }

    auto keyOf = [&](const std::uint8_t* px) {
        if (keyName == "hue") return hue(px);
        if (keyName == "saturation") return saturation(px);
        return luma(px);
    };

    parallelFor(static_cast<std::size_t>(lines), [&](std::size_t begin, std::size_t end) {
        std::vector<std::pair<int, std::array<std::uint8_t, 3>>> run;
        for (std::size_t line = begin; line < end; ++line) {
            if (!plan[line].sorted) continue;
            Rng runRng(plan[line].seed);
            std::uint8_t* row = canvas.row(static_cast<int>(line));

            auto flush = [&](int start) {
                // A stable sort, so equal keys keep their order on every
                // standard library.
                std::stable_sort(run.begin(), run.end(), [&](const auto& a, const auto& b) {
                    return reverse ? a.first > b.first : a.first < b.first;
                });
                for (std::size_t i = 0; i < run.size(); ++i) {
                    std::memcpy(row + (static_cast<std::size_t>(start) + i) * 3, run[i].second.data(), 3);
                }
                run.clear();
            };

            int start = 0;
            int limit = maxRun > 0 ? static_cast<int>(runRng.inRange(1, maxRun)) : span;
            for (int x = 0; x < span; ++x) {
                const std::uint8_t* px = row + static_cast<std::size_t>(x) * 3;
                const int brightness = luma(px);
                if (brightness < low || brightness > high) {
                    flush(start);
                    start = x + 1;
                    continue;
                }
                if (static_cast<int>(run.size()) >= limit) {
                    flush(start);
                    start = x;
                    if (maxRun > 0) limit = static_cast<int>(runRng.inRange(1, maxRun));
                }
                run.push_back({keyOf(px), {px[0], px[1], px[2]}});
            }
            flush(start);
        }
    });

    return vertical ? transposed(canvas) : canvas;
}

}  // namespace effector