// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! Finding a linear predictor for a block of samples, and reducing it to the
//! integers a subframe can carry.
//!
//! This is the one part of the library that computes in floating point. The
//! specification allows it: "Use of floating-point representations in analysis
//! (e.g., finding a good predictor or Rice parameter) is not a concern as long
//! as the process of using the found predictor and Rice parameter to encode
//! audio samples is implemented with only integer math" (RFC 9639,
//! Appendix A). Nothing here touches a sample that is written; it only
//! proposes coefficients, which `quantize` then turns into the integers the
//! subframe stores and the decoder multiplies by.

const std = @import("std");
const math = std.math;

const format = @import("format.zig");
const limits = format.limits;

/// A window is applied before the autocorrelation so that the block's edges,
/// where the predictor has the least to go on, count for less. A Tukey window
/// with this parameter is what the reference encoder uses by default, and it
/// is the shape the coefficients here are tuned against.
pub const default_tukey_parameter = 0.5;

/// Multiply `samples` by a Tukey window into `out`, which is a cosine taper
/// over the outermost `p` of the block and flat in between. `p` of 0 gives a
/// rectangular window and 1 a Hann window.
pub fn tukeyWindow(out: []f64, samples: []const i64, p: f64) void {
    std.debug.assert(out.len == samples.len);
    if (samples.len < 2) {
        for (out, samples) |*o, s| o.* = @floatFromInt(s);
        return;
    }
    if (!(p > 0)) {
        for (out, samples) |*o, s| o.* = @floatFromInt(s);
        return;
    }
    const n: f64 = @floatFromInt(samples.len - 1);
    const taper = @min(p, 1.0) / 2;
    for (out, samples, 0..) |*o, sample, i| {
        const x: f64 = @as(f64, @floatFromInt(i)) / n;
        const weight: f64 = if (x < taper)
            0.5 * (1 + @cos(math.pi * (x / taper - 1)))
        else if (x > 1 - taper)
            0.5 * (1 + @cos(math.pi * ((x - 1) / taper + 1)))
        else
            1;
        o.* = weight * @as(f64, @floatFromInt(sample));
    }
}

/// Autocorrelation of `windowed` at lags 0 through `out.len - 1`.
pub fn autocorrelation(out: []f64, windowed: []const f64) void {
    for (out, 0..) |*value, lag| {
        var sum: f64 = 0;
        if (lag < windowed.len) {
            for (windowed[lag..], 0..) |sample, i| sum += sample * windowed[i];
        }
        value.* = sum;
    }
}

/// What `levinson` found, for one predictor order.
pub const Candidate = struct {
    /// The coefficients, most recent sample first, before quantization.
    coefficients: []const f64,
    /// The prediction error the recursion expects at this order, which is what
    /// `expectedBitsPerResidual` turns into a size.
    err: f64,
};

/// Solve for the best predictor of every order from 1 to `max_order` at once,
/// which is what the Levinson-Durbin recursion gives for free: each step
/// extends the previous order's solution.
///
/// `coefficients` is a `max_order` by `max_order` matrix in row-major order,
/// row `i` holding the order `i + 1` solution; `errors` takes one entry per
/// order. `autoc` must hold `max_order + 1` lags. Returns however many orders
/// the recursion got through before the residual error stopped being positive,
/// which is fewer than asked for when the block is short or exactly periodic.
pub fn levinson(
    coefficients: []f64,
    errors: []f64,
    autoc: []const f64,
    max_order: usize,
) usize {
    std.debug.assert(autoc.len > max_order);
    std.debug.assert(coefficients.len >= max_order * max_order);
    std.debug.assert(errors.len >= max_order);

    var lpc: [limits.max_lpc_order]f64 = undefined;
    var err = autoc[0];
    if (!(err > 0)) return 0;

    for (0..max_order) |i| {
        // The reflection coefficient that makes order i+1 out of order i.
        var r = -autoc[i + 1];
        for (0..i) |j| r -= lpc[j] * autoc[i - j];
        r /= err;

        lpc[i] = r;
        var j: usize = 0;
        while (j < i / 2) : (j += 1) {
            const tmp = lpc[j];
            lpc[j] += r * lpc[i - 1 - j];
            lpc[i - 1 - j] += r * tmp;
        }
        if (i & 1 != 0) lpc[i / 2] += lpc[i / 2] * r;

        err *= 1 - r * r;
        // The prediction is subtracted from the sample, so the coefficients
        // the format stores are the negation of what the recursion carries.
        for (0..i + 1) |k| coefficients[i * max_order + k] = -lpc[k];
        errors[i] = err;
        if (!(err > 0)) return i + 1;
    }
    return max_order;
}

