#include "PointCloudFilter.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <unordered_map>

namespace
{
struct GridKey
{
    long long x;
    long long y;
    long long z;

    bool operator==(const GridKey& other) const
    {
        return x == other.x && y == other.y && z == other.z;
    }
};

struct GridKeyHash
{
    std::size_t operator()(const GridKey& key) const
    {
        const std::size_t h1 = std::hash<long long>()(key.x);
        const std::size_t h2 = std::hash<long long>()(key.y);
        const std::size_t h3 = std::hash<long long>()(key.z);
        return h1 ^ (h2 + static_cast<std::size_t>(0x9e3779b9) + (h1 << 6) + (h1 >> 2))
            ^ (h3 + static_cast<std::size_t>(0x9e3779b9) + (h2 << 6) + (h2 >> 2));
    }
};

struct FlatPoint
{
    SVzNL3DPosition point;
};

bool IsFinite(const SVzNL3DPosition& point)
{
    return std::isfinite(point.pt3D.x) && std::isfinite(point.pt3D.y)
        && std::isfinite(point.pt3D.z);
}

double SquaredDistance(const SVzNL3DPosition& first, const SVzNL3DPosition& second)
{
    const double dx = first.pt3D.x - second.pt3D.x;
    const double dy = first.pt3D.y - second.pt3D.y;
    const double dz = first.pt3D.z - second.pt3D.z;
    return dx * dx + dy * dy + dz * dz;
}

long long ToCell(double value, double cellSize)
{
    const double coordinate = std::floor(value / cellSize);
    if (coordinate <= static_cast<double>(std::numeric_limits<long long>::min())) {
        return std::numeric_limits<long long>::min();
    }
    if (coordinate >= static_cast<double>(std::numeric_limits<long long>::max())) {
        return std::numeric_limits<long long>::max();
    }
    return static_cast<long long>(coordinate);
}

GridKey GetKey(const SVzNL3DPosition& point, double cellSize)
{
    GridKey key;
    key.x = ToCell(point.pt3D.x, cellSize);
    key.y = ToCell(point.pt3D.y, cellSize);
    key.z = ToCell(point.pt3D.z, cellSize);
    return key;
}

void ResetStats(PointCloudFilterStats* stats, std::size_t inputCount)
{
    if (stats != nullptr) {
        stats->inputPointCount = inputCount;
        stats->outputPointCount = 0;
        stats->removedPointCount = 0;
    }
}

void FinishStats(PointCloudFilterStats* stats, std::size_t outputCount)
{
    if (stats != nullptr) {
        stats->outputPointCount = outputCount;
        stats->removedPointCount = stats->inputPointCount >= outputCount
            ? stats->inputPointCount - outputCount : 0;
    }
}

bool FilterScanLine(const std::vector<SVzNL3DPosition>& input,
                    std::vector<SVzNL3DPosition>& output,
                    const PointCloudFilterParams& params)
{
    output.clear();
    if (input.empty()) {
        return true;
    }

    const std::size_t minLength = std::max<std::size_t>(1, params.minSegmentLength);
    std::size_t segmentBegin = 0;
    for (std::size_t index = 1; index <= input.size(); ++index) {
        bool connected = index < input.size() && IsFinite(input[index - 1])
            && IsFinite(input[index])
            && std::abs(input[index].pt3D.z - input[index - 1].pt3D.z)
                <= std::max(0.0, params.continuityThreshold);
        if (index == input.size() || !connected) {
            if (index - segmentBegin >= minLength) {
                output.insert(output.end(), input.begin() + static_cast<std::ptrdiff_t>(segmentBegin),
                              input.begin() + static_cast<std::ptrdiff_t>(index));
            }
            segmentBegin = index;
        }
    }
    return true;
}

double EstimateCellSize(const std::vector<FlatPoint>& points)
{
    if (points.size() < 2) {
        return 1.0;
    }
    double minX = points[0].point.pt3D.x;
    double maxX = minX;
    double minY = points[0].point.pt3D.y;
    double maxY = minY;
    double minZ = points[0].point.pt3D.z;
    double maxZ = minZ;
    for (std::size_t index = 1; index < points.size(); ++index) {
        const SVzNL3DPoint& value = points[index].point.pt3D;
        minX = std::min(minX, value.x);
        maxX = std::max(maxX, value.x);
        minY = std::min(minY, value.y);
        maxY = std::max(maxY, value.y);
        minZ = std::min(minZ, value.z);
        maxZ = std::max(maxZ, value.z);
    }
    const double diagonal = std::sqrt((maxX - minX) * (maxX - minX)
        + (maxY - minY) * (maxY - minY) + (maxZ - minZ) * (maxZ - minZ));
    const double scale = diagonal / std::cbrt(static_cast<double>(points.size()));
    return std::max(scale, 1e-6);
}

template <typename KeepPredicate>
void FilterWithGrid(const std::vector<FlatPoint>& points, double cellSize,
                    double radius, KeepPredicate keepPoint,
                    std::vector<bool>& keep)
{
    std::unordered_map<GridKey, std::vector<std::size_t>, GridKeyHash> grid;
    grid.reserve(points.size() * 2 + 1);
    for (std::size_t index = 0; index < points.size(); ++index) {
        grid[GetKey(points[index].point, cellSize)].push_back(index);
    }

    const long long cellRange = std::max<long long>(1, static_cast<long long>(std::ceil(radius / cellSize)));
    const double radiusSquared = radius * radius;
    std::vector<std::size_t> candidates;
    candidates.reserve(32);
    for (std::size_t index = 0; index < points.size(); ++index) {
        const GridKey center = GetKey(points[index].point, cellSize);
        candidates.clear();
        for (long long x = -cellRange; x <= cellRange; ++x) {
            for (long long y = -cellRange; y <= cellRange; ++y) {
                for (long long z = -cellRange; z <= cellRange; ++z) {
                    GridKey key = {center.x + x, center.y + y, center.z + z};
                    const auto found = grid.find(key);
                    if (found != grid.end()) {
                        candidates.insert(candidates.end(), found->second.begin(), found->second.end());
                    }
                }
            }
        }
        keep[index] = keepPoint(index, candidates, radiusSquared);
    }
}
}

