#include "effects/shatter.hpp"

#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

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

namespace effector {

const std::vector<ParamSpec>& ShatterEffect::params() const {
    static const std::vector<ParamSpec> specs = {
        ParamSpec::makeInt("shatter_cells", "Shards", 40, 2, 5000,
                           "How many shards the image breaks into."),
        ParamSpec::makeInt("shatter_displace", "Displacement", 40, 0, 4096,
                           "Furthest in pixels a shard's picture may slide, along either axis."),
        ParamSpec::makeDouble("shatter_rotate", "Rotation", 8.0, 0.0, 180.0,
                              "Largest turn in degrees of a shard's picture, either way.", 0.5, 1),
        ParamSpec::makeDouble("shatter_zoom", "Zoom", 0.05, 0.0, 0.9,
                              "Largest zoom in or out of a shard's picture, as a fraction.", 0.01),
        ParamSpec::makeInt("shatter_crack_width", "Crack width", 2, 0, 64,
                           "Width in pixels of the cracks between the shards (0 = none)."),
        ParamSpec::makeColor("shatter_crack_colour", "Crack colour", "#101010",
                             "Colour the cracks are drawn in."),
    };
    return specs;
}

namespace {

struct Shard {
    double x = 0.0;  // the point the shard grows around
    double y = 0.0;
    double dx = 0.0;
    double dy = 0.0;
    double sine = 0.0;
    double cosine = 1.0;
    double zoom = 1.0;
};

// The shard points bucketed on a coarse grid, so finding the two nearest to a
// pixel looks at a few buckets instead of every point.
class ShardGrid {
public:
    ShardGrid(const std::vector<Shard>& shards, int width, int height) : shards_(shards) {
        cell_ = std::max(1.0, std::sqrt(static_cast<double>(width) * height /
                                        static_cast<double>(shards.size())));
        cols_ = static_cast<int>(width / cell_) + 1;
        rows_ = static_cast<int>(height / cell_) + 1;
        buckets_.resize(static_cast<std::size_t>(cols_) * rows_);
        for (std::size_t i = 0; i < shards.size(); ++i) {
            buckets_[bucketOf(shards[i].x, shards[i].y)].push_back(static_cast<int>(i));
        }
    }

    // The nearest and second nearest shard to (x, y).
    void nearest(double x, double y, int& first, int& second) const {
        const int bx = std::clamp(static_cast<int>(x / cell_), 0, cols_ - 1);
        const int by = std::clamp(static_cast<int>(y / cell_), 0, rows_ - 1);
        double best = std::numeric_limits<double>::max();
        double next = best;
        first = second = -1;
        const int maxRing = std::max(cols_, rows_);
        for (int ring = 0; ring <= maxRing; ++ring) {
            // Anything in this ring is at least (ring - 1) cells away, so once
            // both answers are closer than that the search is over.
            const double reach = std::max(0, ring - 1) * cell_;
            if (second >= 0 && next < reach * reach) break;
            for (int gy = by - ring; gy <= by + ring; ++gy) {
                if (gy < 0 || gy >= rows_) continue;
                const bool edgeRow = gy == by - ring || gy == by + ring;
                for (int gx = bx - ring; gx <= bx + ring; gx += edgeRow ? 1 : 2 * std::max(ring, 1)) {
                    if (gx < 0 || gx >= cols_) continue;
                    for (int index : buckets_[static_cast<std::size_t>(gy) * cols_ + gx]) {
                        const Shard& shard = shards_[static_cast<std::size_t>(index)];
                        const double d = (shard.x - x) * (shard.x - x) + (shard.y - y) * (shard.y - y);
                        if (d < best || (d == best && index < first)) {
                            next = best;
                            second = first;
                            best = d;
                            first = index;
                        } else if (d < next || (d == next && index < second)) {
                            next = d;
                            second = index;
                        }
                    }
                    if (ring == 0) break;
                }
            }
        }
    }

private:
    std::size_t bucketOf(double x, double y) const {
        const int bx = std::clamp(static_cast<int>(x / cell_), 0, cols_ - 1);
        const int by = std::clamp(static_cast<int>(y / cell_), 0, rows_ - 1);
        return static_cast<std::size_t>(by) * cols_ + bx;
    }

    const std::vector<Shard>& shards_;
    double cell_ = 1.0;
    int cols_ = 1;
    int rows_ = 1;
    std::vector<std::vector<int>> buckets_;
};

}  // namespace

Image ShatterEffect::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 long count = std::max(2L, paramInt(values, "shatter_cells", 40));
    const int displace = static_cast<int>(std::max(0L, paramInt(values, "shatter_displace", 40)));
    const double rotate = std::clamp(paramDouble(values, "shatter_rotate", 8.0), 0.0, 180.0);
    const double zoom = std::clamp(paramDouble(values, "shatter_zoom", 0.05), 0.0, 0.9);
    const double crackWidth = static_cast<double>(std::max(0L, paramInt(values, "shatter_crack_width", 2)));
    std::uint8_t crack[3];
    paramColor(values, "shatter_crack_colour", "#101010", crack);

    const Image input = canvasFromSources(sources, width, height, rng);

    std::vector<Shard> shards(static_cast<std::size_t>(count));
    for (auto& shard : shards) {
        shard.x = unitDouble(rng.next()) * width;
        shard.y = unitDouble(rng.next()) * height;
        shard.dx = static_cast<double>(rng.inRange(-displace, displace));
        shard.dy = static_cast<double>(rng.inRange(-displace, displace));
        sinCosDegrees((unitDouble(rng.next()) * 2.0 - 1.0) * rotate, shard.sine, shard.cosine);
        shard.zoom = 1.0 + (unitDouble(rng.next()) * 2.0 - 1.0) * zoom;
    }
    const ShardGrid grid(shards, width, height);

    Image canvas(width, height);
    parallelFor(static_cast<std::size_t>(height), [&](std::size_t begin, std::size_t end) {
        for (std::size_t row = begin; row < end; ++row) {
            std::uint8_t* out = canvas.row(static_cast<int>(row));
            const double y = static_cast<double>(row) + 0.5;
            for (int col = 0; col < width; ++col) {
                std::uint8_t* px = out + static_cast<std::size_t>(col) * 3;
                const double x = col + 0.5;
                int first = 0;
                int second = 0;
                grid.nearest(x, y, first, second);
                const Shard& a = shards[static_cast<std::size_t>(first)];

                if (crackWidth > 0.0 && second >= 0) {
                    // Distance to the edge between the two nearest shards,
                    // which is the perpendicular bisector of their points.
                    const Shard& b = shards[static_cast<std::size_t>(second)];
                    const double da = (x - a.x) * (x - a.x) + (y - a.y) * (y - a.y);
                    const double db = (x - b.x) * (x - b.x) + (y - b.y) * (y - b.y);
                    const double gap = std::sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
                    if (gap > 0.0 && (db - da) / (2.0 * gap) < crackWidth * 0.5) {
                        px[0] = crack[0];
                        px[1] = crack[1];
                        px[2] = crack[2];
                        continue;
                    }
                }

                // Undo the shard's slide, then its turn and zoom about its point.
                const double rx = x - a.dx - a.x;
                const double ry = y - a.dy - a.y;
                const double sx = a.x + (rx * a.cosine + ry * a.sine) / a.zoom;
                const double sy = a.y + (-rx * a.sine + ry * a.cosine) / a.zoom;
                sampleBilinear(input, sx - 0.5, sy - 0.5, px);
            }
        }
    });

    return canvas;
}

}  // namespace effector