// 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.
double biquad_step(double x, double b0, double b1, double b2, double a1, double a2, double& z1,
                   double& z2) {
  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";
  }
  return "?";
}

int math_op_count() { return 12; }

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;
    default: return MathOp::Div;
  }
}

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; }

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;
  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 (!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;
      }
    }
    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