#include "doa_algorithms.hpp"

#include <algorithm>
#include <cmath>
#include <complex>
#include <limits>
#include <stdexcept>
#include <utility>

namespace esp32_doa {
namespace {

constexpr double kPi = 3.141592653589793238462643383279502884;
constexpr double kTolerance = 1e-14;

std::size_t index(const Stft &s, int snapshot, int frequency, int channel) {
    return (static_cast<std::size_t>(snapshot) * s.bins + frequency) *
           s.channels + channel;
}

void fft(std::vector<Complex> &x, bool inverse) {
    const std::size_t n = x.size();
    if (n == 0 || (n & (n - 1)) != 0) {
        throw std::invalid_argument("nfft must be a power of two");
    }
    for (std::size_t i = 1, j = 0; i < n; ++i) {
        std::size_t bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) std::swap(x[i], x[j]);
    }
    for (std::size_t len = 2; len <= n; len <<= 1) {
        const double sign = inverse ? 1.0 : -1.0;
        const Complex root = std::polar(1.0, sign * 2.0 * kPi / len);
        for (std::size_t start = 0; start < n; start += len) {
            Complex w(1.0, 0.0);
            const std::size_t half = len >> 1;
            for (std::size_t j = 0; j < half; ++j) {
                const Complex u = x[start + j];
                const Complex v = x[start + j + half] * w;
                x[start + j] = u + v;
                x[start + j + half] = u - v;
                w *= root;
            }
        }
    }
    if (inverse) {
        for (Complex &v : x) v /= static_cast<double>(n);
    }
}

double hann(int n, int length) {
    if (length <= 1) return 1.0;
    return 0.5 - 0.5 * std::cos(2.0 * kPi * n / (length - 1));
}

std::vector<int> frequency_bins(const Config &c) {
    const int first = static_cast<int>(std::floor(
        c.min_frequency_hz * c.nfft / c.sample_rate + 0.5));
    const int last = static_cast<int>(std::floor(
        c.max_frequency_hz * c.nfft / c.sample_rate + 0.5));
    std::vector<int> bins;
    for (int k = first; k < last && k <= c.nfft / 2; ++k)
        if (k >= 0) bins.push_back(k);
    return bins;
}

void validate(const Config &c) {
    if (c.sample_rate <= 0 || c.channels <= 0 || c.nfft <= 0 ||
        (c.nfft & (c.nfft - 1)) != 0 || c.hop <= 0 || c.hop > c.nfft)
        throw std::invalid_argument("invalid STFT configuration");
    if (static_cast<int>(c.microphone_azimuths_deg.size()) != c.channels)
        throw std::invalid_argument("one microphone azimuth is required per channel");
    if (!(c.min_frequency_hz >= 0.0 && c.min_frequency_hz < c.max_frequency_hz &&
          c.max_frequency_hz < c.sample_rate / 2.0))
        throw std::invalid_argument("invalid frequency range");
}

std::vector<std::pair<double, double>> microphone_positions(const Config &c) {
    std::vector<std::pair<double, double>> p(c.channels);
    for (int m = 0; m < c.channels; ++m) {
        const double a = c.microphone_azimuths_deg[m] * kPi / 180.0;
        p[m] = {c.radius_m * std::cos(a), c.radius_m * std::sin(a)};
    }
    return p;
}

double wrap_degrees(double angle) {
    angle = std::fmod(angle, 360.0);
    return angle < 0.0 ? angle + 360.0 : angle;
}

std::vector<Complex> mode_vector(const Config &c, int frequency,
                                 double angle_degrees) {
    const auto positions = microphone_positions(c);
    const double a = angle_degrees * kPi / 180.0;
    const double dx = std::cos(a), dy = std::sin(a);
    const double omega = 2.0 * kPi * c.sample_rate * frequency / c.nfft;
    std::vector<Complex> mode(c.channels);
    for (int m = 0; m < c.channels; ++m) {
        const double tau = (dx * positions[m].first + dy * positions[m].second) /
                           c.sound_speed_m_s;
        mode[m] = std::polar(1.0, omega * tau);
    }
    return mode;
}

std::vector<Complex> solve_linear(std::vector<Complex> a, std::vector<Complex> b,
                                  int n) {
    for (int col = 0; col < n; ++col) {
        int pivot = col;
        double best = std::abs(a[col * n + col]);
        for (int row = col + 1; row < n; ++row) {
            const double value = std::abs(a[row * n + col]);
            if (value > best) best = value, pivot = row;
        }
        if (best < 1e-20) throw std::runtime_error("singular covariance matrix");
        if (pivot != col) {
            for (int j = col; j < n; ++j)
                std::swap(a[col * n + j], a[pivot * n + j]);
            std::swap(b[col], b[pivot]);
        }
        const Complex diagonal = a[col * n + col];
        for (int j = col; j < n; ++j) a[col * n + j] /= diagonal;
        b[col] /= diagonal;
        for (int row = 0; row < n; ++row) {
            if (row == col) continue;
            const Complex factor = a[row * n + col];
            if (std::abs(factor) == 0.0) continue;
            for (int j = col; j < n; ++j) a[row * n + j] -= factor * a[col * n + j];
            b[row] -= factor * b[col];
        }
    }
    return b;
}

struct EigenResult {
    std::vector<double> values;
    std::vector<Complex> vectors;  // column-major eigenvectors
};

// Complex Hermitian Jacobi eigensolver. This small implementation is enough
// for the fixed 8-microphone covariance matrices and avoids BLAS/LAPACK.
EigenResult hermitian_eigen(const std::vector<Complex> &input, int n) {
    std::vector<Complex> a = input;
    std::vector<Complex> v(static_cast<std::size_t>(n) * n, Complex(0.0));
    for (int i = 0; i < n; ++i) v[i * n + i] = 1.0;
    for (int iteration = 0; iteration < 80 * n * n; ++iteration) {
        int p = 0, q = 1;
        double largest = 0.0;
        for (int i = 0; i < n; ++i)
            for (int j = i + 1; j < n; ++j)
                if (std::abs(a[i * n + j]) > largest)
                    largest = std::abs(a[i * n + j]), p = i, q = j;
        if (largest < 1e-12) break;
        const Complex b = a[p * n + q];
        const double phase = std::arg(b);
        const double magnitude = std::abs(b);
        const double theta = 0.5 * std::atan2(2.0 * magnitude,
                                               a[q * n + q].real() - a[p * n + p].real());
        const double cs = std::cos(theta), sn = std::sin(theta);
        const Complex ph = std::polar(1.0, phase);
        const double app = a[p * n + p].real(), aqq = a[q * n + q].real();
        for (int k = 0; k < n; ++k) {
            if (k == p || k == q) continue;
            const Complex akp = a[k * n + p], akq = a[k * n + q];
            a[k * n + p] = akp * ph * cs - akq * sn;
            a[p * n + k] = std::conj(a[k * n + p]);
            a[k * n + q] = akp * ph * sn + akq * cs;
            a[q * n + k] = std::conj(a[k * n + q]);
        }
        a[p * n + p] = app * cs * cs + aqq * sn * sn - 2.0 * magnitude * cs * sn;
        a[q * n + q] = app * sn * sn + aqq * cs * cs + 2.0 * magnitude * cs * sn;
        a[p * n + q] = a[q * n + p] = Complex(0.0);
        for (int k = 0; k < n; ++k) {
            const Complex vkp = v[k * n + p], vkq = v[k * n + q];
            v[k * n + p] = vkp * ph * cs - vkq * sn;
            v[k * n + q] = vkp * ph * sn + vkq * cs;
        }
    }
    std::vector<int> order(n);
    for (int i = 0; i < n; ++i) order[i] = i;
    std::sort(order.begin(), order.end(), [&](int x, int y) {
        return a[x * n + x].real() < a[y * n + y].real();
    });
    EigenResult result;
    result.values.resize(n);
    result.vectors.resize(static_cast<std::size_t>(n) * n);
    for (int col = 0; col < n; ++col) {
        result.values[col] = a[order[col] * n + order[col]].real();
        for (int row = 0; row < n; ++row)
            result.vectors[row * n + col] = v[row * n + order[col]];
    }
    return result;
}

double spectrum_sharpness(const std::vector<double> &spectrum) {
    if (spectrum.empty()) return 0.0;
    std::vector<double> sorted = spectrum;
    std::sort(sorted.begin(), sorted.end());
    const double peak = sorted.back();
    const double floor = sorted[sorted.size() / 2];
    const double spread = peak - sorted.front();
    return std::clamp((peak - floor) / (spread + 1e-12), 0.0, 1.0);
}

DoaResult finish_spectrum(std::vector<double> spectrum) {
    auto peak = std::max_element(spectrum.begin(), spectrum.end());
    const int index = static_cast<int>(peak - spectrum.begin());
    DoaResult result;
    result.angle_degrees = wrap_degrees(360.0 * index / spectrum.size());
    result.peak_sharpness = spectrum_sharpness(spectrum);
    result.spectrum = std::move(spectrum);
    return result;
}

}  // namespace