/// The bits per residual sample a predictor with this expected error should
/// produce over a block of `samples` samples. This is the usual estimate: the
/// residual is treated as Laplacian, whose entropy is half the log of its
/// variance, with the constant that makes it come out in bits.
pub fn expectedBitsPerResidual(err: f64, samples: usize) f64 {
    if (samples == 0) return 0;
    const scale = 0.5 * math.ln2 * math.ln2 / @as(f64, @floatFromInt(samples));
    const scaled = scale * err;
    if (!(scaled > 0)) return 0;
    const bits = 0.5 * @log(scaled) / math.ln2;
    return @max(bits, 0);
}

/// A predictor reduced to what a subframe can hold.
pub const Quantized = struct {
    /// `order` coefficients, each fitting in `precision` bits, signed.
    coefficients: [limits.max_lpc_order]i32 = @splat(0),
    order: usize = 0,
    precision: u5 = limits.min_qlp_precision,
    /// How far right to shift the sum of the products. The field is five bits
    /// signed and the format forbids a negative shift (RFC 9639, Section 9.2.6
    /// and Appendix B.4), so this is 0 to 15 -- four bits' worth, which is
    /// what its type says.
    shift: u4 = 0,
};

pub const QuantizeError = error{
    /// Every coefficient came out zero, or the largest of them is so large
    /// that no non-negative shift can hold it at this precision. Either way
    /// there is no predictor here and the caller should use another model.
    Unquantizable,
};

/// Turn real coefficients into integers of `precision` bits and a shift, so
/// that `sum(coefficient * sample) >> shift` approximates the prediction.
///
/// The error each rounding makes is carried into the next coefficient rather
/// than dropped, which keeps the quantized predictor closer to the one that
/// was solved for than rounding each coefficient on its own would.
pub fn quantize(coefficients: []const f64, precision: u5) QuantizeError!Quantized {
    std.debug.assert(precision >= limits.min_qlp_precision);
    std.debug.assert(precision <= limits.max_qlp_precision);
    std.debug.assert(coefficients.len <= limits.max_lpc_order);

    // One bit goes to the sign, so only the rest holds magnitude.
    const magnitude_bits: u5 = precision - 1;
    const qmax: i32 = (@as(i32, 1) << magnitude_bits) - 1;
    const qmin: i32 = -(@as(i32, 1) << magnitude_bits);

    var largest: f64 = 0;
    for (coefficients) |c| largest = @max(largest, @abs(c));
    if (!(largest > 0)) return error.Unquantizable;

    // Scale so the largest coefficient just fills the magnitude bits. The
    // shift is capped at the largest a five-bit signed field holds without
    // going negative, and a shift that would have to be negative means the
    // coefficients are too large to store. Coefficients small enough to want
    // more than the cap simply do not fill their bits; the predictor is a
    // little coarser and the stream is still valid.
    const log2_largest: i32 = @intFromFloat(@floor(@log2(largest)));
    const wanted: i32 = @as(i32, magnitude_bits) - log2_largest - 1;
    if (wanted < 0) return error.Unquantizable;
    const shift: u4 = @intCast(@min(wanted, limits.max_lpc_shift));

    var result: Quantized = .{
        .order = coefficients.len,
        .precision = precision,
        .shift = shift,
    };
    const scale = @as(f64, @floatFromInt(@as(i64, 1) << shift));
    var carried: f64 = 0;
    for (coefficients, 0..) |c, i| {
        carried += c * scale;
        const rounded: i32 = @intFromFloat(math.clamp(@round(carried), -2147483648.0, 2147483647.0));
        const clamped = math.clamp(rounded, qmin, qmax);
        carried -= @floatFromInt(clamped);
        result.coefficients[i] = clamped;
    }
    return result;
}

/// The coefficient precision the reference encoder would choose for a block of
/// this size and depth, which is a compromise between the space the
/// coefficients take and how well they describe the signal. Shorter blocks
/// cannot amortize wide coefficients, so they get fewer bits.
pub fn defaultPrecision(block_size: u32, bits_per_sample: u8) u5 {
    if (bits_per_sample < 16) {
        return @intCast(math.clamp(2 + bits_per_sample / 2, limits.min_qlp_precision, limits.max_qlp_precision));
    }
    if (bits_per_sample == 16) return switch (block_size) {
        0...192 => 8,
        193...384 => 10,
        385...576 => 12,
        577...1152 => 14,
        else => 15,
    };
    return switch (block_size) {
        0...384 => limits.max_qlp_precision - 2,
        385...1152 => limits.max_qlp_precision - 1,
        else => limits.max_qlp_precision,
    };
}

const testing = std.testing;

