#include "actors/VoxelActor.hpp"

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

#include <vtkActor.h>
#include <vtkColorTransferFunction.h>
#include <vtkDoubleArray.h>
#include <vtkFloatArray.h>
#include <vtkFlyingEdges3D.h>
#include <vtkGPUVolumeRayCastMapper.h>
#include <vtkNew.h>
#include <vtkPolyData.h>
#include <vtkPolyDataMapper.h>
#include <vtkProperty.h>
#include <vtkShortArray.h>
#include <vtkSmartVolumeMapper.h>
#include <vtkPiecewiseFunction.h>
#include <vtkPointData.h>
#include <vtkUnsignedCharArray.h>
#include <vtkVolumeMapper.h>
#include <vtkVolumeProperty.h>

namespace geopro::render {

namespace {

// 颜色/不透明度传递函数采样级数。
constexpr int kTransferSamples = 64;

// 是否支持 GPU 体绘制（光线投射）。默认 true（有独显的常态）；无 GPU 机器由 setVolumeGpuSupported(false)
//   设回退。影响：mask 真白化只有 GPU mapper 支持 → 无 GPU 时不建 mask、改用 SmartVolumeMapper(自动 CPU 回退)，
//   空值仍靠不透明度传函(哨兵→0)透明，仅交界处少了 mask 的干净边（重现一圈细渗色）。
bool g_gpuVolumeSupported = true;

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

// 二值 mask 体（UCHAR，255=有效、0=空值）。与标量同维同 origin/spacing、同点序
//   （id=(k*ny+j)*nx+i）。空值格 mask=0 → 喂给 GPU ray cast 后被完全跳过：不着色、不参与
//   三线性插值，对齐 Surfer Blanking 的真白化语义（消除"空白处沿数据边界渗蓝")。
//   调用方在填标量的同一循环里写 m->SetValue(id, valid?255:0)。
vtkSmartPointer<vtkImageData> makeMaskLike(int nx, int ny, int nz,
                                           double ox, double oy, double oz,
                                           double dx, double dy, double dz,
                                           vtkUnsignedCharArray*& outArr)
{
    auto mask = vtkSmartPointer<vtkImageData>::New();
    mask->SetDimensions(nx, ny, nz);
    mask->SetOrigin(ox, oy, oz);
    mask->SetSpacing(dx, dy, dz);
    vtkNew<vtkUnsignedCharArray> m;
    m->SetName("mask");
    m->SetNumberOfTuples(static_cast<vtkIdType>(nx) * ny * nz);
    mask->GetPointData()->SetScalars(m);
    outArr = m;  // image 持有引用，循环结束前有效
    return mask;
}

// double/int16 两版公用的 mapper+property+volume 组装。mask 非空 → 用 GPU ray cast + 二值 mask
//   做真白化（SmartVolumeMapper 不转发 mask，故走 GPU mapper；桌面端恒有 GL 上下文）；
//   mask 为空 → 保留 SmartVolumeMapper（GPU/CPU 自适应）。
vtkSmartPointer<vtkVolume> assembleVolume(vtkImageData* img,
                                          vtkColorTransferFunction* color,
                                          vtkPiecewiseFunction* opacity,
                                          vtkImageData* mask)
{
    // 采样距离 + 不透明度单位距离用到几何尺度。
    double sp[3];
    img->GetSpacing(sp);
    const double minSp =
        std::min({std::abs(sp[0]), std::abs(sp[1]), std::abs(sp[2])});  // 最细体素维度
    double bnd[6];
    img->GetBounds(bnd);
    const double ex = std::abs(bnd[1] - bnd[0]), ey = std::abs(bnd[3] - bnd[2]),
                 ez = std::abs(bnd[5] - bnd[4]);
    const double diag = std::sqrt(ex * ex + ey * ey + ez * ez);
    // 不透明度单位距离的尺度基准：
    //   - 【各向异性】体(细长，如雷达 375×1.4×5m)：对角线被长轴主宰(375m)→ 单位距离过大 →
    //     不透明度只在 100% 才实心、稍降即很透(用户实测)。改用【特征尺度=三轴几何平均 cbrt】，
    //     对各向异性稳健。
    //   - 【近立方】体(反演)：维持原对角线，观感不变(门控：长短轴比 ≤ kAnisoRatio 走对角线)。
    constexpr double kAnisoRatio = 4.0;
    const double maxE = std::max({ex, ey, ez}), minE = std::min({ex, ey, ez});
    const bool anisotropic = (minE > 0.0) && (maxE / minE > kAnisoRatio);
    const double charLen = anisotropic ? std::cbrt(ex * ey * ez) : diag;

    // 大体(如雷达：24M 体素、深度采样距 mm 级 → 单条光线上千采样步)开启【交互期】采样距自适应：
    //   旋转时 VTK 自动加大采样步(变粗)保流畅，停手即恢复设定的细采样距(0.3×minSp)出全质量帧。
    //   只是渲染期降采样、【绝不动数据】；切片/异常取自全分辨率体，保真不受影响。
    //   小体(反演，实测~7ms/帧)保持全程全质量，避免"停手补高清"的视觉突跳。
    constexpr vtkIdType kInteractiveLodVoxels = 4'000'000;
    const int interactiveAdjust = (img->GetNumberOfPoints() > kInteractiveLodVoxels) ? 1 : 0;

    vtkSmartPointer<vtkVolumeMapper> mapper;
    if (mask && g_gpuVolumeSupported) {
        // 真白化：mask=0 体素被光线投射完全跳过，杜绝空值格沿边界渗蓝。需 GPU 光线投射支持。
        vtkNew<vtkGPUVolumeRayCastMapper> gpu;
        gpu->SetInputData(img);
        gpu->SetMaskInput(mask);
        gpu->SetMaskTypeToBinary();
        gpu->SetAutoAdjustSampleDistances(interactiveAdjust);  // 大体:交互降采样保流畅,停手全质量;小体:全程全质量
        // 静止帧用【细】采样距离(0.3×minSp)：否则用粗默认值 → 看到一层层体素(分层伪影)。
        if (minSp > 0) gpu->SetSampleDistance(static_cast<float>(0.3 * minSp));
        // 抖动：用噪声纹理微扰每条光线的采样起点，消除规则采样面造成的「木纹/分层」伪影（VTK 官方此用途）。
        gpu->SetUseJittering(1);
        mapper = gpu;
    } else {
        // SmartVolumeMapper：有 GPU 走 GPU ray cast，否则自动回退 CPU，避免无 GPU 时卡死/失败。
        vtkNew<vtkSmartVolumeMapper> sm;
        sm->SetInputData(img);
        // 大体交互降采样保流畅(停手恢复全质量)；小体全程全质量(GPU 足够快, 实测 ~7ms/帧)避免突跳。
        sm->SetAutoAdjustSampleDistances(interactiveAdjust);
        sm->SetInteractiveAdjustSampleDistances(interactiveAdjust);
        mapper = sm;
    }

    vtkNew<vtkVolumeProperty> prop;
    prop->SetColor(color);
    prop->SetScalarOpacity(opacity);
    prop->SetInterpolationTypeToLinear();
    prop->ShadeOff();
    // 不透明度单位距离 = 尺度基准(charLen：各向异性=特征尺度 / 近立方=对角线) × kOpacityUnitFraction。
    //   控制累积速度使「不透明度」滑块有层次；细长体走特征尺度后不再"只有 100% 实心、99% 即很透"。
    constexpr double kOpacityUnitFraction = 0.1;
    if (charLen > 0) prop->SetScalarOpacityUnitDistance(kOpacityUnitFraction * charLen);

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

// int16 量化域传函组装：颜色对每量化级 qv 用 q.toPhys(qv) 反查 ColorScale；
// 不透明度 kBlank→0(透明)、[qmin,qmax] 线性递增到 kMaxOpacity。
// 与 buildVoxelI16 内联逻辑一致（被 buildVoxelI16FromImage 复用）。
vtkSmartPointer<vtkVolume> assembleVolumeI16(vtkImageData* img,
                                             const geopro::core::Quant& q,
                                             const geopro::core::ColorScale& cs,
                                             double vminPhys, double vmaxPhys)
{
    if (vminPhys >= vmaxPhys) vmaxPhys = vminPhys + 1.0;

    const std::int16_t qmin = q.toQ(vminPhys);
    const std::int16_t qmax = q.toQ(vmaxPhys);
    const double qminD = static_cast<double>(qmin);
    const double qmaxD = static_cast<double>(qmax);

    vtkNew<vtkColorTransferFunction> color;
    for (int t = 0; t < kTransferSamples; ++t) {
        const double qd = qminD + (qmaxD - qminD) * t / (kTransferSamples - 1);
        const auto qvLevel = static_cast<std::int16_t>(std::lround(qd));
        const double phys = q.toPhys(qvLevel);
        const auto c = cs.colorAt(phys);
        color->AddRGBPoint(qd, c.r / 255.0, c.g / 255.0, c.b / 255.0);
    }

    vtkNew<vtkPiecewiseFunction> opacity;
    opacity->AddPoint(static_cast<double>(geopro::core::ScalarVolumeI16::kBlank), 0.0);
    for (int t = 0; t < kTransferSamples; ++t) {
        const double qd = qminD + (qmaxD - qminD) * t / (kTransferSamples - 1);
        const auto qvLevel = static_cast<std::int16_t>(std::lround(qd));
        const double phys = q.toPhys(qvLevel);
        opacity->AddPoint(qd, cs.colorAt(phys).a / 255.0 * cs.globalOpacity());
    }

    // 由预建 short 体扫出二值 mask（kBlank→0 跳过）。稠密体（无 kBlank）→ 全 255，等价无 mask。
    int dims[3];
    img->GetDimensions(dims);
    vtkUnsignedCharArray* mArr = nullptr;
    auto mask = makeMaskLike(dims[0], dims[1], dims[2], img->GetOrigin()[0], img->GetOrigin()[1],
                             img->GetOrigin()[2], img->GetSpacing()[0], img->GetSpacing()[1],
                             img->GetSpacing()[2], mArr);
    if (auto* sc = vtkShortArray::SafeDownCast(img->GetPointData()->GetScalars())) {
        const vtkIdType n = sc->GetNumberOfTuples();
        for (vtkIdType id = 0; id < n; ++id)
            mArr->SetValue(id, sc->GetValue(id) == geopro::core::ScalarVolumeI16::kBlank ? 0 : 255);
    }
    return assembleVolume(img, color, opacity, mask);
}

}  // namespace

void setVolumeGpuSupported(bool ok) { g_gpuVolumeSupported = ok; }

void updateVolumeColors(vtkVolume* volume, const geopro::core::ColorScale& cs, double vmin,
                        double vmax) {
    if (!volume || !volume->GetProperty()) return;
    if (vmin >= vmax) vmax = vmin + 1.0;
    const double blank = sentinel(vmin);
    // 与 buildVoxel(float 路径) 同口径重建颜色/不透明度传函，原地换到已有 actor 上（不重建 image →
    //   切片基底不变、不被关闭）。
    vtkNew<vtkColorTransferFunction> color;
    vtkNew<vtkPiecewiseFunction> opacity;
    opacity->AddPoint(blank, 0.0);
    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);
        opacity->AddPoint(val, c.a / 255.0 * cs.globalOpacity());
    }
    volume->GetProperty()->SetColor(color);
    volume->GetProperty()->SetScalarOpacity(opacity);
}

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);

    // 标量用 float（非 double）：OpenGL 无原生 double 体纹理，GPU 体绘制对 double 处理不稳/部分驱动间歇
    //   出空（偶发不渲染根因之一），且省一半显存。float 精度对可视化足够。
    // 二值 mask：NaN 空格→0（光线投射跳过，真白化），有值→255。与标量同循环填，免二次扫描。
    vtkUnsignedCharArray* mArr = nullptr;
    auto mask = makeMaskLike(nx, ny, nz, ox, oy, oz, dx, dy, dz, mArr);

    vtkNew<vtkFloatArray> 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;
                const bool isBlank = std::isnan(v);
                sc->SetValue(id, static_cast<float>(isBlank ? blank : v));  // NaN → 哨兵
                mArr->SetValue(id, isBlank ? 0 : 255);
            }
    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（透明）；区间内由色阶 alpha 驱动，再乘体密度主控 kMaxOpacity。
    //   体素不透明度 = (色阶 alpha/255) × kMaxOpacity（整体透明度已在配置时乘进 alpha）。
    //   alpha=0 → 真透明；alpha=255（无 alpha 色阶默认）→ 维持 kMaxOpacity 的通透手感，不回归。
    vtkNew<vtkPiecewiseFunction> opacity;
    opacity->AddPoint(blank, 0.0);
    for (int t = 0; t < kTransferSamples; ++t) {
        const double val = vmin + (vmax - vmin) * t / (kTransferSamples - 1);
        opacity->AddPoint(val, cs.colorAt(val).a / 255.0 * cs.globalOpacity());
    }

    return assembleVolume(img, color, opacity, mask);
}

