#include "LocalizationService.h"

#include <vtkCellData.h>
#include <vtkExtractEdges.h>
#include <vtkGeometryFilter.h>
#include <vtkImageData.h>
#include <vtkMarchingCubes.h>
#include <vtkNew.h>
#include <vtkPointData.h>
#include <vtkPointDataToCellData.h>
#include <vtkPoints.h>
#include <vtkPolyData.h>
#include <vtkPolyDataConnectivityFilter.h>
#include <vtkPolyDataNormals.h>
#include <vtkSphereSource.h>
#include <vtkThreshold.h>
#include <vtkWindowedSincPolyDataFilter.h>
#include <vtkCleanPolyData.h>

#include <cmath>
#include <iostream>

namespace
{
struct Surf
{
    double length{0.0};
    double hausdorff{0.0};
    double bounds[6]{0, 0, 0, 0, 0, 0};
    double center[3]{0, 0, 0};
    double confronto_lati{0.0};
    int regionId{-1};
    Point3d centroid;
};

static Point3d centroidFromPoints(vtkPoints* pts)
{
    Point3d c;
    if (!pts) return c;
    const vtkIdType n = pts->GetNumberOfPoints();
    if (n <= 0) return c;

    double sx = 0.0, sy = 0.0, sz = 0.0;
    double p[3];
    for (vtkIdType i = 0; i < n; ++i)
    {
        // vtkPoints::GetPoint is non-const in VTK 9.x, so keep pts non-const here.
        pts->GetPoint(i, p);
        sx += p[0];
        sy += p[1];
        sz += p[2];
    }
    c.x = sx / static_cast<double>(n);
    c.y = sy / static_cast<double>(n);
    c.z = sz / static_cast<double>(n);
    return c;
}

static double sideDifferenceMetric(const double b[6])
{
    // Legacy "lati" metric: compare extents.
    const double dx = std::abs(b[1] - b[0]);
    const double dy = std::abs(b[3] - b[2]);
    const double dz = std::abs(b[5] - b[4]);

    // difference between the two largest extents (simple, stable)
    const double a = std::max({dx, dy, dz});
    const double c = std::min({dx, dy, dz});
    const double mid = dx + dy + dz - a - c;
    return std::max({std::abs(a - mid), std::abs(mid - c), std::abs(a - c)});
}

static double hausdorffToSphere(vtkPolyData* surface, const double center[3])
{
    if (!surface || surface->GetNumberOfPoints() == 0) return 0.0;

    vtkNew<vtkSphereSource> source_sphere;
    source_sphere->SetCenter(center[0], center[1], center[2]);
    source_sphere->SetRadius(5.0);
    source_sphere->SetThetaResolution(24);
    source_sphere->SetPhiResolution(24);
    source_sphere->Update();

    vtkPolyData* sphere = source_sphere->GetOutput();
    if (!sphere || sphere->GetNumberOfPoints() == 0) return 0.0;

    const vtkIdType nSource = surface->GetNumberOfPoints();
    const vtkIdType nTarget = sphere->GetNumberOfPoints();

    double maxMin = 0.0;
    double p1[3], p2[3];

    for (vtkIdType i = 0; i < nSource; ++i)
    {
        surface->GetPoint(i, p1);
        double smallest = std::numeric_limits<double>::infinity();

        for (vtkIdType j = 0; j < nTarget; ++j)
        {
            sphere->GetPoint(j, p2);
            const double dx = p1[0] - p2[0];
            const double dy = p1[1] - p2[1];
            const double dz = p1[2] - p2[2];
            const double dist2 = dx * dx + dy * dy + dz * dz;
            if (dist2 < smallest) smallest = dist2;
        }

        if (smallest > maxMin) maxMin = smallest;
    }

    return maxMin;
}

} // namespace

