// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2026 Martino Ferrari

#include "math_channel.hpp"

#include <algorithm>
#include <cmath>
#include <limits>

namespace udpscope {
namespace {

constexpr double kPi = 3.14159265358979323846;
/// A value the step can return without a library call: `std::nan("")` parses its argument, and this
/// runs once per stored slot per node.
constexpr double kNaN = std::numeric_limits<double>::quiet_NaN();

/// A biquad's coefficients for one node, from the RBJ audio-EQ cookbook. `q` is 1/sqrt(2)
/// (Butterworth) for the plain low/high pass, and the node's own for the band pass -- which is
/// built here as the constant-0-dB-peak form, because a derived trace wants the band itself and
/// not whatever gain its width happened to produce.
void biquad(MathOp op, double fc, double q, double dt, double& b0, double& b1, double& b2,
            double& a1, double& a2) {
  // Two Nyquist guards and a floor: a cutoff at or above Nyquist has no meaning to give, and one
  // at zero would give coefficients of zero -- a trace that is flat because the filter is broken
  // rather than because the signal is.
  const double nyquist = 0.5 / std::max(dt, 1e-12);
  fc = std::clamp(fc, nyquist * 1e-6, nyquist * 0.99);
  q = std::clamp(q, 1e-3, 100.0);
  const double w0 = 2.0 * kPi * fc * dt;
  const double cos_w0 = std::cos(w0);
  const double alpha = std::sin(w0) / (2.0 * q);
  const double a0 = 1.0 + alpha;
  switch (op) {
    case MathOp::LowPass:
      b0 = (1.0 - cos_w0) * 0.5;
      b1 = 1.0 - cos_w0;
      b2 = b0;
      break;
    case MathOp::HighPass:
      b0 = (1.0 + cos_w0) * 0.5;
      b1 = -(1.0 + cos_w0);
      b2 = b0;
      break;
    default:  // BandPass
      b0 = alpha;
      b1 = 0.0;
      b2 = -alpha;
      break;
  }
  a1 = -2.0 * cos_w0;
  a2 = 1.0 - alpha;
  b0 /= a0;
  b1 /= a0;
  b2 /= a0;
  a1 /= a0;
  a2 /= a0;
}

/// One biquad step, transposed direct form II -- the form whose state words are the two things
/// that actually need carrying between samples.
/// One pole's coefficient for a corner of `fc` Hz at `dt` seconds a sample: `a = 1 - exp(-2*pi*fc*dt)`,
/// clamped just under Nyquist so `a` cannot exceed 1. `fc == 0` returns 1, which is a pass-through -- so
/// "no pole" costs the step no branch.
double pole_a(double fc, double dt) {
  if (!(fc > 0.0) || !(dt > 0.0)) return 1.0;
  const double nyquist = 0.5 / dt;
  const double use = fc < nyquist * 0.99 ? fc : nyquist * 0.99;
  return 1.0 - std::exp(-2.0 * kPi * use * dt);
}

/// The median window's length: odd, and between 3 and 9.
///
/// Odd because an even window's median is the average of two samples, which is a value the signal never
/// took -- and the whole point of a median here is that its output is a sample that was really there.
/// Rounded *up*, so a request for 4 gets 5 rather than 3. The ceiling of 9 is what `math_step`'s fixed
/// `double sorted[9]` is sized for.
int median_n(int want) {
  int n = want < 3 ? 3 : (want > 9 ? 9 : want);
  if ((n & 1) == 0) ++n;
  return n;
}

double biquad_step(double x, double b0, double b1, double b2, double a1, double a2, double& z1,
                   double& z2) {
  // A reserved hole is NaN in the source, and a NaN integrated into the state words is
  // permanent: every later sample would come out NaN, so a filtered trace died at the first
  // gap and never came back -- which is what a dropped packet did to every low pass in
  // relative mode, where a counter gap reserves NaN slots by construction. A gap in the input
  // is a gap in the output (the hole stays visible), and the filter's history is the last
  // *real* samples it saw, so the state is left alone and the trace resumes after the hole.
  if (!std::isfinite(x)) return kNaN;
  const double y = b0 * x + z1;
  z1 = b1 * x - a1 * y + z2;
  z2 = b2 * x - a2 * y;
  return y;
}

}  // namespace

const char* math_op_name(MathOp op) {
  switch (op) {
    case MathOp::Source: return "channel";
    case MathOp::Gain: return "gain + offset";
    case MathOp::Abs: return "abs";
    case MathOp::Sqrt: return "sqrt";
    case MathOp::Clamp: return "clamp";
    case MathOp::LowPass: return "low pass";
    case MathOp::HighPass: return "high pass";
    case MathOp::BandPass: return "band pass";
    case MathOp::Sum: return "sum";
    case MathOp::Sub: return "difference";
    case MathOp::Mul: return "product";
    case MathOp::Div: return "ratio";
    case MathOp::DeNoise: return "de-noise";
  }
  return "?";
}

int math_op_count() { return 13; }

MathOp math_op_at(int index) {
  switch (index) {
    case 0: return MathOp::Source;
    case 1: return MathOp::Gain;
    case 2: return MathOp::Abs;
    case 3: return MathOp::Sqrt;
    case 4: return MathOp::Clamp;
    case 5: return MathOp::LowPass;
    case 6: return MathOp::HighPass;
    case 7: return MathOp::BandPass;
    case 8: return MathOp::Sum;
    case 9: return MathOp::Sub;
    case 10: return MathOp::Mul;
    case 11: return MathOp::Div;
    default: return MathOp::DeNoise;
  }
}

int math_op_index(MathOp op) {
  for (int i = 0; i < math_op_count(); ++i) {
    if (math_op_at(i) == op) return i;
  }
  return 0;
}

int math_op_arity(MathOp op) {
  switch (op) {
    case MathOp::Source: return 0;
    case MathOp::Sum:
    case MathOp::Sub:
    case MathOp::Mul:
    case MathOp::Div: return 2;
    default: return 1;
  }
}

bool math_op_uses_gain(MathOp op) { return op == MathOp::Gain; }
bool math_op_uses_filter(MathOp op) {
  return op == MathOp::LowPass || op == MathOp::HighPass || op == MathOp::BandPass;
}
bool math_op_uses_clamp(MathOp op) { return op == MathOp::Clamp; }
bool math_op_uses_fc(MathOp op) { return math_op_uses_filter(op) || op == MathOp::DeNoise; }
bool math_op_uses_median(MathOp op) { return op == MathOp::DeNoise; }

void math_prepare(const MathChannel& channel, MathRuntime& rt, double dt,
                  const std::function<int(const std::string&)>& index_of) {
  const std::size_t n = channel.nodes.size();
  // A rate change *keeps* the filter state: the state words are in the signal's units, not in
  // time, so a decimation change moves the cutoffs (which is why the coefficients are rebuilt)
  // without making the trace jump. A node that is new gets zeros, which is a filter starting from
  // rest -- the only honest starting point there is.
  const std::vector<MathNodeRt> was = std::move(rt.node);
  rt.node.assign(n, MathNodeRt{});
  for (std::size_t i = 0; i < n && i < was.size(); ++i) {
    rt.node[i].z1 = was[i].z1;
    rt.node[i].z2 = was[i].z2;
  }
  rt.dt = dt;
  rt.ready = true;
  // The median windows, pooled: one pass to size the pool, then the node loop hands out the slices. A
  // window starts from rest after a re-prepare -- the same choice the biquads' `z2` makes -- while the
  // pole's `z1` is carried, so a decimation change does not make the trace jump.
  std::size_t med_total = 0;
  for (std::size_t i = 0; i < n; ++i) {
    if (channel.nodes[i].op == MathOp::DeNoise) {
      med_total += static_cast<std::size_t>(median_n(channel.nodes[i].med_n));
    }
  }
  rt.median.assign(med_total, 0.0);
  std::size_t med_next = 0;
  for (std::size_t i = 0; i < n; ++i) {
    const MathNode& node = channel.nodes[i];
    MathNodeRt& out = rt.node[i];
    if (node.op == MathOp::Source) {
      out.src = index_of ? index_of(node.source) : -1;
      continue;
    }
    if (node.op == MathOp::DeNoise) {
      out.med_n = median_n(node.med_n);
      out.med_at = med_next;
      out.med_pos = 0;
      out.med_have = 0;
      med_next += static_cast<std::size_t>(out.med_n);
      out.b0 = pole_a(node.fc, dt);
      continue;
    }
    if (!math_op_uses_filter(node.op)) continue;
    biquad(node.op, node.fc, node.q, dt, out.b0, out.b1, out.b2, out.a1, out.a2);
  }
}

double math_step(const MathChannel& channel, MathRuntime& rt, const float* values,
                 std::size_t count) {
  const std::size_t n = channel.nodes.size();
  if (n == 0 || rt.node.size() != n) return kNaN;
  MathNodeRt* rts = rt.node.data();
  const auto at = [&](int index) -> double {
    if (index < 0 || static_cast<std::size_t>(index) >= n) return kNaN;
    return rts[static_cast<std::size_t>(index)].value;
  };
  for (std::size_t i = 0; i < n; ++i) {
    const MathNode& node = channel.nodes[i];
    MathNodeRt& self = rts[i];
    double x = 0.0;
    switch (node.op) {
      case MathOp::Source: {
        const int c = self.src;
        x = c >= 0 && static_cast<std::size_t>(c) < count
                ? static_cast<double>(values[static_cast<std::size_t>(c)])
                : kNaN;
        break;
      }
      case MathOp::Gain:
        x = at(node.in_a) * node.gain + node.offset;
        break;
      case MathOp::Abs:
        x = std::fabs(at(node.in_a));
        break;
      case MathOp::Sqrt:
        // Of the magnitude: the square root of a negative sample has no real answer, and a derived
        // trace that turns into NaN for half a cycle is not a useful one.
        x = std::sqrt(std::fabs(at(node.in_a)));
        break;
      case MathOp::Clamp:
        x = std::clamp(at(node.in_a), std::min(node.lo, node.hi), std::max(node.lo, node.hi));
        break;
      case MathOp::LowPass:
      case MathOp::HighPass:
      case MathOp::BandPass:
        x = biquad_step(at(node.in_a), self.b0, self.b1, self.b2, self.a1, self.a2, self.z1,
                        self.z2);
        break;
      case MathOp::Sum:
        x = at(node.in_a) + at(node.in_b);
        break;
      case MathOp::Sub:
        x = at(node.in_a) - at(node.in_b);
        break;
      case MathOp::Mul:
        x = at(node.in_a) * at(node.in_b);
        break;
      case MathOp::Div: {
        const double d = at(node.in_b);
        // A scope trace must not go to infinity or NaN: a division by zero reads as zero, which is
        // visibly wrong on the plot instead of invisibly wrong in the number.
        x = std::fabs(d) > 0.0 ? at(node.in_a) / d : 0.0;
        break;
      }
      case MathOp::DeNoise: {
        const double in = at(node.in_a);
        // A hole stays a hole, and the window is left alone: folding a NaN into the median would smear one
        // gap across `med_n` samples, and the whole trace after it in the pole. `med_n <= 0` is a node the
        // runtime has not prepared -- the window pointer below would be a read past the pool, so it is
        // checked rather than assumed.
        if (!std::isfinite(in) || self.med_n <= 0) {
          x = kNaN;
          break;
        }
        const std::size_t len = static_cast<std::size_t>(self.med_n);
        double* win = rt.median.data() + self.med_at;
        win[self.med_pos] = in;
        self.med_pos = self.med_pos + 1 == len ? 0 : self.med_pos + 1;
        if (self.med_have < len) ++self.med_have;
        // An insertion sort of at most nine doubles into a local. The window is tiny and fixed, so this
        // beats anything that allocates, and it leaves the ring itself in arrival order -- which is what
        // lets the oldest sample be overwritten in place. Before the window has filled, the median is taken
        // over what there is: a shorter window, deliberately, rather than NaN. A derived trace that reads
        // "no data" for its first few samples after every rebuild looks like a broken graph.
        double sorted[9];
        const std::size_t m = self.med_have;
        for (std::size_t j = 0; j < m; ++j) {
          const double v = win[j];
          std::size_t p = j;
          while (p > 0 && sorted[p - 1] > v) {
            sorted[p] = sorted[p - 1];
            --p;
          }
          sorted[p] = v;
        }
        const double med = sorted[m / 2];
        // Median *then* pole, and the order is the design. The median kills an impulse outright; the pole
        // then takes the broadband hash it leaves behind. The other way round feeds the pole the impulse,
        // which smears it over the pole's whole settling time and leaves the median a wide bump it cannot
        // reject.
        self.z1 += self.b0 * (med - self.z1);
        x = self.z1;
        break;
      }
    }
    self.value = x;
  }
  const int out = channel.output;
  if (out < 0 || static_cast<std::size_t>(out) >= n) return kNaN;
  return rts[static_cast<std::size_t>(out)].value;
}

bool math_channel_valid_nodes(const std::vector<MathNode>& nodes, int output, std::string* why) {
  const std::size_t n = nodes.size();
  const auto fail = [&](const std::string& text) {
    if (why != nullptr) *why = text;
    return false;
  };
  if (n == 0) return fail("empty graph");
  for (std::size_t i = 0; i < n; ++i) {
    const MathNode& node = nodes[i];
    if (node.op == MathOp::Source && node.source.empty()) {
      return fail("node " + std::to_string(i + 1) + ": no channel chosen");
    }
    const int arity = math_op_arity(node.op);
    const int ins[2] = {node.in_a, node.in_b};
    for (int k = 0; k < arity; ++k) {
      // Only nodes *before* this one: the evaluator walks the list in order, so a forward or
      // self reference is not a cycle to detect, it is a node reading a value that is not there.
      if (ins[k] < 0 || ins[k] >= static_cast<int>(i)) {
        return fail("node " + std::to_string(i + 1) + ": input " + std::to_string(k + 1) +
                    " is not connected");
      }
    }
  }
  if (output < 0 || static_cast<std::size_t>(output) >= n) {
    return fail("no output chosen");
  }
  if (why != nullptr) why->clear();
  return true;
}

bool math_channel_valid(const MathChannel& channel, std::string* why) {
  return math_channel_valid_nodes(channel.nodes, channel.output, why);
}

}  // namespace udpscope