#include "acoustic_analyzer/core/gcc_phat_localizer.h"
#include <cmath>
#include <vector>
#include <Eigen/SVD>
#include <memory>

namespace acoustic {

namespace {

// Simple iterative Cooley-Tukey FFT for power-of-2 sizes.
// Sufficient for GCC-PHAT delay estimation.
struct Complex { float r = 0.0f, i = 0.0f; };

void fft_iterative(Complex* data, int n, bool inverse) {
    int j = 0;
    for (int i = 1; i < n; ++i) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) std::swap(data[i], data[j]);
    }
    for (int len = 2; len <= n; len <<= 1) {
        float ang = 2.0f * static_cast<float>(M_PI) / len * (inverse ? 1.0f : -1.0f);
        float wlen_r = std::cos(ang);
        float wlen_i = std::sin(ang);
        for (int i = 0; i < n; i += len) {
            float wr = 1.0f, wi = 0.0f;
            for (int k = 0; k < len / 2; ++k) {
                int u = i + k;
                int v = i + k + len / 2;
                float ur = data[u].r, ui = data[u].i;
                float vr = data[v].r * wr - data[v].i * wi;
                float vi = data[v].r * wi + data[v].i * wr;
                data[u].r = ur + vr;
                data[u].i = ui + vi;
                data[v].r = ur - vr;
                data[v].i = ui - vi;
                float wnr = wr * wlen_r - wi * wlen_i;
                float wni = wr * wlen_i + wi * wlen_r;
                wr = wnr;
                wi = wni;
            }
        }
    }
    if (inverse) {
        float inv_n = 1.0f / n;
        for (int i = 0; i < n; ++i) {
            data[i].r *= inv_n;
            data[i].i *= inv_n;
        }
    }
}

} // anonymous namespace

struct GccPhatLocalizer::Impl {
    MicArrayConfig mic_config;
    int sample_rate;
    float max_tdoa;
    int interp_factor;

    Impl(const MicArrayConfig& config, int sr, float max_t, int interp)
        : mic_config(config), sample_rate(sr), max_tdoa(max_t), interp_factor(interp) {
        if (mic_config.layout == "cross" && mic_config.positions.empty()) {
            float s = mic_config.spacing;
            mic_config.positions = {
                {0.0f,  s, 0.0f},
                { s, 0.0f, 0.0f},
                {0.0f, -s, 0.0f},
                {-s, 0.0f, 0.0f}
            };
        } else if (mic_config.layout == "linear" && mic_config.positions.empty()) {
            float s = mic_config.spacing;
            for (uint32_t i = 0; i < mic_config.num_mics; ++i) {
                mic_config.positions.push_back({i * s, 0.0f, 0.0f});
            }
        } else if (mic_config.layout == "circular" && mic_config.positions.empty()) {
            float s = mic_config.spacing;
            float radius = s / (2.0f * std::sin(static_cast<float>(M_PI) / mic_config.num_mics));
            for (uint32_t i = 0; i < mic_config.num_mics; ++i) {
                float angle = 2.0f * static_cast<float>(M_PI) * i / mic_config.num_mics;
                mic_config.positions.push_back({radius * std::cos(angle), radius * std::sin(angle), 0.0f});
            }
        }
    }

    float ComputeGccPhatDelay(const float* ch1, const float* ch2, std::size_t n) {
        if (n == 0) return 0.0f;
        // 防止 nfft 溢出：上限 1<<22（约 400 万点，足够 16kHz@120s 音频）
        constexpr std::size_t MAX_NFFT = static_cast<std::size_t>(1) << 22;
        std::size_t nfft = 1;
        while (nfft < 2 * n && nfft < MAX_NFFT) nfft <<= 1;
        if (nfft > MAX_NFFT) nfft = MAX_NFFT;

        std::vector<Complex> a(nfft), b(nfft), c(nfft);
        for (std::size_t i = 0; i < nfft; ++i) {
            a[i].r = (i < n) ? ch1[i] : 0.0f; a[i].i = 0;
            b[i].r = (i < n) ? ch2[i] : 0.0f; b[i].i = 0;
        }
        fft_iterative(a.data(), static_cast<int>(nfft), false);
        fft_iterative(b.data(), static_cast<int>(nfft), false);

        for (std::size_t i = 0; i < nfft; ++i) {
            float real = a[i].r * b[i].r + a[i].i * b[i].i;
            float imag = a[i].i * b[i].r - a[i].r * b[i].i;
            float mag = std::sqrt(real * real + imag * imag) + 1e-12f;
            c[i].r = real / mag;
            c[i].i = imag / mag;
        }
        fft_iterative(c.data(), static_cast<int>(nfft), true);

        float max_val = -1e30f;
        int max_idx = 0;
        for (std::size_t i = 0; i < nfft; ++i) {
            float v = c[i].r * c[i].r + c[i].i * c[i].i;
            if (v > max_val) { max_val = v; max_idx = static_cast<int>(i); }
        }
        if (max_idx > static_cast<int>(nfft) / 2) max_idx -= static_cast<int>(nfft);

        // Parabolic interpolation
        int idx_left = max_idx - 1;
        int idx_right = max_idx + 1;
        float y0 = c[(idx_left + static_cast<int>(nfft)) % static_cast<int>(nfft)].r;
        float y1 = c[(max_idx + static_cast<int>(nfft)) % static_cast<int>(nfft)].r;
        float y2 = c[(idx_right + static_cast<int>(nfft)) % static_cast<int>(nfft)].r;
        float denom = y0 - 2.0f * y1 + y2;
        float p = 0.0f;
        if (std::abs(denom) > 1e-6f) {
            p = 0.5f * (y0 - y2) / denom;
        }
        float delay_samples = max_idx + p;
        float delay_sec = delay_samples / static_cast<float>(sample_rate);
        // Note: GCC-PHAT IFFT peak at tau corresponds to ch2[n] = ch1[n - tau],
        // so tau is actually the negative of the physical delay of ch2 relative to ch1.
        return -delay_sec;
    }