test "a Tukey window tapers the edges and leaves the middle alone" {
    var samples: [64]i64 = @splat(1000);
    var windowed: [64]f64 = undefined;
    tukeyWindow(&windowed, &samples, default_tukey_parameter);

    try testing.expectApproxEqAbs(@as(f64, 0), windowed[0], 1e-9);
    try testing.expectApproxEqAbs(@as(f64, 0), windowed[63], 1e-9);
    try testing.expectApproxEqAbs(@as(f64, 1000), windowed[32], 1e-9);
    // Symmetric about the middle.
    for (0..32) |i| try testing.expectApproxEqAbs(windowed[i], windowed[63 - i], 1e-9);
    // A rectangular window is the degenerate case and changes nothing.
    tukeyWindow(&windowed, &samples, 0);
    for (windowed) |w| try testing.expectApproxEqAbs(@as(f64, 1000), w, 1e-9);
}

test "the recursion recovers a predictor that generated the signal" {
    // x(n) = 1.6 x(n-1) - 0.8 x(n-2) driven by an impulse: a decaying
    // oscillation whose second-order predictor is exactly those two numbers.
    var samples: [512]i64 = undefined;
    var history: [2]f64 = .{ 0, 0 };
    for (&samples, 0..) |*sample, i| {
        const drive: f64 = if (i == 0) 10000 else 0;
        const value = 1.6 * history[0] - 0.8 * history[1] + drive;
        history[1] = history[0];
        history[0] = value;
        sample.* = @intFromFloat(@round(value));
    }

    var windowed: [512]f64 = undefined;
    tukeyWindow(&windowed, &samples, 0);
    var autoc: [9]f64 = undefined;
    autocorrelation(&autoc, &windowed);

    var coefficients: [64]f64 = undefined;
    var errors: [8]f64 = undefined;
    const orders = levinson(&coefficients, &errors, &autoc, 8);
    try testing.expect(orders >= 2);

    // Row 1 is the order-2 solution.
    try testing.expectApproxEqAbs(@as(f64, 1.6), coefficients[8], 0.01);
    try testing.expectApproxEqAbs(@as(f64, -0.8), coefficients[9], 0.01);
    // What is left over is the impulse that drove the signal, which no
    // predictor can foresee: its energy is 10000 squared, and orders past the
    // second buy nothing at all.
    try testing.expectApproxEqRel(@as(f64, 1e8), errors[1], 0.001);
    try testing.expect(errors[1] < errors[0] * 0.4);
    try testing.expect(errors[orders - 1] > errors[1] * 0.99);
}

test "quantization keeps the predictor it was given" {
    const coefficients = [_]f64{ 1.6, -0.8 };
    const q = try quantize(&coefficients, 15);
    try testing.expectEqual(@as(usize, 2), q.order);

    const scale = @as(f64, @floatFromInt(@as(i64, 1) << q.shift));
    for (coefficients, q.coefficients[0..2]) |wanted, got| {
        try testing.expectApproxEqAbs(wanted, @as(f64, @floatFromInt(got)) / scale, 0.001);
    }
}

test "coefficients that no shift can hold are refused" {
    // All zeroes is not a predictor.
    try testing.expectError(error.Unquantizable, quantize(&.{ 0, 0, 0 }, 15));
    // A coefficient this large needs more magnitude bits than the precision
    // has, which would mean a negative shift, which the format forbids.
    try testing.expectError(error.Unquantizable, quantize(&.{1e6}, 4));
}

test "small coefficients get the largest shift the format allows, not more" {
    // At fifteen bits of precision a largest coefficient of about 2^-10 would
    // want a shift of 23 to fill them. Five signed bits hold 15 at most, and 16
    // or more would be written as a negative shift (RFC 9639, Section 9.2.6).
    // Encoding 24-bit audio at -8 found this: a quiet, smooth signal gave
    // exactly such a predictor, and the stream could not be decoded.
    const q = try quantize(&.{ 0.001, -0.0007, 0.0003 }, 15);
    try testing.expectEqual(@as(u4, limits.max_lpc_shift), q.shift);
    // Coarser than it wanted to be, but still the predictor it was given.
    try testing.expectEqualSlices(i32, &.{ 33, -23, 10 }, q.coefficients[0..3]);
}

test "every coefficient fits the precision it was quantized to" {
    var prng: std.Random.DefaultPrng = .init(0x4c5043);
    const random = prng.random();
    for (0..2000) |_| {
        const precision = random.intRangeAtMost(u5, limits.min_qlp_precision, limits.max_qlp_precision);
        var coefficients: [limits.max_lpc_order]f64 = undefined;
        const order = random.intRangeAtMost(usize, 1, limits.max_lpc_order);
        for (coefficients[0..order]) |*c| c.* = (random.float(f64) - 0.5) * 4;

        const q = quantize(coefficients[0..order], precision) catch continue;
        const bound = @as(i64, 1) << (precision - 1);
        for (q.coefficients[0..order]) |c| {
            try testing.expect(c >= -bound and c < bound);
        }
        try testing.expect(q.shift <= limits.max_lpc_shift);
    }
}