vtkSmartPointer<vtkVolume> buildVoxelI16(const geopro::core::ScalarVolumeI16& vol,
                                         const geopro::core::Quant& q,
                                         const geopro::core::ColorScale& cs,
                                         double ox, double oy, double oz,
                                         double dx, double dy, double dz,
                                         double vminPhys, double vmaxPhys,
                                         vtkSmartPointer<vtkImageData>& outImage)
{
    const int nx = vol.nx(), ny = vol.ny(), nz = vol.nz();

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

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

    // 二值 mask：kBlank 空格→0（真白化跳过），有值→255。与标量同循环填。
    vtkUnsignedCharArray* mArr = nullptr;
    auto mask = makeMaskLike(nx, ny, nz, ox, oy, oz, dx, dy, dz, mArr);

    vtkNew<vtkShortArray> sc;
    sc->SetName("v");
    sc->SetNumberOfTuples(static_cast<vtkIdType>(nx) * ny * nz);
    // 点序 i 最快、j 次之、k 最慢（匹配 vtkImageData 与 ScalarVolumeI16::idx）。
    // kBlank 原样保留，由量化域传递函数映射为透明。
    for (int k = 0; k < nz; ++k)
        for (int j = 0; j < ny; ++j)
            for (int i = 0; i < nx; ++i) {
                const std::int16_t qv = vol.at(i, j, k);
                const vtkIdType id = (static_cast<vtkIdType>(k) * ny + j) * nx + i;
                sc->SetValue(id, qv);
                mArr->SetValue(id, qv == geopro::core::ScalarVolumeI16::kBlank ? 0 : 255);
            }
    img->GetPointData()->SetScalars(sc);
    outImage = img;

    // 传递函数在量化域取（标量本身是 int16 量化值）。
    const std::int16_t qmin = q.toQ(vminPhys);
    const std::int16_t qmax = q.toQ(vmaxPhys);
    const double qminD = static_cast<double>(qmin);
    const double qmaxD = static_cast<double>(qmax);

    // 颜色传递函数：对每个量化级 qv，物理值 phys=q.toPhys(qv)，用 double 版相同方式取色。
    vtkNew<vtkColorTransferFunction> color;
    for (int t = 0; t < kTransferSamples; ++t) {
        const double qd = qminD + (qmaxD - qminD) * t / (kTransferSamples - 1);
        const auto qvLevel = static_cast<std::int16_t>(std::lround(qd));
        const double phys = q.toPhys(qvLevel);
        const auto c = cs.colorAt(phys);
        color->AddRGBPoint(qd, c.r / 255.0, c.g / 255.0, c.b / 255.0);
    }

    // 不透明度传递函数（量化域）：kBlank → 0（透明）；区间内由色阶 alpha 驱动 × 体密度主控 kMaxOpacity。
    vtkNew<vtkPiecewiseFunction> opacity;
    opacity->AddPoint(static_cast<double>(geopro::core::ScalarVolumeI16::kBlank), 0.0);
    for (int t = 0; t < kTransferSamples; ++t) {
        const double qd = qminD + (qmaxD - qminD) * t / (kTransferSamples - 1);
        const auto qvLevel = static_cast<std::int16_t>(std::lround(qd));
        const double phys = q.toPhys(qvLevel);
        opacity->AddPoint(qd, cs.colorAt(phys).a / 255.0 * cs.globalOpacity());
    }

    return assembleVolume(img, color, opacity, mask);
}

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);
}

