#include "actors/VoxelActor.hpp"

#include <cmath>
#include <limits>

#include <vtkColorTransferFunction.h>
#include <vtkDoubleArray.h>
#include <vtkNew.h>
#include <vtkSmartVolumeMapper.h>
#include <vtkPiecewiseFunction.h>
#include <vtkPointData.h>
#include <vtkVolumeProperty.h>

namespace geopro::render {

namespace {

// 颜色/不透明度传递函数采样级数。
constexpr int kTransferSamples = 64;
// 体绘制最大不透明度（[vmin,vmax] 线性 0→kMaxOpacity）。
constexpr double kMaxOpacity = 0.15;

// NaN/留空格的哨兵：落在 [vmin,vmax] 之外，传递函数把它映射为完全透明。
double sentinel(double vmin) { return vmin - 1.0; }

}  // namespace

vtkSmartPointer<vtkVolume> buildVoxel(const geopro::core::ScalarVolume& vol,
                                      const geopro::core::ColorScale& cs,
                                      double ox, double oy, double oz,
                                      double dx, double dy, double dz,
                                      double vmin, double vmax,
                                      vtkSmartPointer<vtkImageData>& outImage)
{
    const int nx = vol.nx(), ny = vol.ny(), nz = vol.nz();

    // vmin/vmax 退化兜底，避免传递函数区间为零。
    if (vmin >= vmax) vmax = vmin + 1.0;
    const double blank = sentinel(vmin);

    auto img = vtkSmartPointer<vtkImageData>::New();
    img->SetDimensions(nx, ny, nz);
    img->SetOrigin(ox, oy, oz);
    img->SetSpacing(dx, dy, dz);

    vtkNew<vtkDoubleArray> sc;
    sc->SetName("v");
    sc->SetNumberOfTuples(static_cast<vtkIdType>(nx) * ny * nz);
    // 点序 i 最快、j 次之、k 最慢（匹配 vtkImageData 与 ScalarVolume::idx）。
    for (int k = 0; k < nz; ++k)
        for (int j = 0; j < ny; ++j)
            for (int i = 0; i < nx; ++i) {
                const double v = vol.at(i, j, k);
                const vtkIdType id = (static_cast<vtkIdType>(k) * ny + j) * nx + i;
                sc->SetValue(id, std::isnan(v) ? blank : v);  // NaN → 哨兵
            }
    img->GetPointData()->SetScalars(sc);
    outImage = img;

    // 颜色传递函数：在 [vmin,vmax] 按 ColorScale 采样若干 RGB 点。
    vtkNew<vtkColorTransferFunction> color;
    for (int t = 0; t < kTransferSamples; ++t) {
        const double val = vmin + (vmax - vmin) * t / (kTransferSamples - 1);
        const auto c = cs.colorAt(val);
        color->AddRGBPoint(val, c.r / 255.0, c.g / 255.0, c.b / 255.0);
    }

    // 不透明度传递函数：哨兵 → 0（透明）；[vmin,vmax] 线性递增到 kMaxOpacity。
    vtkNew<vtkPiecewiseFunction> opacity;
    opacity->AddPoint(blank, 0.0);
    opacity->AddPoint(vmin, 0.0);
    opacity->AddPoint(vmax, kMaxOpacity);

    // SmartVolumeMapper：有 GPU 走 GPU ray cast，否则自动回退 CPU，避免无 GPU 时卡死/失败。
    vtkNew<vtkSmartVolumeMapper> mapper;
    mapper->SetInputData(img);

    vtkNew<vtkVolumeProperty> prop;
    prop->SetColor(color);
    prop->SetScalarOpacity(opacity);
    prop->SetInterpolationTypeToLinear();
    prop->ShadeOff();

    auto volume = vtkSmartPointer<vtkVolume>::New();
    volume->SetMapper(mapper);
    volume->SetProperty(prop);
    return volume;
}

vtkSmartPointer<vtkVolume> buildVoxel(const geopro::core::ScalarVolume& vol,
                                      const geopro::core::ColorScale& cs,
                                      double ox, double oy, double oz,
                                      double dx, double dy, double dz,
                                      double vmin, double vmax)
{
    vtkSmartPointer<vtkImageData> ignored;
    return buildVoxel(vol, cs, ox, oy, oz, dx, dy, dz, vmin, vmax, ignored);
}

}  // namespace geopro::render