#include "itkFourierStripeArtifactImageFilter.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkResampleImageFilter.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkCartesianToPolarTransform.h"
#include "itkPolarToCartesianTransform.h"
#include"itkTranslationTransform.h"
#include "itkRescaleIntensityImageFilter.h"
#include <iostream>
#include "itkFFTPadImageFilter.h"
#include <cmath>
#include <random>
#include <vector>
#include <sstream>
#include <iomanip>

int main(int argc, char* argv[])
{
    if (argc < 3)
    {
        std::cerr << "Usage: " << argv[0] << " inputImage outputBaseName" << std::endl;
        std::cerr << "Note: outputBaseName will be appended with parameter values." << std::endl;
        return EXIT_FAILURE;
    }

    constexpr unsigned int Dimension = 2;
    using PixelType = float;
    using ImageType = itk::Image< PixelType, Dimension >;

    // Read input image
    using ReaderType = itk::ImageFileReader< ImageType >;
    ReaderType::Pointer reader = ReaderType::New();
    reader->SetFileName(argv[1]);
    reader->Update();

    ImageType::Pointer inputImage = reader->GetOutput();

    // 极坐标变换
    ImageType::RegionType region = inputImage->GetLargestPossibleRegion();
    ImageType::SizeType size = region.GetSize();
    ImageType::IndexType centerIndex;
    centerIndex[0] = size[0] / 2;
    centerIndex[1] = size[1] / 2;
    ImageType::PointType centerPoint;
    inputImage->TransformIndexToPhysicalPoint(centerIndex, centerPoint);

    using PolarToCartesianType = itk::PolarToCartesianTransform<double, Dimension>;
    auto toCartesian = PolarToCartesianType::New();
    toCartesian->SetCenter(centerPoint);

    // Resample: 变换必须是 输出(θ,r) → 输入(x,y)
    using ResampleFilterType = itk::ResampleImageFilter< ImageType, ImageType >;
    auto polarResampler = ResampleFilterType::New();
    polarResampler->SetTransform(toCartesian);
    polarResampler->SetInput(inputImage);

    // 角向采样
    const unsigned int numAngles = 1000;
    const double dtheta = 2.0 * itk::Math::pi / static_cast<double>(numAngles);

    // 计算最大半径（从中心到四个角的最大距离）
    double rmax = 0.0;
    for (int dx : {0, static_cast<int>(size[0]) - 1})
        for (int dy : {0, static_cast<int>(size[1]) - 1}) {
            ImageType::IndexType idx{ {dx, dy} };
            ImageType::PointType p; inputImage->TransformIndexToPhysicalPoint(idx, p);
            rmax = std::max(rmax, centerPoint.EuclideanDistanceTo(p));
        }

    // r 方向步长：用输入图像的最小 spacing，避免漏采样
    const auto inSpacing = inputImage->GetSpacing();
    const double dr = std::min(inSpacing[0], inSpacing[1]);

    ImageType::SizeType polarSize;
    polarSize[0] = numAngles;                         // θ 采样数
    polarSize[1] = static_cast<size_t>(std::ceil(rmax / dr)); // r 采样数

    ImageType::SpacingType polarSpacing;
    polarSpacing[0] = dtheta; // 弧度/像素
    polarSpacing[1] = dr;     // 物理长度/像素

    ImageType::PointType polarOrigin;
    polarOrigin[0] = 0.0;     // θ 从 0 开始（弧度）
    polarOrigin[1] = 0.0;     // r 从 0 开始

    polarResampler->SetSize(polarSize);
    polarResampler->SetOutputOrigin(polarOrigin);
    polarResampler->SetOutputSpacing(polarSpacing);
    polarResampler->SetInterpolator(itk::LinearInterpolateImageFunction<ImageType, double>::New());
    polarResampler->SetDefaultPixelValue(-1.f); // 调试用：越界就会看到 -1

    polarResampler->Update();
    ImageType::Pointer polarImage = polarResampler->GetOutput();

    // 保存极坐标图像
    using WriterType = itk::ImageFileWriter<ImageType>;
    WriterType::Pointer polarWriter = WriterType::New();
    polarWriter->SetFileName("polar_image.mha");
    polarWriter->SetInput(polarImage);
    polarWriter->Update();

    // 定义多组参数
    std::vector<double> sigmaValues = { 10,20,30,40,50 }; // 您可以修改这些值  值越大滤波效果越强
    std::vector<double> frequencyValues = { 0.02,0.03,0.04, 0.05 }; // 您可以修改这些值  值越小滤波效果越强

    // 循环遍历所有参数组合
    for (double sigma : sigmaValues) {
        for (double frequency : frequencyValues) {
            std::cout << "Processing with Sigma=" << sigma << ", Frequency=" << frequency << std::endl;

            // ========== 新增部分：条纹伪影滤波 ==========
            using PadFilterType = itk::FFTPadImageFilter< ImageType >;
            PadFilterType::Pointer padFilter = PadFilterType::New();
            padFilter->SetInput(polarImage);

            using FourierFilterType = itk::FourierStripeArtifactImageFilter<ImageType>;
            auto fourierFilter = FourierFilterType::New();
            fourierFilter->SetInput(padFilter->GetOutput());
            fourierFilter->SetDirection(0);  // 沿角度方向滤波（θ轴）
            fourierFilter->SetSigma(sigma);
            fourierFilter->SetStartFrequency(frequency);
            fourierFilter->Update();
            ImageType::Pointer filteredPolarImage = fourierFilter->GetOutput();

            // 创建包含参数的文件名
            std::ostringstream oss;
            oss << "filtered_polar_s" << std::fixed << std::setprecision(2) << sigma
                << "_f" << std::fixed << std::setprecision(3) << frequency << ".mha";

            // 保存滤波后的极坐标图像
            WriterType::Pointer filteredPolarWriter = WriterType::New();
            filteredPolarWriter->SetFileName(oss.str());
            filteredPolarWriter->SetInput(filteredPolarImage);
          //  filteredPolarWriter->Update();

            // ========== 新增部分：变换回笛卡尔坐标系 ==========
            // 创建笛卡尔→极坐标变换（输出坐标系→输入坐标系）
            using CartesianToPolarType = itk::CartesianToPolarTransform<double, Dimension>;
            auto toPolar = CartesianToPolarType::New();
            toPolar->SetCenter(centerPoint);  // 使用相同的中心点

            // 创建重采样器将图像转回笛卡尔坐标系
            auto cartesianResampler = ResampleFilterType::New();
            cartesianResampler->SetTransform(toPolar);  // 输出(笛卡尔)→输入(极坐标)
            cartesianResampler->SetInput(filteredPolarImage);

            // 设置输出参数为原始图像的空间参数
            cartesianResampler->SetSize(size);                 // 原始图像尺寸
            cartesianResampler->SetOutputOrigin(inputImage->GetOrigin());
            cartesianResampler->SetOutputSpacing(inputImage->GetSpacing());
            cartesianResampler->SetOutputDirection(inputImage->GetDirection());
            cartesianResampler->SetInterpolator(itk::LinearInterpolateImageFunction<ImageType, double>::New());
            cartesianResampler->SetDefaultPixelValue(0);  // 默认背景值

            cartesianResampler->Update();
            ImageType::Pointer outputImage = cartesianResampler->GetOutput();

            // ========== 保存最终结果 ==========
            // 创建包含参数的输出文件名
            std::ostringstream outputOss;
            outputOss <<"image/"<< argv[2] << "_s" << std::fixed << std::setprecision(2) << sigma
                << "_f" << std::fixed << std::setprecision(3) << frequency << ".mha";

            WriterType::Pointer writer = WriterType::New();
            writer->SetFileName(outputOss.str());
            writer->SetInput(outputImage);
            writer->Update();

            std::cout << "Saved: " << outputOss.str() << std::endl;
        }
    }

    return EXIT_SUCCESS;
}