    bool SolveDirection(const std::vector<float>& delays, float& azimuth, float& elevation) {
        int M = mic_config.num_mics;
        if (M < 2 || delays.size() != static_cast<std::size_t>(M * (M - 1) / 2)) {
            return false;
        }

        const float c = 343.0f;
        int pair_idx = 0;
        int pairs = M * (M - 1) / 2;
        Eigen::MatrixXf A(pairs, 3);
        Eigen::VectorXf b(pairs);

        for (int i = 0; i < M; ++i) {
            for (int j = i + 1; j < M; ++j) {
                Eigen::Vector3f d;
                d.x() = mic_config.positions[j][0] - mic_config.positions[i][0];
                d.y() = mic_config.positions[j][1] - mic_config.positions[i][1];
                d.z() = mic_config.positions[j][2] - mic_config.positions[i][2];
                A.row(pair_idx) = d.transpose();
                b(pair_idx) = -delays[pair_idx] * c;
                pair_idx++;
            }
        }

        // SVD least squares
        Eigen::JacobiSVD<Eigen::MatrixXf> svd(A, Eigen::ComputeThinU | Eigen::ComputeThinV);
        Eigen::Vector3f dir = svd.solve(b);
        if (dir.norm() < 1e-6f) return false;
        dir.normalize();

        // Azimuth: angle from y-axis (North) clockwise
        azimuth = std::atan2(dir.x(), dir.y()) * 180.0f / static_cast<float>(M_PI);
        if (azimuth < 0.0f) azimuth += 360.0f;

        // Elevation
        float horiz = std::sqrt(dir.x() * dir.x() + dir.y() * dir.y());
        elevation = std::atan2(dir.z(), horiz) * 180.0f / static_cast<float>(M_PI);
        return true;
    }
};

GccPhatLocalizer::GccPhatLocalizer(const MicArrayConfig& mic_config,
                                   int sample_rate,
                                   float max_tdoa,
                                   int interp_factor)
    : impl_(std::make_unique<Impl>(mic_config, sample_rate, max_tdoa, interp_factor)) {}

GccPhatLocalizer::~GccPhatLocalizer() = default;
GccPhatLocalizer::GccPhatLocalizer(GccPhatLocalizer&&) noexcept = default;
GccPhatLocalizer& GccPhatLocalizer::operator=(GccPhatLocalizer&&) noexcept = default;

std::pair<float, float> GccPhatLocalizer::Localize(const Eigen::MatrixXf& multi_channel_audio) {
    int M = impl_->mic_config.num_mics;
    if (M < 2 || multi_channel_audio.cols() < M) return {0.0f, 0.0f};
    std::size_t n = multi_channel_audio.rows();
    if (n == 0) return {0.0f, 0.0f};

    std::vector<float> delays;
    for (int i = 0; i < M; ++i) {
        for (int j = i + 1; j < M; ++j) {
            std::vector<float> ch1(n), ch2(n);
            for (std::size_t s = 0; s < n; ++s) {
                ch1[s] = multi_channel_audio(static_cast<int>(s), i);
                ch2[s] = multi_channel_audio(static_cast<int>(s), j);
            }
            float d = impl_->ComputeGccPhatDelay(ch1.data(), ch2.data(), n);
            if (std::abs(d) > impl_->max_tdoa) d = (d > 0 ? impl_->max_tdoa : -impl_->max_tdoa);
            delays.push_back(d);
        }
    }

    float azimuth = 0.0f, elevation = 0.0f;
    if (impl_->SolveDirection(delays, azimuth, elevation)) {
        return {azimuth, elevation};
    }
    return {0.0f, 0.0f};
}

std::pair<float, float> GccPhatLocalizer::Localize(const std::vector<float>& flat_samples,
                                                   std::size_t num_channels) {
    if (num_channels == 0) return {0.0f, 0.0f};
    std::size_t total = flat_samples.size();
    std::size_t samples_per_channel = total / num_channels;
    if (samples_per_channel == 0) return {0.0f, 0.0f};
    Eigen::MatrixXf mat(static_cast<int>(samples_per_channel), static_cast<int>(num_channels));
    for (std::size_t ch = 0; ch < num_channels; ++ch) {
        for (std::size_t i = 0; i < samples_per_channel; ++i) {
            mat(static_cast<int>(i), static_cast<int>(ch)) = flat_samples[i * num_channels + ch];
        }
    }
    return Localize(mat);
}

const MicArrayConfig& GccPhatLocalizer::MicConfig() const noexcept {
    return impl_->mic_config;
}

}  // namespace acoustic