Complex &Stft::at(int snapshot, int frequency, int channel) {
    return spectra[index(*this, snapshot, frequency, channel)];
}
const Complex &Stft::at(int snapshot, int frequency, int channel) const {
    return spectra[index(*this, snapshot, frequency, channel)];
}

Stft analyze_stft(const std::vector<double> &audio, int sample_count,
                 const Config &config) {
    validate(config);
    if (sample_count < config.nfft ||
        static_cast<std::size_t>(sample_count) * config.channels > audio.size())
        throw std::invalid_argument("audio does not contain enough samples");
    Stft result;
    result.sample_count = sample_count;
    result.snapshots = (sample_count - config.nfft) / config.hop + 1;
    result.bins = config.nfft / 2 + 1;
    result.channels = config.channels;
    result.spectra.assign(static_cast<std::size_t>(result.snapshots) * result.bins *
                              result.channels,
                          Complex(0.0));
    std::vector<double> means(config.channels, 0.0);
    for (int n = 0; n < sample_count; ++n)
        for (int m = 0; m < config.channels; ++m)
            means[m] += audio[static_cast<std::size_t>(n) * config.channels + m] /
                        sample_count;
    for (int s = 0; s < result.snapshots; ++s) {
        const int start = s * config.hop;
        for (int m = 0; m < config.channels; ++m) {
            std::vector<Complex> frame(config.nfft);
            for (int n = 0; n < config.nfft; ++n)
                frame[n] = (audio[static_cast<std::size_t>(start + n) * config.channels + m] -
                            means[m]) * hann(n, config.nfft);
            fft(frame, false);
            for (int k = 0; k < result.bins; ++k) result.at(s, k, m) = frame[k];
        }
    }
    return result;
}

