#include "gpr/ascan.hpp"

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

#include "pocketfft_hdronly.h"

namespace gpr {

TxWindows uniform_tx_windows(const std::vector<int>& tx_order, double acq_s, int tx_per_segment) {
    TxWindows windows;
    if (tx_order.empty()) return windows;
    int per_seg = std::max(1, tx_per_segment);
    int n_seg = std::max(1, (int)((tx_order.size() + per_seg - 1) / per_seg));
    double t_seg = acq_s / n_seg;
    for (size_t k = 0; k < tx_order.size(); ++k) {
        int seg = (int)k / per_seg;
        windows[tx_order[k]] = {seg * t_seg, seg * t_seg + t_seg};
    }
    return windows;
}

PairTimestamps compute_pair_timestamps(const TxWindows& tx_windows_s, double acq_s,
                                       const std::vector<int>& tx_order,
                                       const GeometryConfig& geom, const MotionConfig& mot) {
    // Аппаратный порядок: сперва tx_order (run.combos), затем прочие ключи окон.
    std::vector<int> hw_order;
    for (int tx : tx_order)
        if (tx_windows_s.count(tx) && std::find(hw_order.begin(), hw_order.end(), tx) == hw_order.end())
            hw_order.push_back(tx);
    for (auto& [tx, _] : tx_windows_s)
        if (std::find(hw_order.begin(), hw_order.end(), tx) == hw_order.end())
            hw_order.push_back(tx);

    // Tx с одинаковым окном — один сегмент; сегменты по возрастанию окна,
    // слот — порядок появления в hw_order (порт группировки ноутбука).
    std::vector<std::pair<std::pair<double, double>, std::vector<int>>> groups;
    for (int tx : hw_order) {
        auto win = tx_windows_s.at(tx);
        auto found = std::find_if(groups.begin(), groups.end(),
                                  [&](const auto& g) { return g.first == win; });
        if (found == groups.end())
            groups.push_back({win, {tx}});
        else
            found->second.push_back(tx);
    }
    std::sort(groups.begin(), groups.end(),
              [](const auto& a, const auto& b) { return a.first < b.first; });

    struct SegSlot { double t0, t1; int slot, n_slots; };
    std::map<int, SegSlot> seg_of;
    for (auto& [win, txs] : groups)
        for (size_t slot = 0; slot < txs.size(); ++slot)
            seg_of[txs[slot]] = {win.first, win.second, (int)slot, (int)txs.size()};

    double cos_theta = std::cos(mot.look_angle_deg * M_PI / 180.0);
    double t_ref = 0.5 * acq_s;
    PairTimestamps ts;
    for (int i = 0; i < geom.n_tx; ++i) {
        auto it = seg_of.find(geom.output_idx[i]);
        if (it == seg_of.end())
            throw std::runtime_error("pair timestamps: нет окна для Tx");
        const SegSlot& s = it->second;
        PairTiming timing;
        timing.tx_interleave_slot = s.slot;
        timing.tx_interleave_slots = s.n_slots;
        timing.segment_start_s = s.t0;
        timing.segment_time_s = s.t1 - s.t0;
        timing.t_center_s = 0.5 * (s.t0 + s.t1);
        double dz = mot.direction_sign * mot.speed_m_s * (timing.t_center_s - t_ref) * cos_theta;
        timing.dtau_pair_s = 2.0 * dz / geom.v;
        for (int j = 0; j < geom.n_rx; ++j) ts[{i, j}] = timing;
    }
    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) {
    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;

    // Сегмент проходит все n_full частот за своё окно; на каждой частоте по
    // очереди работают его n_slots передатчиков, поэтому событие частоты i для
    // слота slot имеет номер (n_slots*i + slot) из (n_slots*n_full).
    int n_slots = std::max(1, info.tx_interleave_slots);
    int slot = info.tx_interleave_slot;
    double dt_event = info.segment_time_s / ((double)n_slots * n_full);
    double cos_theta = std::cos(cfg.look_angle_deg * M_PI / 180.0);
    double pair_dtau = cfg.apply_pair_motion_correction ? info.dtau_pair_s : 0.0;
    if (!cfg.apply_freq_phase_correction && pair_dtau == 0.0) return out;
    for (int i : idx) {
        double dtau_intra = 0.0;
        if (cfg.apply_freq_phase_correction) {
            double t_abs = info.segment_start_s + (double)(n_slots * i + slot) * dt_event;
            double dt_intra = t_abs - info.t_center_s;
            dtau_intra = 2.0 * cfg.direction_sign * cfg.speed_m_s * cos_theta * dt_intra / wave_speed;
        }
        double phi = 2.0 * M_PI * freq_full[i] * (dtau_intra + pair_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);
        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