bool PointCloudFilter::Filter(const std::vector<std::vector<SVzNL3DPosition>>& input,
                              std::vector<std::vector<SVzNL3DPosition>>& output,
                              const PointCloudFilterParams& params,
                              PointCloudFilterStats* stats)
{
    std::size_t inputCount = 0;
    for (const auto& line : input) {
        inputCount += line.size();
    }
    ResetStats(stats, inputCount);
    output.clear();
    output.resize(input.size());

    if (params.method == PointCloudFilterMethod::None) {
        output = input;
        FinishStats(stats, inputCount);
        return true;
    }
    if (params.method == PointCloudFilterMethod::ScanLineContinuity) {
        std::size_t outputCount = 0;
        for (std::size_t lineIndex = 0; lineIndex < input.size(); ++lineIndex) {
            FilterScanLine(input[lineIndex], output[lineIndex], params);
            outputCount += output[lineIndex].size();
        }
        FinishStats(stats, outputCount);
        return true;
    }
    if (params.method != PointCloudFilterMethod::RadiusOutlier
        && params.method != PointCloudFilterMethod::StatisticalOutlier) {
        output.clear();
        return false;
    }
    if (params.method == PointCloudFilterMethod::RadiusOutlier && params.radius <= 0.0) {
        output.clear();
        return false;
    }

    std::vector<FlatPoint> points;
    points.reserve(inputCount);
    for (const auto& line : input) {
        for (const SVzNL3DPosition& point : line) {
            if (IsFinite(point)) {
                points.push_back({point});
            }
        }
    }
    std::vector<SVzNL3DPosition> filtered;
    std::vector<SVzNL3DPosition> flatInput;
    flatInput.reserve(points.size());
    for (std::size_t index = 0; index < points.size(); ++index) {
        flatInput.push_back(points[index].point);
        // Encode the flattened position only for reconstructing scan lines.
        flatInput.back().nPointIdx = static_cast<int>(index);
    }
    if (!Filter(flatInput, filtered, params, nullptr)) {
        output.clear();
        return false;
    }

    std::vector<bool> retained(points.size(), false);
    for (const SVzNL3DPosition& point : filtered) {
        if (point.nPointIdx >= 0
            && static_cast<std::size_t>(point.nPointIdx) < retained.size()) {
            retained[static_cast<std::size_t>(point.nPointIdx)] = true;
        }
    }
    std::size_t pointIndex = 0;
    for (std::size_t lineIndex = 0; lineIndex < input.size(); ++lineIndex) {
        for (const SVzNL3DPosition& point : input[lineIndex]) {
            if (IsFinite(point)) {
                if (pointIndex < retained.size() && retained[pointIndex]) {
                    output[lineIndex].push_back(point);
                }
                ++pointIndex;
            }
        }
    }
    std::size_t outputCount = 0;
    for (const auto& line : output) {
        outputCount += line.size();
    }
    FinishStats(stats, outputCount);
    return true;
}

