#include "gpr/ascan.hpp"

#include <algorithm>
#include <cmath>
#include <numeric>
#include <stdexcept>

#include "pocketfft_hdronly.h"

namespace gpr {

PairTimestamps compute_pair_timestamps(const MotionConfig& mot, const GeometryConfig& geom) {
    // interleaved_tx_by_frequency: слот Tx = порядок в отсортированном output_idx.
    std::vector<int> tx = geom.output_idx, rx = geom.input_idx;
    std::sort(tx.begin(), tx.end());
    std::sort(rx.begin(), rx.end());
    int n_slots = (int)tx.size();
    double frame_center = 0.5 * n_slots * mot.tx_sweep_time_s;
    PairTimestamps ts;
    for (int slot = 0; slot < (int)tx.size(); ++slot) {
        int i_tx = slot;  // индекс в отсортированном == логический
        for (int j = 0; j < (int)rx.size(); ++j) {
            ts[{i_tx, j}] = PairTiming{slot, n_slots, frame_center};
        }
    }
    return ts;
}

static double median_of_diff(const std::vector<double>& f) {
    std::vector<double> d(f.size() - 1);
    for (size_t i = 0; i + 1 < f.size(); ++i) d[i] = f[i + 1] - f[i];
    std::sort(d.begin(), d.end());
    size_t n = d.size();
    return (n % 2) ? d[n / 2] : 0.5 * (d[n / 2 - 1] + d[n / 2]);
}

std::vector<c128> apply_intra_sweep_phase_correction(
    const std::vector<c128>& s21, const std::vector<double>& freq_full,
    const PairTiming& info, const MotionConfig& cfg, double wave_speed,
    double f_start, double f_stop, int n_slots_default) {
    std::vector<c128> out = s21;
    int n_full = (int)freq_full.size();
    // индексы в полосе
    std::vector<int> idx;
    for (int i = 0; i < n_full; ++i)
        if (freq_full[i] >= f_start && freq_full[i] <= f_stop) idx.push_back(i);
    if (idx.empty()) throw std::runtime_error("phase corr: пустая полоса");
    if (n_full < 2) return out;

    int n_slots = info.tx_interleave_slots ? info.tx_interleave_slots : n_slots_default;
    int slot = info.tx_interleave_slot;
    double dt_base = cfg.tx_sweep_time_s / (n_full - 1);
    // t_all по всему свипу; t_center — середина ПОЛНОГО свипа (как в preprocess.py)
    double t_first = (double)(n_slots * 0 + slot) * dt_base;
    double t_last = (double)(n_slots * (n_full - 1) + slot) * dt_base;
    double t_center = 0.5 * (t_first + t_last);
    double theta = cfg.look_angle_deg * M_PI / 180.0;
    if (!cfg.apply_freq_phase_correction) return out;
    for (int i : idx) {
        double t_abs = (double)(n_slots * i + slot) * dt_base;
        double dt_intra = t_abs - t_center;
        double delta_range = cfg.direction_sign * cfg.speed_m_s * std::cos(theta) * dt_intra;
        double dtau = 2.0 * delta_range / wave_speed;
        double phi = 2.0 * M_PI * freq_full[i] * dtau;
        double ang = cfg.motion_phase_sign * phi;
        out[i] *= c128(std::cos(ang), std::sin(ang));
    }
    return out;
}

static int next_pow2(int x) {
    int p = 1;
    while (p < x) p <<= 1;
    return p;
}

void compute_ascan_bp(const std::vector<c128>& s21, const std::vector<double>& freq,
                      double f_start, double f_stop, bool window, int oversample,
                      std::vector<double>& t_out, std::vector<c128>& h_out) {
    std::vector<double> freq_cut;
    std::vector<c128> s_cut;
    for (size_t i = 0; i < freq.size(); ++i)
        if (freq[i] >= f_start && freq[i] <= f_stop) {
            freq_cut.push_back(freq[i]);
            s_cut.push_back(s21[i]);
        }
    if (freq_cut.size() < 2) throw std::runtime_error("ascan: <2 точек в полосе");

    double df = median_of_diff(freq_cut);
    int n = (int)freq_cut.size();
    int k0 = (int)std::nearbyint(freq_cut[0] / df);  // round-half-to-even как python round
    int min_len = 2 * (k0 + n - 1);
    int n_fft_base = next_pow2(min_len);
    int n_fft = n_fft_base * oversample;
    double dt = 1.0 / ((double)n_fft * df);

    // окно Ханна (numpy np.hanning: 0.5 - 0.5*cos(2*pi*i/(N-1)))
    std::vector<c128> H(n_fft, c128(0.0, 0.0));
    for (int i = 0; i < n; ++i) {
        c128 v = s_cut[i];
        if (window) {
            double w = 0.5 - 0.5 * std::cos(2.0 * M_PI * i / (n - 1));
            v *= w;
        }
        H[k0 + i] = v;
    }

    // IFFT через pocketfft (backward + 1/n_fft == numpy ifft)
    h_out.assign(n_fft, c128(0.0, 0.0));
    pocketfft::shape_t shape{(size_t)n_fft};
    pocketfft::stride_t stride{(ptrdiff_t)sizeof(c128)};
    pocketfft::c2c(shape, stride, stride, {0}, pocketfft::BACKWARD,
                   H.data(), h_out.data(), (double)(1.0 / n_fft));

    t_out.resize(n_fft);
    for (int i = 0; i < n_fft; ++i) t_out[i] = i * dt;
}

double pair_norm_scale(const std::vector<c128>& h, const std::vector<double>& t,
                       double wave_speed, const BPConfig& bp) {
    std::vector<double> vals;
    vals.reserve(h.size());
    for (size_t i = 0; i < h.size(); ++i) {
        double z = t[i] * wave_speed / 2.0;
        if (z >= bp.pair_norm_depth_min && z <= bp.pair_norm_depth_max)
            vals.push_back(std::abs(h[i]));
    }
    if (vals.empty())
        for (auto& c : h) vals.push_back(std::abs(c));
    std::sort(vals.begin(), vals.end());
    // numpy.percentile linear interpolation
    double q = bp.pair_norm_percentile / 100.0;
    double pos = q * (vals.size() - 1);
    int lo = (int)std::floor(pos);
    int hi = std::min(lo + 1, (int)vals.size() - 1);
    double scale = vals[lo] + (pos - lo) * (vals[hi] - vals[lo]);
    if (!std::isfinite(scale) || scale <= bp.pair_norm_eps) scale = 1.0;
    return scale;
}

AScanSet preprocess_snapshot(
    const std::map<PairKey, std::vector<c128>>& s21,
    const std::map<PairKey, std::vector<double>>& freq,
    const std::map<PairKey, std::vector<c128>>* background,
    const MotionConfig& mot, const GeometryConfig& geom, const BPConfig& bp,
    const PairTimestamps& ts, bool normalize) {
    AScanSet out;
    out.pairs = geom.pair_keys();
    out.n_pairs = (int)out.pairs.size();

    std::vector<std::vector<c128>> Hs;
    std::vector<std::vector<double>> Ts;
    for (auto& key : out.pairs) {
        auto its = s21.find(key);
        if (its == s21.end()) throw std::runtime_error("preprocess: нет пары в s21");
        std::vector<c128> s = its->second;  // копия
        if (background) {
            auto ib = background->find(key);
            if (ib != background->end())
                for (size_t i = 0; i < s.size(); ++i) s[i] -= ib->second[i];
        }
        const std::vector<double>& fr = freq.at(key);
        auto itt = ts.find(key);
        if (itt != ts.end())
            s = apply_intra_sweep_phase_correction(s, fr, itt->second, mot, geom.v,
                                                   bp.f_start, bp.f_stop, geom.n_tx);
        std::vector<double> t;
        std::vector<c128> h;
        compute_ascan_bp(s, fr, bp.f_start, bp.f_stop, bp.window, bp.oversample, t, h);
        Hs.push_back(std::move(h));
        Ts.push_back(std::move(t));
    }

    // общая ось времени — проверка
    int n_fft = (int)Ts[0].size();
    double t0 = Ts[0][0];
    double dt = (Ts[0].back() - t0) / (n_fft - 1);
    for (auto& t : Ts)
        if ((int)t.size() != n_fft || std::abs(t[0] - t0) > 1e-18)
            throw std::runtime_error("A-сканы: разная ось времени у пар");

    if (normalize) {
        for (int p = 0; p < out.n_pairs; ++p) {
            double scale = pair_norm_scale(Hs[p], Ts[p], geom.v, bp);
            for (auto& c : Hs[p]) c /= (scale + bp.pair_norm_eps);
        }
    }

    out.n_fft = n_fft;
    out.t0 = t0;
    out.dt = dt;
    out.H.resize((size_t)out.n_pairs * n_fft);
    for (int p = 0; p < out.n_pairs; ++p)
        std::copy(Hs[p].begin(), Hs[p].end(), out.H.begin() + (size_t)p * n_fft);
    return out;
}

}  // namespace gpr