//
// Created by david on 04.03.2026.
//

#include "include/pd_signal.h"
#include <stdexcept>
#include <algorithm>
#include <iostream>
#include <numeric>

namespace pd_signal {

static constexpr double kPi = 3.14159265358979323846;

static void throwIfNotAscending(std::vector<double>& xp) {
    size_t N = xp.size();
    for(int i = 0; i < N - 1; i++)
        if((xp[i+1] - xp[i]) <= 0.0)
            throw std::invalid_argument("xp is not strictly monotonically ascending");
}

static int binarySearch(std::vector<double>& data, double x) {
    auto p = std::equal_range(data.begin(), data.end(), x);
    if(p.first == data.end()) return -((int) data.size()) - 1; // out of range on upper end
    int insertion_point = (int) std::distance(data.begin(), p.first);
    if(p.first != p.second) return insertion_point; // element found
    return -(insertion_point) - 1; // no element found directly
}

void linspace(std::vector<double>& data, double start, double stop, int num) {
    linspace(data, start, stop, num, true);
}

// `num` evenly spaced numbers over interval [start,stop] with endpoint=true or [start,stop) with endpoint=false
void linspace(std::vector<double>& data, double start, double stop, int num, bool endpoint) {
    if(num < 0) throw std::invalid_argument("num must be >= 0");

    int end = endpoint ? num : (num-1);
    double step = (stop - start) / (double) end;

    double d = start;
    data.resize(num);

    for(int i = 0; i < num; i++) {
        data[i] = d;
        d += step;
    }
}

// Evaluate at points x the function given by the samples fp[xp[n]].
void interp(std::vector<double>& y, std::vector<double>& x, std::vector<double>& xp, std::vector<double>& fp) {
    if (xp.size() != fp.size()) throw std::invalid_argument("xp.size() != fp.size()");
    size_t N = xp.size();
    size_t M = x.size();
    if (N < 2) throw std::invalid_argument("N < 2");
    throwIfNotAscending(xp);

    std::vector<double> slope, intercept;
    y.resize(M);
    slope.resize(N-1);
    intercept.resize(N-1);

    for(int i = 0; i < N-1; i++) {
        double dxp = xp[i+1] - xp[i];
        double dfp = fp[i+1] - fp[i];
        slope[i] = dfp / dxp;
        intercept[i] = fp[i] - slope[i] * xp[i];
    }

    for(int i = 0; i < M; i++) {
        if(x[i] < xp[0]) {
            y[i] = fp[0];
        } else if(x[i] > xp[N-1]) {
            y[i] = fp[N-1];
        } else {
            int idx = binarySearch(xp, x[i]);
            if (idx <= -1) {
                // not found precisely. interpolate
                idx = -idx - 1; // convert back to array index (insertion point)
                // idx==0 or idx==M cannot happen, handled above.
                //assert(idx > 0 && idx < N);
                idx = idx - 1; // one below insertion point (left point)
                y[i] = slope[idx] * x[i] + intercept[idx];
            } else {
                // found precisely
                y[i] = fp[idx];
            }
        }
    }
}

// resample to BEAT_LEN
void resample(std::vector<double> &out, std::vector<double> &x, int beat_len) {
    std::vector<double> t;
    std::vector<double> i;
    linspace(t, 0, (double) (x.size()-1), beat_len, false);
    linspace(i, 0, (int) (x.size()-1), (int) x.size(), false);
    interp(out, t, i, x);
}

// normalized cross-correlation of the two signals of same length
double crossCorr(std::vector<double> &x, std::vector<double> &y) {
    if (x.size() != y.size()) throw std::invalid_argument("x.size() != y.size()");
    double xs = 0.0, ys = 0.0, cs = 0.0;
    for (size_t i = 0; i < x.size(); i++) {
        xs += x[i] * x[i];
        ys += y[i] * y[i];
        cs += x[i] * y[i];
    }
    return cs / sqrt(xs * ys);
}

// clip 'val' to between 'a_min' and 'a_max'.
double clip(double val, double a_min, double a_max) {
    return std::min(std::max(val, a_min), a_max);
}

// two-dimensional mean of a collection of signals
template<class T> void mean_tpl(std::vector<double> &out, T& m) {
    if (m.empty()) {
        out.resize(0);
        return;
    }
    const size_t sz = m[0].size();
    out.resize(sz);
    out.assign(sz, 0.0);
    const size_t N = m.size();
    for (size_t i = 0; i < N; i++) {
        for (size_t j = 0; j < sz; j++) {
            out[j] += m[i][j];
        }
    }
    for (size_t j = 0; j < sz; j++) {
        out[j] /= static_cast<double>(N);
    }
}

void mean(std::vector<double> &out, std::vector<std::vector<double> >& m) {
    mean_tpl(out, m);
}
void mean(std::vector<double> &out, std::deque<std::vector<double> >& m) {
    mean_tpl(out, m);
}

double mean(const std::vector<double> &in) {
    if (in.empty()) {
        throw std::invalid_argument("mean: input vector is empty");
    }
    double sum = std::accumulate(in.begin(), in.end(), 0.0);
    return sum / static_cast<double>(in.size());
}

void diff(std::vector<double>& out, const std::vector<double>& in) {
    if (in.size() < 2) {
        out.clear();
        return;
    }
    out.resize(in.size() - 1);
    for (std::size_t i = 1; i < in.size(); ++i) {
        out[i - 1] = in[i] - in[i - 1];
    }
}

// Convolution of two polynomials in ascending power order.
void polymul(std::vector<cplx>& out,
             const std::vector<cplx>& p, const std::vector<cplx>& q) {
    if (p.empty() || q.empty()) {
        out.clear();
        return;
    }
    out.assign(p.size() + q.size() - 1, cplx(0.0, 0.0));
    for (size_t i = 0; i < p.size(); ++i)
        for (size_t j = 0; j < q.size(); ++j)
            out[i + j] += p[i] * q[j];
}

// Build a monic polynomial from its roots, returned in descending power order.
void poly(std::vector<cplx>& out, const std::vector<cplx>& roots) {
    // Accumulate the product (x - r_0)(x - r_1)...(x - r_{N-1}) in ascending
    // power order, then reverse to descending order to match numpy.poly().
    std::vector<cplx> acc{cplx(1.0, 0.0)};
    std::vector<cplx> tmp;
    for (const cplx& root : roots) {
        const std::vector<cplx> factor{-root, cplx(1.0, 0.0)};  // -root + 1*x
        polymul(tmp, acc, factor);
        acc.swap(tmp);
    }
    out.assign(acc.rbegin(), acc.rend());
}

void iirHighpass(std::vector<double>& b, std::vector<double>& a,
                 int N, double fc, double fs) {
    if (N < 1) throw std::invalid_argument("N must be >= 1");
    if (!(fc > 0.0 && fc < 0.5 * fs))
        throw std::invalid_argument("require 0 < fc < fs/2");

    // 1) Analog Butterworth low-pass prototype poles (Omega_c = 1 rad/s):
    //    equally spaced on the unit circle in the left-half s-plane.
    // 2) Bilinear pre-warp so the digital cutoff lands exactly at fc:
    //        Omega_c = 2*fs*tan(pi*fc/fs)  =>  F_c = Omega_c / (2 pi).
    // 3) Scale prototype poles to the desired analog cutoff.
    // 4) Bilinear transform s -> z, with T = 1/fs:
    //        z = (1 + s T / 2) / (1 - s T / 2).
    const double Fc = (fs / kPi) * std::tan(kPi * fc / fs);
    const double T = 1.0 / fs;
    const double scale = 2.0 * kPi * Fc;

    std::vector<cplx> p_z(N);
    for (int k = 1; k <= N; ++k) {
        const double phase = kPi * (2 * k + N - 1) / (2.0 * N);
        const cplx p_proto(std::cos(phase), std::sin(phase));
        const cplx half_sT = (scale * T / 2.0) * p_proto;
        p_z[k - 1] = (1.0 + half_sT) / (1.0 - half_sT);
    }

    // 5) High-pass: place all N zeros at z = +1 so |H(z=1)| = 0 (DC null).
    const std::vector<cplx> z_z(N, cplx(1.0, 0.0));

    // 6) Build numerator and denominator polynomials (highest power of z first).
    //    Poles come in conjugate pairs, so the imaginary part is numerical noise.
    std::vector<cplx> a_c, b_c;
    poly(a_c, p_z);
    poly(b_c, z_z);

    a.resize(N + 1);
    b.resize(N + 1);
    for (int i = 0; i <= N; ++i) {
        a[i] = a_c[i].real();
        b[i] = b_c[i].real();
    }

    // 7) Normalize for unit gain at Nyquist (z = -1), the high-pass passband peak.
    //    Sum_i b[i] * (-1)^(N-i)  /  Sum_i a[i] * (-1)^(N-i).
    double num = 0.0, den = 0.0;
    double sign = ((N & 1) == 0) ? 1.0 : -1.0;  // (-1)^N for i=0
    for (int i = 0; i <= N; ++i) {
        num += b[i] * sign;
        den += a[i] * sign;
        sign = -sign;
    }
    const double gain = num / den;
    for (int i = 0; i <= N; ++i) b[i] /= gain;
}

}