DoaResult srp_phat(const Stft &stft, const Config &config) {
    validate(config);
    const std::vector<int> bins = frequency_bins(config);
    const int pairs = config.channels * (config.channels - 1) / 2;
    std::vector<double> response(config.grid_points, 0.0);
    const auto positions = microphone_positions(config);
    for (int g = 0; g < config.grid_points; ++g) {
        const double angle = 360.0 * g / config.grid_points * kPi / 180.0;
        const double dx = std::cos(angle), dy = std::sin(angle);
        double score = 0.0;
        for (int k : bins) {
            std::vector<std::vector<Complex>> cc(config.channels,
                                                  std::vector<Complex>(config.channels));
            for (int i = 0; i < config.channels; ++i)
                for (int j = 0; j < config.channels; ++j)
                    for (int s = 0; s < stft.snapshots; ++s) {
                        Complex xi = stft.at(s, k, i), xj = stft.at(s, k, j);
                        const double ai = std::max(std::abs(xi), kTolerance);
                        const double aj = std::max(std::abs(xj), kTolerance);
                        cc[i][j] += (xi / ai) * std::conj(xj / aj);
                    }
            const double omega = 2.0 * kPi * config.sample_rate * k / config.nfft;
            for (int i = 0; i < config.channels; ++i) {
                const double ti = (dx * positions[i].first + dy * positions[i].second) /
                                  config.sound_speed_m_s;
                for (int j = i + 1; j < config.channels; ++j) {
                    const double tj = (dx * positions[j].first + dy * positions[j].second) /
                                      config.sound_speed_m_s;
                    const Complex mode = std::polar(1.0, omega * (tj - ti));
                    score += 2.0 * std::real(cc[i][j] * mode);
                }
            }
        }
        const double dc = static_cast<double>(stft.snapshots * config.channels * bins.size());
        response[g] = (score + dc) /
                      (stft.snapshots * bins.size() * static_cast<double>(pairs));
    }
    return finish_spectrum(std::move(response));
}

