// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

/**
 * @file ct_reconstruction.hpp
 * @brief CT 再構成パイプライン — ラドン変換 + FBP + 反復法
 *
 * X 線 CT の投影データ (サイノグラム) から断面画像を再構成する。
 *
 * 提供機能:
 *   radonTransform(image, angles)        — ラドン変換 (前方投影)
 *   filteredBackProjection(sinogram, ...) — FBP (フィルタ補正逆投影法)
 *   artReconstruct(sinogram, ...)         — ART (代数的再構成法)
 *   sirtReconstruct(sinogram, ...)        — SIRT (同時反復再構成法)
 *
 * 用途: 医療 CT, 産業 CT, 電子顕微鏡トモグラフィ
 */

#ifndef CALX_IMAGE_CT_RECONSTRUCTION_HPP
#define CALX_IMAGE_CT_RECONSTRUCTION_HPP

#include <math/fft/fft.hpp>
#include <cmath>
#include <vector>
#include <complex>
#include <algorithm>
#include <numeric>
#include <cassert>
#include <cstddef>

namespace calx::image::ct {

/// サイノグラム: [angle][detector] の 2D 配列
template<typename T>
using Sinogram = std::vector<std::vector<T>>;

// ========================================================================
// ラドン変換 (前方投影)
// ========================================================================

/// 角度列を生成 [0, π) を numAngles 等分
template<typename T>
std::vector<T> defaultAngles(int numAngles) {
    std::vector<T> angles(numAngles);
    for (int i = 0; i < numAngles; ++i)
        angles[i] = T(i) * T(3.14159265358979323846) / T(numAngles);
    return angles;
}

/// ラドン変換: 画像 → サイノグラム
/// @param image    N×N の正方画像 (行優先)
/// @param N        画像サイズ
/// @param angles   投影角度の列 [rad]
/// @return         サイノグラム [numAngles][numDetectors]
template<typename T>
Sinogram<T> radonTransform(const std::vector<T>& image, int N,
                            const std::vector<T>& angles) {
    int numAngles = static_cast<int>(angles.size());
    int numDetectors = N;
    T center = T(N - 1) / T(2);

    Sinogram<T> sinogram(numAngles, std::vector<T>(numDetectors, T(0)));

    for (int ai = 0; ai < numAngles; ++ai) {
        T cosA = std::cos(angles[ai]);
        T sinA = std::sin(angles[ai]);

        for (int di = 0; di < numDetectors; ++di) {
            T det = T(di) - center;  // 検出器位置 (中心からのオフセット)
            T sum = T(0);

            // レイの方向: (-sinA, cosA), 検出器位置: det * (cosA, sinA)
            // レイ上の点: (det*cosA - t*sinA, det*sinA + t*cosA)
            for (int ti = 0; ti < N; ++ti) {
                T t = T(ti) - center;
                T x = det * cosA - t * sinA + center;
                T y = det * sinA + t * cosA + center;

                // バイリニア補間
                int x0 = static_cast<int>(std::floor(x));
                int y0 = static_cast<int>(std::floor(y));
                if (x0 < 0 || x0 >= N - 1 || y0 < 0 || y0 >= N - 1) continue;

                T fx = x - T(x0), fy = y - T(y0);
                sum += image[y0 * N + x0] * (T(1) - fx) * (T(1) - fy)
                     + image[y0 * N + x0 + 1] * fx * (T(1) - fy)
                     + image[(y0 + 1) * N + x0] * (T(1) - fx) * fy
                     + image[(y0 + 1) * N + x0 + 1] * fx * fy;
            }
            sinogram[ai][di] = sum;
        }
    }
    return sinogram;
}

// ========================================================================
// FBP フィルタ
// ========================================================================

/// FBP 用フィルタの種類
enum class FBPFilter {
    RamLak,      ///< Ram-Lak (ramp) — 基本フィルタ
    SheppLogan,  ///< Shepp-Logan — 穏やかな高周波抑制
    Cosine,      ///< コサイン窓
    Hamming,     ///< ハミング窓
    Hann         ///< ハン窓
};

/// フィルタ付き逆投影法 (Filtered Back Projection)
/// @param sinogram   サイノグラム [numAngles][numDetectors]
/// @param angles     投影角度 [rad]
/// @param N          再構成画像サイズ (N×N)
/// @param filter     FBP フィルタ種類
/// @return           再構成画像 (N×N, 行優先)
template<typename T>
std::vector<T> filteredBackProjection(const Sinogram<T>& sinogram,
                                       const std::vector<T>& angles,
                                       int N,
                                       FBPFilter filter = FBPFilter::RamLak) {
    int numAngles = static_cast<int>(sinogram.size());
    int numDetectors = static_cast<int>(sinogram[0].size());

    // FFT 長 (2のべき乗にゼロ詰め)
    int fftLen = 1;
    while (fftLen < numDetectors * 2) fftLen *= 2;

    // フィルタ構築 (周波数領域)
    std::vector<T> H(fftLen, T(0));
    for (int i = 0; i < fftLen; ++i) {
        T freq = (i <= fftLen / 2) ? T(i) / T(fftLen) : T(fftLen - i) / T(fftLen);
        T ramp = T(2) * freq;  // Ram-Lak

        T window = T(1);
        switch (filter) {
        case FBPFilter::RamLak:     window = T(1); break;
        case FBPFilter::SheppLogan:
            if (freq > T(0)) window = std::sin(T(3.14159265358979) * freq)
                                    / (T(3.14159265358979) * freq);
            break;
        case FBPFilter::Cosine:     window = std::cos(T(3.14159265358979) * freq / T(2)); break;
        case FBPFilter::Hamming:    window = T(0.54) + T(0.46) * std::cos(T(2) * T(3.14159265358979) * freq); break;
        case FBPFilter::Hann:       window = T(0.5) + T(0.5) * std::cos(T(2) * T(3.14159265358979) * freq); break;
        }
        H[i] = ramp * window;
    }

    // 各投影をフィルタリング
    using Complex = std::complex<T>;
    Sinogram<T> filtered(numAngles);

    for (int ai = 0; ai < numAngles; ++ai) {
        // ゼロ詰め + FFT
        std::vector<Complex> buf(fftLen, Complex(0));
        for (int i = 0; i < numDetectors; ++i) buf[i] = Complex(sinogram[ai][i]);

        FFT<Complex, T>::fft(std::span<Complex>(buf));

        // フィルタ適用
        for (int i = 0; i < fftLen; ++i) buf[i] *= H[i];

        // IFFT
        FFT<Complex, T>::ifft(std::span<Complex>(buf));

        filtered[ai].resize(numDetectors);
        for (int i = 0; i < numDetectors; ++i)
            filtered[ai][i] = buf[i].real();
    }

    // 逆投影
    std::vector<T> image(N * N, T(0));
    T center = T(N - 1) / T(2);
    T detCenter = T(numDetectors - 1) / T(2);

    for (int ai = 0; ai < numAngles; ++ai) {
        T cosA = std::cos(angles[ai]);
        T sinA = std::sin(angles[ai]);

        for (int y = 0; y < N; ++y) {
            for (int x = 0; x < N; ++x) {
                // (x, y) を検出器座標に投影
                T dx = T(x) - center;
                T dy = T(y) - center;
                T det = dx * cosA + dy * sinA + detCenter;

                // 線形補間
                int d0 = static_cast<int>(std::floor(det));
                if (d0 < 0 || d0 >= numDetectors - 1) continue;
                T frac = det - T(d0);
                T val = filtered[ai][d0] * (T(1) - frac)
                      + filtered[ai][d0 + 1] * frac;
                image[y * N + x] += val;
            }
        }
    }

    // 正規化
    T scale = T(3.14159265358979323846) / T(numAngles);
    for (auto& v : image) v *= scale;

    return image;
}

// ========================================================================
// ART (Algebraic Reconstruction Technique)
// ========================================================================

/// ART: Kaczmarz 法ベースの反復再構成
/// @param sinogram   サイノグラム
/// @param angles     投影角度
/// @param N          再構成サイズ
/// @param iterations 反復回数
/// @param relaxation 緩和係数 (0 < λ ≤ 1, デフォルト 0.25)
template<typename T>
std::vector<T> artReconstruct(const Sinogram<T>& sinogram,
                               const std::vector<T>& angles,
                               int N, int iterations = 10,
                               T relaxation = T(0.25)) {
    int numAngles = static_cast<int>(sinogram.size());
    int numDetectors = static_cast<int>(sinogram[0].size());
    T center = T(N - 1) / T(2);

    std::vector<T> image(N * N, T(0));

    for (int iter = 0; iter < iterations; ++iter) {
        for (int ai = 0; ai < numAngles; ++ai) {
            T cosA = std::cos(angles[ai]);
            T sinA = std::sin(angles[ai]);

            for (int di = 0; di < numDetectors; ++di) {
                T det = T(di) - center;

                // レイ上のピクセルを収集
                std::vector<std::pair<int, T>> ray;  // (pixel_index, weight)
                T norm2 = T(0);

                for (int ti = 0; ti < N; ++ti) {
                    T t = T(ti) - center;
                    T x = det * cosA - t * sinA + center;
                    T y = det * sinA + t * cosA + center;

                    int x0 = static_cast<int>(std::floor(x));
                    int y0 = static_cast<int>(std::floor(y));
                    if (x0 < 0 || x0 >= N - 1 || y0 < 0 || y0 >= N - 1) continue;

                    T fx = x - T(x0), fy = y - T(y0);
                    auto addPixel = [&](int px, int py, T w) {
                        if (w > T(1e-10)) {
                            ray.push_back({py * N + px, w});
                            norm2 += w * w;
                        }
                    };
                    addPixel(x0, y0, (T(1) - fx) * (T(1) - fy));
                    addPixel(x0 + 1, y0, fx * (T(1) - fy));
                    addPixel(x0, y0 + 1, (T(1) - fx) * fy);
                    addPixel(x0 + 1, y0 + 1, fx * fy);
                }

                if (norm2 < T(1e-15)) continue;

                // 現在の投影値を計算
                T proj = T(0);
                for (auto& [idx, w] : ray) proj += w * image[idx];

                // 補正
                T correction = relaxation * (sinogram[ai][di] - proj) / norm2;
                for (auto& [idx, w] : ray) image[idx] += correction * w;
            }
        }
    }
    return image;
}

/// SIRT: Simultaneous Iterative Reconstruction Technique
/// ART の全レイを同時に更新するバージョン (より滑らかな結果)
template<typename T>
std::vector<T> sirtReconstruct(const Sinogram<T>& sinogram,
                                const std::vector<T>& angles,
                                int N, int iterations = 50,
                                T relaxation = T(1)) {
    int numAngles = static_cast<int>(sinogram.size());
    int numDetectors = static_cast<int>(sinogram[0].size());
    T center = T(N - 1) / T(2);

    std::vector<T> image(N * N, T(0));

    for (int iter = 0; iter < iterations; ++iter) {
        // 前方投影 → 残差
        auto currentSinogram = radonTransform(image, N, angles);

        // 残差の逆投影
        std::vector<T> update(N * N, T(0));
        std::vector<T> weights(N * N, T(0));

        for (int ai = 0; ai < numAngles; ++ai) {
            T cosA = std::cos(angles[ai]);
            T sinA = std::sin(angles[ai]);

            for (int di = 0; di < numDetectors; ++di) {
                T residual = sinogram[ai][di] - currentSinogram[ai][di];
                T det = T(di) - center;

                for (int ti = 0; ti < N; ++ti) {
                    T t = T(ti) - center;
                    T x = det * cosA - t * sinA + center;
                    T y = det * sinA + t * cosA + center;

                    int x0 = static_cast<int>(std::floor(x));
                    int y0 = static_cast<int>(std::floor(y));
                    if (x0 < 0 || x0 >= N - 1 || y0 < 0 || y0 >= N - 1) continue;

                    T fx = x - T(x0), fy = y - T(y0);
                    auto addUpdate = [&](int px, int py, T w) {
                        int idx = py * N + px;
                        update[idx] += w * residual;
                        weights[idx] += w;
                    };
                    addUpdate(x0, y0, (T(1) - fx) * (T(1) - fy));
                    addUpdate(x0 + 1, y0, fx * (T(1) - fy));
                    addUpdate(x0, y0 + 1, (T(1) - fx) * fy);
                    addUpdate(x0 + 1, y0 + 1, fx * fy);
                }
            }
        }

        // 重み付き更新
        for (int i = 0; i < N * N; ++i) {
            if (weights[i] > T(1e-15))
                image[i] += relaxation * update[i] / weights[i];
        }
    }
    return image;
}

} // namespace calx::image::ct

#endif // CALX_IMAGE_CT_RECONSTRUCTION_HPP
