#include "SkullRemovalService.h"

#include "connectITKVTK.h"

#include <vtkImageData.h>
#include <vtkImageExport.h>
#include <vtkImageImport.h>
#include <vtkSmartPointer.h>

#include <itkBinaryDilateImageFilter.h>
#include <itkBinaryErodeImageFilter.h>
#include <itkBinaryThresholdImageFilter.h>
#include <itkCastImageFilter.h>
#include <itkConnectedThresholdImageFilter.h>
#include <itkCurvatureFlowImageFilter.h>
#include <itkImage.h>
#include <itkMaskNegatedImageFilter.h>
#include <itkVTKImageImport.h>
#include <itkExceptionObject.h>

#include <stdexcept>

namespace
{
using InternalPixelType = float;
constexpr unsigned int Dimension = 3;

using InternalImageType = itk::Image<InternalPixelType, Dimension>;
using MaskPixelType = unsigned char;
using MaskImageType = itk::Image<MaskPixelType, Dimension>;

using ImportFilterType = itk::VTKImageImport<InternalImageType>;
using CastingINVFilterType = itk::CastImageFilter<InternalImageType, InternalImageType>;
using CurvatureFlowImageFilterType = itk::CurvatureFlowImageFilter<InternalImageType, InternalImageType>;
using ConnectedThresholdImageFilterType = itk::ConnectedThresholdImageFilter<InternalImageType, MaskImageType>;
using CastingFilterType = itk::CastImageFilter<MaskImageType, MaskImageType>;

using StructuringElementType = itk::BinaryBallStructuringElement<MaskPixelType, Dimension>;
using DilateFilterType = itk::BinaryDilateImageFilter<MaskImageType, MaskImageType, StructuringElementType>;
using ErodeFilterType  = itk::BinaryErodeImageFilter<MaskImageType, MaskImageType, StructuringElementType>;
using MaskNegatedFilterType = itk::MaskNegatedImageFilter<InternalImageType, MaskImageType, InternalImageType>;
} // namespace

vtkImageData* SkullRemovalService::run(vtkImageData* inputVolume,
                                      double thr_low,
                                      double thr_up,
                                      std::atomic_bool* abortFlag,
                                      std::function<void(const std::string&, double)> progressCb)
{
    if (!inputVolume)
        throw std::runtime_error("SkullRemovalService::run: inputVolume is null");

    auto isAborted = [&]() -> bool { return abortFlag && abortFlag->load(); };

    if (progressCb) progressCb("Skull masking: preparing pipeline", 0.0);
    if (isAborted()) return nullptr;

    // VTK -> ITK
    vtkSmartPointer<vtkImageExport> vtkExporter = vtkSmartPointer<vtkImageExport>::New();
    vtkExporter->SetInputData(inputVolume);

    auto itkImporter = ImportFilterType::New();
    ConnectPipelines(vtkExporter.GetPointer(), itkImporter);

    // Pipeline (mirrors old gSkullRemoval::runFilter in a pure form)
    auto smoothing = CurvatureFlowImageFilterType::New();
    auto connectedThreshold = ConnectedThresholdImageFilterType::New();
    auto caster_mask = CastingFilterType::New();
    auto binaryDilateFilter = DilateFilterType::New();
    auto binaryErodeFilter = ErodeFilterType::New();
    auto maskNegatedFilter = MaskNegatedFilterType::New();

    smoothing->SetNumberOfIterations(10);
    smoothing->SetTimeStep(0.125);

    smoothing->SetInput(itkImporter->GetOutput());
    connectedThreshold->SetInput(smoothing->GetOutput());

    connectedThreshold->SetLower(static_cast<InternalPixelType>(thr_low));
    connectedThreshold->SetUpper(static_cast<InternalPixelType>(thr_up));
    connectedThreshold->SetReplaceValue(255);

    // Seed at volume center
    connectedThreshold->UpdateOutputInformation();
    InternalImageType::Pointer in = itkImporter->GetOutput();
    const auto region = in->GetBufferedRegion();
    const auto size = region.GetSize();
    const auto start = region.GetIndex();

    InternalImageType::IndexType seed;
    seed[0] = start[0] + static_cast<long>(size[0] / 2);
    seed[1] = start[1] + static_cast<long>(size[1] / 2);
    seed[2] = start[2] + static_cast<long>(size[2] / 2);
    connectedThreshold->SetSeed(seed);

    caster_mask->SetInput(connectedThreshold->GetOutput());

    StructuringElementType kernel;
    kernel.SetRadius(10);
    kernel.CreateStructuringElement();

    binaryDilateFilter->SetKernel(kernel);
    binaryErodeFilter->SetKernel(kernel);
    binaryDilateFilter->SetDilateValue(255);
    binaryErodeFilter->SetErodeValue(255);

    binaryDilateFilter->SetInput(caster_mask->GetOutput());
    binaryErodeFilter->SetInput(binaryDilateFilter->GetOutput());

    maskNegatedFilter->SetInput1(itkImporter->GetOutput());
    maskNegatedFilter->SetInput2(binaryErodeFilter->GetOutput());

    if (progressCb) progressCb("Skull masking: running", 0.5);
    if (isAborted()) return nullptr;

    try {
        maskNegatedFilter->Update();
    } catch (const itk::ExceptionObject& ex) {
        if (isAborted()) return nullptr;
        throw std::runtime_error(std::string("SkullRemovalService: ITK exception: ") + ex.GetDescription());
    }

    if (isAborted()) return nullptr;

    // ITK -> VTK
    vtkSmartPointer<vtkImageImport> vtkImporter = vtkSmartPointer<vtkImageImport>::New();
    auto itkExporter = itk::VTKImageExport<InternalImageType>::New();
    itkExporter->SetInput(maskNegatedFilter->GetOutput());
    ConnectPipelines(itkExporter, vtkImporter);

    vtkImporter->Update();

    vtkImageData* out = vtkImageData::New();
    out->DeepCopy(vtkImporter->GetOutput());

    if (progressCb) progressCb("Skull masking: done", 1.0);
    return out;
}