MarkerList LocalizationService::localize(
    vtkImageData* volume,
    const LocalizationParams& params,
    const ProgressFn& onProgress,
    const AbortFn& shouldAbort)
{
    MarkerList result;
    if (!volume)
        return result;

    // Marching cubes
    vtkNew<vtkMarchingCubes> marchingcubes;
    marchingcubes->SetInputData(volume);
    marchingcubes->SetValue(0, params.marchThreshold);
    marchingcubes->SetComputeNormals(0);
    marchingcubes->Update();

    if (marchingcubes->GetNumberOfContours() <= 0)
        return result;

    // Smoothing
    vtkNew<vtkWindowedSincPolyDataFilter> smoother;
    smoother->SetInputConnection(marchingcubes->GetOutputPort());
    smoother->SetNumberOfIterations(20);
    smoother->NormalizeCoordinatesOn();
    smoother->SetPassBand(0.01);
    smoother->BoundarySmoothingOn();
    smoother->Update();

    // Normals
    vtkNew<vtkPolyDataNormals> normals;
    normals->SetInputConnection(smoother->GetOutputPort());
    normals->Update();

    // 1) Label / count regions
    vtkNew<vtkPolyDataConnectivityFilter> connAll;
    connAll->SetInputConnection(normals->GetOutputPort());
    connAll->SetExtractionModeToAllRegions();
    connAll->ColorRegionsOn();
    connAll->Update();

    const int nRegions = connAll->GetNumberOfExtractedRegions();
    std::cout << "nRegions " << nRegions << std::endl;
    if (nRegions <= 0)
        return result;

    // 2) Extract each region
    vtkNew<vtkPolyDataConnectivityFilter> connOne;
    connOne->SetInputConnection(normals->GetOutputPort());
    connOne->SetExtractionModeToSpecifiedRegions();
    connOne->ColorRegionsOff();
    connOne->Update();

    for (int regionId = 0; regionId < nRegions; ++regionId)
    {
        if (shouldAbort && shouldAbort())
            break;

        if (onProgress)
            onProgress(100.0 * double(regionId + 1) / double(nRegions));

        connOne->InitializeSpecifiedRegionList();
        connOne->AddSpecifiedRegion(regionId);
        connOne->Modified();
        connOne->Update();

        vtkPolyData* raw = connOne->GetOutput();
        if (!raw || raw->GetNumberOfCells() == 0)
        {
            std::cout << "regionId " << regionId << " => cells 0\n";
            continue;
        }

        // IMPORTANT: compact points so bounds/length reflect ONLY this region
        vtkNew<vtkCleanPolyData> clean;
        clean->SetInputData(raw);
        clean->PointMergingOff();     // keep geometry identical; just drop unused points
        clean->Update();

        vtkPolyData* out = clean->GetOutput();
        if (!out || out->GetNumberOfCells() == 0 || out->GetNumberOfPoints() == 0)
        {
            std::cout << "regionId " << regionId << " => empty after clean\n";
            continue;
        }

        out->ComputeBounds();
        double b[6];
        out->GetBounds(b);
        const double len = out->GetLength();

        std::cout << "regionId " << regionId
                  << " => cells " << out->GetNumberOfCells()
                  << " points " << out->GetNumberOfPoints()
                  << " len " << len
                  << " bounds "
                  << b[0] << " " << b[1] << " "
                  << b[2] << " " << b[3] << " "
                  << b[4] << " " << b[5] << "\n";

        bool ok = true;

        // 0: bounding-box diagonal length
        if (params.selectedFilters[0])
        {
            if (len < params.thrDown_D || len > params.thrUp_D)
                ok = false;
        }

        // 1: Hausdorff
        if (ok && params.selectedFilters[1])
        {
            double center[3];
            out->GetCenter(center);
            const double hd = hausdorffToSphere(out, center);
            if (hd > params.thr_HAUSD)
                ok = false;
        }

        // 2: side-difference
        if (ok && params.selectedFilters[2])
        {
            const double metric = sideDifferenceMetric(b);
            if (metric > params.thr_S)
                ok = false;
        }

        if (!ok)
            continue;

        // Centroid from edges, to match legacy behavior.
        vtkNew<vtkExtractEdges> edge_extractor;
        edge_extractor->SetInputData(out);
        edge_extractor->Update();

        vtkPolyData* edges = edge_extractor->GetOutput();
        if (!edges || edges->GetNumberOfPoints() == 0)
            continue;

        const Point3d c_mm = centroidFromPoints(edges->GetPoints());

        Marker m;
        m.regionId = regionId;
        m.centroid = Point3d{c_mm.x / 1000.0, c_mm.y / 1000.0, c_mm.z / 1000.0};
        result.push_back(m);
    }

    return result;
}