bool PointCloudFilter::Filter(const std::vector<SVzNL3DPosition>& input,
                              std::vector<SVzNL3DPosition>& output,
                              const PointCloudFilterParams& params,
                              PointCloudFilterStats* stats)
{
    ResetStats(stats, input.size());
    output.clear();
    if (params.method == PointCloudFilterMethod::None) {
        output = input;
        FinishStats(stats, output.size());
        return true;
    }
    if (params.method == PointCloudFilterMethod::ScanLineContinuity) {
        FilterScanLine(input, output, params);
        FinishStats(stats, output.size());
        return true;
    }
    if (params.method != PointCloudFilterMethod::RadiusOutlier
        && params.method != PointCloudFilterMethod::StatisticalOutlier) {
        return false;
    }
    if (params.method == PointCloudFilterMethod::RadiusOutlier && params.radius <= 0.0) {
        return false;
    }

    std::vector<FlatPoint> points;
    points.reserve(input.size());
    for (const SVzNL3DPosition& point : input) {
        if (IsFinite(point)) {
            points.push_back({point});
        }
    }
    if (points.empty()) {
        FinishStats(stats, 0);
        return true;
    }

    const double requestedRadius = params.method == PointCloudFilterMethod::RadiusOutlier
        ? params.radius : params.searchRadius;
    const double radius = requestedRadius > 0.0 ? requestedRadius : EstimateCellSize(points) * 2.0;
    const double cellSize = std::max(radius, 1e-6);
    std::vector<bool> keep(points.size(), false);

    if (params.method == PointCloudFilterMethod::RadiusOutlier) {
        const std::size_t minNeighbors = params.minNeighbors;
        FilterWithGrid(points, cellSize, radius,
            [&points, minNeighbors](std::size_t index, const std::vector<std::size_t>& candidates,
                           double radiusSquared) {
                std::size_t neighbors = 0;
                for (std::size_t candidate : candidates) {
                    if (candidate != index && SquaredDistance(points[index].point, points[candidate].point) <= radiusSquared) {
                        ++neighbors;
                    }
                }
                return neighbors >= minNeighbors;
            }, keep);
    }

    if (params.method == PointCloudFilterMethod::StatisticalOutlier) {
        std::vector<double> meanDistances(points.size(), std::numeric_limits<double>::infinity());
        FilterWithGrid(points, cellSize, radius,
            [&points, &meanDistances, &params](std::size_t index,
                                               const std::vector<std::size_t>& candidates,
                                               double radiusSquared) {
                std::vector<double> distances;
                distances.reserve(candidates.size());
                for (std::size_t candidate : candidates) {
                    if (candidate != index) {
                        const double distance = SquaredDistance(points[index].point, points[candidate].point);
                        if (distance <= radiusSquared) {
                            distances.push_back(std::sqrt(distance));
                        }
                    }
                }
                if (distances.empty()) {
                    return false;
                }
                const std::size_t count = std::min(params.meanK == 0 ? std::size_t(1) : params.meanK,
                                                  distances.size());
                if (count < distances.size()) {
                    std::nth_element(distances.begin(),
                                     distances.begin() + static_cast<std::ptrdiff_t>(count),
                                     distances.end());
                }
                meanDistances[index] = std::accumulate(distances.begin(), distances.begin() + static_cast<std::ptrdiff_t>(count), 0.0)
                    / static_cast<double>(count);
                return true;
            }, keep);
        std::vector<double> validMeans;
        for (double value : meanDistances) {
            if (std::isfinite(value)) {
                validMeans.push_back(value);
            }
        }
        if (!validMeans.empty()) {
            const double mean = std::accumulate(validMeans.begin(), validMeans.end(), 0.0)
                / static_cast<double>(validMeans.size());
            double variance = 0.0;
            for (double value : validMeans) {
                const double delta = value - mean;
                variance += delta * delta;
            }
            const double deviation = std::sqrt(variance / static_cast<double>(validMeans.size()));
            const double threshold = mean + std::max(0.0, params.standardDeviationMultiplier) * deviation;
            for (std::size_t index = 0; index < keep.size(); ++index) {
                keep[index] = std::isfinite(meanDistances[index]) && meanDistances[index] <= threshold;
            }
        }
    }

    output.reserve(points.size());
    for (std::size_t index = 0; index < points.size(); ++index) {
        if (keep[index]) {
            output.push_back(points[index].point);
        }
    }
    FinishStats(stats, output.size());
    return true;
}