DoaResult norm_music(const Stft &stft, const Config &config) {
    validate(config);
    const std::vector<int> bins = frequency_bins(config);
    std::vector<double> per_frequency(static_cast<std::size_t>(bins.size()) *
                                      config.grid_points, 0.0);
    for (std::size_t frequency_index = 0; frequency_index < bins.size(); ++frequency_index) {
        const int k = bins[frequency_index];
        std::vector<Complex> covariance(static_cast<std::size_t>(config.channels) * config.channels,
                                         Complex(0.0));
        for (int i = 0; i < config.channels; ++i)
            for (int j = 0; j < config.channels; ++j)
                for (int s = 0; s < stft.snapshots; ++s)
                    covariance[i * config.channels + j] +=
                        stft.at(s, k, i) * std::conj(stft.at(s, k, j)) / stft.snapshots;
        const EigenResult eig = hermitian_eigen(covariance, config.channels);
        for (int g = 0; g < config.grid_points; ++g) {
            const auto mode = mode_vector(config, k, 360.0 * g / config.grid_points);
            Complex denom(0.0);
            for (int i = 0; i < config.channels; ++i) {
                for (int j = 0; j < config.channels; ++j) {
                    Complex cross = (i == j ? Complex(1.0) : Complex(0.0));
                    cross -= eig.vectors[i * config.channels + config.channels - 1] *
                             std::conj(eig.vectors[j * config.channels + config.channels - 1]);
                    denom += std::conj(mode[i]) * cross * mode[j];
                }
            }
            per_frequency[frequency_index * config.grid_points + g] =
                1.0 / std::max(std::abs(denom), 1e-20);
        }
    }
    std::vector<double> response(config.grid_points, 0.0);
    // Pyroomacoustics' NormMUSIC normalizes along the frequency axis for
    // every candidate direction (Pssl / max(Pssl, axis=0)).
    for (int g = 0; g < config.grid_points; ++g) {
        double maximum = 0.0;
        for (std::size_t f = 0; f < bins.size(); ++f)
            maximum = std::max(maximum, per_frequency[f * config.grid_points + g]);
        for (std::size_t f = 0; f < bins.size(); ++f)
            response[g] += per_frequency[f * config.grid_points + g] /
                           std::max(maximum, 1e-20);
        response[g] /= bins.size();
    }
    return finish_spectrum(std::move(response));
}

MvdrResult mvdr(const Stft &stft, const Config &config, double angle_degrees) {
    validate(config);
    const std::vector<int> bins = frequency_bins(config);
    const int n = config.channels;
    std::vector<Complex> weights(static_cast<std::size_t>(stft.bins) * n);
    for (int k = 0; k < stft.bins; ++k) {
        const auto steering = mode_vector(config, k, angle_degrees);
        for (int m = 0; m < n; ++m) weights[k * n + m] = steering[m] / n;
    }
    for (int k : bins) {
        std::vector<Complex> covariance(static_cast<std::size_t>(n) * n, Complex(0.0));
        for (int i = 0; i < n; ++i)
            for (int j = 0; j < n; ++j)
                for (int s = 0; s < stft.snapshots; ++s)
                    covariance[i * n + j] += stft.at(s, k, i) * std::conj(stft.at(s, k, j)) /
                                              stft.snapshots;
        double power = 0.0;
        for (int i = 0; i < n; ++i) power += covariance[i * n + i].real();
        const double loading = config.diagonal_loading * std::max(power / n, 1e-12);
        for (int i = 0; i < n; ++i) covariance[i * n + i] += loading;
        std::vector<Complex> solved = solve_linear(covariance, steering, n);
        Complex denominator(0.0);
        for (int m = 0; m < n; ++m) denominator += std::conj(steering[m]) * solved[m];
        if (std::abs(denominator) < 1e-12) denominator = Complex(1e-12);
        for (int m = 0; m < n; ++m) weights[k * n + m] = solved[m] / denominator;
    }
    std::vector<double> output(stft.sample_count, 0.0), normalization(stft.sample_count, 0.0);
    for (int s = 0; s < stft.snapshots; ++s) {
        std::vector<Complex> spectrum(config.nfft, Complex(0.0));
        for (int k = 0; k < stft.bins; ++k) {
            Complex value(0.0);
            for (int m = 0; m < n; ++m) value += std::conj(weights[k * n + m]) * stft.at(s, k, m);
            spectrum[k] = value;
            if (k > 0 && k < config.nfft / 2) spectrum[config.nfft - k] = std::conj(value);
        }
        fft(spectrum, true);
        const int start = s * config.hop;
        for (int i = 0; i < config.nfft; ++i) {
            const double w = hann(i, config.nfft);
            output[start + i] += spectrum[i].real() * w;
            normalization[start + i] += w * w;
        }
    }
    for (int i = 0; i < stft.sample_count; ++i)
        output[i] = normalization[i] > 1e-12 ? output[i] / normalization[i] : 0.0;
    return {std::move(output)};
}

DoaResult estimate(const std::vector<double> &audio, int sample_count,
                   const Config &config, const std::string &algorithm) {
    const Stft stft = analyze_stft(audio, sample_count, config);
    if (algorithm == "SRP-PHAT" || algorithm == "SRP" || algorithm == "SRP-PATH")
        return srp_phat(stft, config);
    if (algorithm == "NormMUSIC" || algorithm == "MUSIC" || algorithm == "NORMMUSIC")
        return norm_music(stft, config);
    throw std::invalid_argument("unknown algorithm: " + algorithm);
}

MvdrResult beamform(const std::vector<double> &audio, int sample_count,
                    const Config &config, double angle_degrees) {
    return mvdr(analyze_stft(audio, sample_count, config), config, angle_degrees);
}

}  // namespace esp32_doa