vtkSmartPointer<vtkVolume> buildVoxelI16FromImage(vtkImageData* shortImg,
                                                  const geopro::core::Quant& q,
                                                  const geopro::core::ColorScale& cs,
                                                  double vminPhys, double vmaxPhys)
{
    // 图像由 source 预构建(VTK_SHORT,带 origin/spacing),直接成体；传函复用量化域逻辑。
    return assembleVolumeI16(shortImg, q, cs, vminPhys, vmaxPhys);
}

vtkSmartPointer<vtkActor> buildIsosurface(vtkImageData* img, const geopro::core::ColorScale& cs,
                                          double vmin, double vmax, double isoValue)
{
    if (!img) return nullptr;
    if (vmin >= vmax) vmax = vmin + 1.0;
    // 阈值钳进 (vmin,vmax)：=vmin 会沿留空哨兵边界成面、=vmax 抽不出。
    const double eps = 1e-6 * (vmax - vmin);
    isoValue = std::max(vmin + eps, std::min(vmax - eps, isoValue));

    vtkNew<vtkFlyingEdges3D> fe;
    fe->SetInputData(img);
    fe->SetValue(0, isoValue);
    fe->ComputeNormalsOn();
    fe->ComputeGradientsOff();
    fe->ComputeScalarsOff();
    fe->Update();
    if (!fe->GetOutput() || fe->GetOutput()->GetNumberOfPoints() == 0) return nullptr;  // 无超阈区

    vtkNew<vtkPolyDataMapper> mapper;
    mapper->SetInputConnection(fe->GetOutputPort());
    mapper->ScalarVisibilityOff();  // 用 actor 实色，不按标量着色

    auto actor = vtkSmartPointer<vtkActor>::New();
    actor->SetMapper(mapper);
    const auto c = cs.colorAt(isoValue);  // 阈值处的色（高值多为暖红，复刻参考图红块）
    actor->GetProperty()->SetColor(c.r / 255.0, c.g / 255.0, c.b / 255.0);
    actor->GetProperty()->SetOpacity(1.0);  // 不透明实心
    return actor;
}

}  // namespace geopro::render