const std = @import("std");

const Matrix = @import("Matrix.zig").Matrix;

/// Controls the size and computation of the left singular vectors matrix (U) in SVD.
pub const Mode = enum {
    /// Skip computation of U matrix entirely. Use when only singular values are needed.
    ///
    /// Note: For dynamic SVD, the returned U matrix is empty (0×0).
    /// For static SVD, the returned U matrix is m×n with undefined contents.
    no_u,
    /// Compute only the first n columns of U (economy/thin SVD). Results in U being m×n.
    /// More memory efficient when m >> n.
    skinny_u,
    /// Compute the full m×m U matrix. Use when all left singular vectors are needed.
    full_u,
};

/// Internal state machine for the SVD algorithm's iterative process.
/// Based on the classical Golub-Reinsch algorithm.
const State = enum {
    /// Test if the superdiagonal element can be set to zero (decoupling test)
    test_splitting,
    /// Cancel the superdiagonal element using Givens rotations
    cancellation,
    /// Check if the algorithm has converged for the current singular value
    test_convergence,
    /// Final convergence check and sign correction
    convergence_check,
};

/// Configuration options for SVD computation.
/// Allows fine-grained control over which matrices are computed.
pub const Options = struct {
    /// Whether to compute the right singular vectors (V matrix).
    with_v: bool = false,
    /// Controls computation and size of the U matrix.
    mode: Mode = .full_u,
};

/// Result type for SVD decomposition: A = U × Σ × V^T
/// where A is the input matrix, U contains left singular vectors,
/// Σ is a diagonal matrix of singular values (stored as a vector),
/// and V contains right singular vectors.
///
/// The dimensions of matrices depend on the options used:
/// - U: m×m (full_u), m×n (skinny_u), or empty (no_u)
/// - s: n×1 vector of singular values in descending order
/// - V: n×n matrix (or empty if with_v=false)
pub fn Result(comptime T: type) type {
    return struct {
        /// Left singular vectors matrix. Each column is a left singular vector.
        /// Dimensions: m×m (full_u), m×n (skinny_u), or 0×0 (no_u)
        u: Matrix(T),
        /// Singular values in descending order as a column vector.
        /// These are the diagonal elements of the Σ matrix.
        s: Matrix(T),
        /// Right singular vectors matrix. Each column is a right singular vector.
        /// The matrix is orthogonal: V^T × V = I
        v: Matrix(T),
        /// Convergence status: 0 if successful, k if failed at k-th singular value.
        /// Non-zero values indicate the iterative algorithm failed to converge.
        converged: usize,

        pub fn deinit(self: *@This()) void {
            self.u.deinit();
            self.s.deinit();
            self.v.deinit();
        }
    };
}

/// Singular value decomposition of `a` (m×n with m ≥ n) into U, Σ, V^T.
///
/// Adapted from dlib's svd4, which translated the original Algol from
/// "Handbook for Automatic Computation, vol. II, Linear Algebra" (Springer-Verlag) into C.
/// An iteration counter is added to prevent stalls.
///
/// Sets `converged = 0` on success, or `k` if the algorithm fails to converge at the k-th singular value.
pub fn svd(
    comptime T: type,
    allocator: std.mem.Allocator,
    a: Matrix(T),
    options: Options,
) !Result(T) {
    std.debug.assert(a.rows >= a.cols);
    const m = a.rows;
    const n = a.cols;

    var u: Matrix(T) = if (options.mode == .full_u)
        try .initAll(allocator, m, m, 0)
    else
        try .initAll(allocator, m, n, 0);
    errdefer u.deinit();

    var v: Matrix(T) = if (options.with_v)
        try .initAll(allocator, n, n, 0)
    else
        try .init(allocator, 0, 0);
    errdefer v.deinit();

    var q: Matrix(T) = try .initAll(allocator, n, 1, 0);
    errdefer q.deinit();

    var e: Matrix(T) = try .initAll(allocator, n, 1, 0);
    defer e.deinit();

    const converged = kernel(a, &u, &v, &q, &e, options.with_v, options.mode);
    if (options.mode == .no_u) {
        const empty_u: Matrix(T) = try .init(allocator, 0, 0);
        u.deinit();
        u = empty_u;
    }
    return .{ .u = u, .s = q, .v = v, .converged = converged };
}

/// Comptime check that `M` is matrix-like: an `at(row, col)` accessor plus
/// `rows` and `cols` members, as exposed by both Matrix and SMatrix.
fn assertMatrixLike(comptime name: []const u8, comptime M: type) void {
    if (@typeInfo(M) != .@"struct")
        @compileError("svd kernel: `" ++ name ++ "` must be a matrix, got " ++ @typeName(M));
    if (!@hasDecl(M, "at"))
        @compileError("svd kernel: `" ++ name ++ "` (" ++ @typeName(M) ++ ") has no `at(row, col)` accessor");
    if (!@hasField(M, "rows") or !@hasField(M, "cols"))
        @compileError("svd kernel: `" ++ name ++ "` (" ++ @typeName(M) ++ ") has no `rows`/`cols` members");
}

/// Comptime check that `P` is a mutable single-item pointer to a matrix-like type.
fn assertMatrixPtr(comptime name: []const u8, comptime P: type) void {
    const info = @typeInfo(P);
    if (info != .pointer or info.pointer.size != .one or info.pointer.attrs.@"const")
        @compileError("svd kernel: `" ++ name ++ "` must be a mutable pointer to a matrix, got " ++ @typeName(P));
    assertMatrixLike(name, info.pointer.child);
}

/// Comptime check that a matrix shares the element type of `a`.
fn assertSameElem(comptime name: []const u8, comptime T: type, comptime E: type) void {
    if (E != T)
        @compileError("svd kernel: `" ++ name ++ "` has element type " ++ @typeName(E) ++
            ", expected " ++ @typeName(T) ++ " (the element type of `a`)");
}

/// Golub-Reinsch kernel shared by the dynamic (`svd` above) and static
/// (`SMatrix.svd`) entry points. `u`, `v`, `q` and `e` point to
/// pre-sized matrices; all element access goes through the `at()` accessor
/// that both matrix types expose. The element type and dimensions are taken
/// from `a`. Returns 0 on success, or k if the iteration fails to converge
/// at the k-th singular value.
pub fn kernel(
    a: anytype,
    u: anytype,
    v: anytype,
    q: anytype,
    e: anytype,
    with_v: bool,
    mode: Mode,
) usize {
    comptime {
        assertMatrixLike("a", @TypeOf(a));
        assertMatrixPtr("u", @TypeOf(u));
        assertMatrixPtr("v", @TypeOf(v));
        assertMatrixPtr("q", @TypeOf(q));
        assertMatrixPtr("e", @TypeOf(e));
    }
    const T = @TypeOf(a.at(0, 0).*);
    comptime {
        if (@typeInfo(T) != .float) {
            @compileError("svd kernel: element type must be a floating point type, got " ++ @typeName(T));
        }
        assertSameElem("u", T, @TypeOf(u.at(0, 0).*));
        assertSameElem("v", T, @TypeOf(v.at(0, 0).*));
        assertSameElem("q", T, @TypeOf(q.at(0, 0).*));
        assertSameElem("e", T, @TypeOf(e.at(0, 0).*));
    }

    const m: usize = a.rows;
    const n: usize = a.cols;
    const max_iterations: usize = 300;
    var eps: T = std.math.floatEps(T);
    const tol: T = std.math.floatMin(T) / eps;

    var l: usize = 0;
    var retval: usize = 0;
    var c: T = undefined;
    var f: T = undefined;
    var g: T = undefined;
    var h: T = undefined;
    var s: T = undefined;
    var x: T = undefined;
    var y: T = undefined;
    var z: T = undefined;

    // Copy a to u.
    for (0..m) |i| {
        for (0..n) |j| {
            u.at(i, j).* = a.at(i, j).*;
        }
    }

    // Householder's reduction to bidiagonal form.
    g = 0;
    x = 0;
    for (0..n) |i| {
        e.at(i, 0).* = g;
        s = 0;
        l = i + 1;

        for (i..m) |j| {
            s += u.at(j, i).* * u.at(j, i).*;
        }

        if (s < tol) {
            g = 0;
        } else {
            f = u.at(i, i).*;
            g = if (f < 0) @sqrt(s) else -@sqrt(s);
            h = f * g - s;
            u.at(i, i).* = f - g;

            for (l..n) |j| {
                s = 0;
                for (i..m) |k| {
                    s += u.at(k, i).* * u.at(k, j).*;
                }
                f = s / h;

                for (i..m) |k| {
                    u.at(k, j).* += f * u.at(k, i).*;
                }
            }
        }

        q.at(i, 0).* = g;
        s = 0;

        for (l..n) |j| {
            s += u.at(i, j).* * u.at(i, j).*;
        }

        if (s < tol) {
            g = 0;
        } else {
            f = u.at(i, i + 1).*;
            g = if (f < 0) @sqrt(s) else -@sqrt(s);
            h = f * g - s;
            u.at(i, i + 1).* = f - g;

            for (l..n) |j| {
                e.at(j, 0).* = u.at(i, j).* / h;
            }

            for (l..m) |j| {
                s = 0;
                for (l..n) |k| {
                    s += u.at(j, k).* * u.at(i, k).*;
                }
                for (l..n) |k| {
                    u.at(j, k).* += s * e.at(k, 0).*;
                }
            }
        }
        y = @abs(q.at(i, 0).*) + @abs(e.at(i, 0).*);
        x = @max(x, y);
    }

    // Accumulation of right-hand transformations.
    if (with_v) {
        for (0..n) |ri| {
            const i = n - 1 - ri;
            if (g != 0) {
                h = u.at(i, i + 1).* * g;
                for (l..n) |j| {
                    v.at(j, i).* = u.at(i, j).* / h;
                }
                for (l..n) |j| {
                    s = 0;
                    for (l..n) |k| {
                        s += u.at(i, k).* * v.at(k, j).*;
                    }
                    for (l..n) |k| {
                        v.at(k, j).* += s * v.at(k, i).*;
                    }
                }
            }
            for (l..n) |j| {
                v.at(i, j).* = 0;
                v.at(j, i).* = 0;
            }
            v.at(i, i).* = 1;
            g = e.at(i, 0).*;
            l = i;
        }
    }

    // Accumulation of left-hand transformations.
    if (mode != .no_u) {
        for (n..u.rows) |i| {
            for (n..u.cols) |j| {
                u.at(i, j).* = 0;
            }
            if (i < u.cols) {
                u.at(i, i).* = 1;
            }
        }
    }

    if (mode != .no_u) {
        for (0..n) |ri| {
            const i = n - 1 - ri;
            l = i + 1;
            g = q.at(i, 0).*;

            for (l..u.cols) |j| {
                u.at(i, j).* = 0;
            }
            if (g != 0) {
                h = u.at(i, i).* * g;
                for (l..u.cols) |j| {
                    s = 0;
                    for (l..m) |k| {
                        s += u.at(k, i).* * u.at(k, j).*;
                    }
                    f = s / h;
                    for (i..m) |k| {
                        u.at(k, j).* += f * u.at(k, i).*;
                    }
                }
                for (i..m) |j| {
                    u.at(j, i).* /= g;
                }
            } else {
                for (i..m) |j| {
                    u.at(j, i).* = 0;
                }
            }
            u.at(i, i).* += 1;
        }
    }

    // Diagonalization of the bidiagonal form.
    eps *= x;

    for (0..n) |rk| {
        const k = n - 1 - rk;
        var iter: usize = 0;

        state: switch (State.test_splitting) {
            .test_splitting => {
                for (0..k + 1) |rl| {
                    l = k - rl;
                    if (@abs(e.at(l, 0).*) <= eps) {
                        continue :state .test_convergence;
                    }
                    if (@abs(q.at(l - 1, 0).*) <= eps) {
                        continue :state .cancellation;
                    }
                }
                continue :state .test_convergence;
            },

            .cancellation => {
                // Cancellation of e.at(l, 0) if l > 0
                c = 0;
                s = 1;
                const l1 = l - 1;
                for (l..k + 1) |i| {
                    f = s * e.at(i, 0).*;
                    e.at(i, 0).* *= c;

                    if (@abs(f) <= eps) {
                        continue :state .test_convergence;
                    }
                    g = q.at(i, 0).*;
                    h = @sqrt(f * f + g * g);
                    q.at(i, 0).* = h;
                    c = g / h;
                    s = -f / h;
                    if (mode != .no_u) {
                        for (0..m) |j| {
                            y = u.at(j, l1).*;
                            z = u.at(j, i).*;
                            u.at(j, l1).* = y * c + z * s;
                            u.at(j, i).* = -y * s + z * c;
                        }
                    }
                }
                continue :state .test_convergence;
            },

            .test_convergence => {
                z = q.at(k, 0).*;
                if (l == k) {
                    continue :state .convergence_check;
                }
                // Shift from bottom 2x2 minor.
                iter += 1;
                if (iter > max_iterations) {
                    retval = k;
                    break :state;
                }
                x = q.at(l, 0).*;
                y = q.at(k - 1, 0).*;
                g = e.at(k - 1, 0).*;
                h = e.at(k, 0).*;
                f = ((y - z) * (y + z) + (g - h) * (g + h)) / (2 * h * y);
                g = @sqrt(f * f + 1.0);
                f = ((x - z) * (x + z) + h * (y / (if (f < 0) (f - g) else (f + g)) - h)) / x;

                // Next QR transformation.
                c = 1;
                s = 1;
                for (l + 1..k + 1) |i| {
                    g = e.at(i, 0).*;
                    y = q.at(i, 0).*;
                    h = s * g;
                    g *= c;
                    z = @sqrt(f * f + h * h);
                    e.at(i - 1, 0).* = z;
                    c = f / z;
                    s = h / z;
                    f = x * c + g * s;
                    g = -x * s + g * c;
                    h = y * s;
                    y *= c;
                    if (with_v) {
                        for (0..n) |j| {
                            x = v.at(j, i - 1).*;
                            z = v.at(j, i).*;
                            v.at(j, i - 1).* = x * c + z * s;
                            v.at(j, i).* = -x * s + z * c;
                        }
                    }
                    z = @sqrt(f * f + h * h);
                    q.at(i - 1, 0).* = z;
                    if (z != 0) {
                        c = f / z;
                        s = h / z;
                    }
                    f = c * g + s * y;
                    x = -s * g + c * y;
                    if (mode != .no_u) {
                        for (0..m) |j| {
                            y = u.at(j, i - 1).*;
                            z = u.at(j, i).*;
                            u.at(j, i - 1).* = y * c + z * s;
                            u.at(j, i).* = -y * s + z * c;
                        }
                    }
                }
                e.at(l, 0).* = 0;
                e.at(k, 0).* = f;
                q.at(k, 0).* = x;
                continue :state .test_splitting;
            },

            .convergence_check => {
                if (z < 0) {
                    q.at(k, 0).* = -z;
                    if (with_v) {
                        for (0..n) |j| {
                            v.at(j, k).* = -v.at(j, k).*;
                        }
                    }
                }
                break :state;
            },
        }
    }
    // Sort singular values in descending order.
    for (0..n) |i| {
        var max_idx = i;
        var max_val = q.at(i, 0).*;
        for (i + 1..n) |j| {
            if (q.at(j, 0).* > max_val) {
                max_idx = j;
                max_val = q.at(j, 0).*;
            }
        }

        if (max_idx != i) {
            std.mem.swap(T, q.at(i, 0), q.at(max_idx, 0));
            if (mode != .no_u) {
                for (0..m) |row| {
                    std.mem.swap(T, u.at(row, i), u.at(row, max_idx));
                }
            }
            if (with_v) {
                for (0..n) |row| {
                    std.mem.swap(T, v.at(row, i), v.at(row, max_idx));
                }
            }
        }
    }

    return retval;
}

test "svd basic" {
    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const m: usize = 5;
    const n: usize = 4;
    // Example matrix taken from Wikipedia
    var a: Matrix(f64) = try .init(allocator, m, n);
    const data = [m][n]f64{
        .{ 1, 0, 0, 0 },
        .{ 0, 0, 0, 2 },
        .{ 0, 3, 0, 0 },
        .{ 0, 0, 0, 0 },
        .{ 2, 0, 0, 0 },
    };
    for (0..m) |i| {
        for (0..n) |j| {
            a.at(i, j).* = data[i][j];
        }
    }

    var res = try a.svd(allocator, .{ .with_v = true, .mode = .full_u });
    defer res.deinit();
    const u = &res.u;
    const s = &res.s;
    const v = &res.v;

    // Check that we got the right dimensions
    try std.testing.expectEqual(@as(usize, m), u.rows);
    try std.testing.expectEqual(@as(usize, m), u.cols);
    try std.testing.expectEqual(@as(usize, n), s.rows);
    try std.testing.expectEqual(@as(usize, 1), s.cols);
    try std.testing.expectEqual(@as(usize, n), v.rows);
    try std.testing.expectEqual(@as(usize, n), v.cols);

    // Check convergence
    try std.testing.expectEqual(@as(usize, 0), res.converged);

    // Check that singular values are non-negative and in descending order
    for (0..n) |i| {
        try std.testing.expect(s.at(i, 0).* >= 0);
        if (i > 0) {
            try std.testing.expect(s.at(i - 1, 0).* >= s.at(i, 0).*);
        }
    }
}

test "svd modes" {
    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const m: usize = 4;
    const n: usize = 4;
    var a: Matrix(f64) = try .fromSlice(allocator, m, n, &.{
        2, 1, 0, 0,
        1, 2, 1, 0,
        0, 1, 2, 1,
        0, 0, 1, 2,
    });

    // Test no_u mode
    var res_no_u = try a.svd(allocator, .{ .with_v = true, .mode = .no_u });
    defer res_no_u.deinit();
    try std.testing.expectEqual(@as(u32, 0), res_no_u.u.rows);
    try std.testing.expectEqual(@as(u32, 0), res_no_u.u.cols);
    const s_no_u = &res_no_u.s;

    // Test skinny_u mode
    var res_skinny = try a.svd(allocator, .{ .with_v = false, .mode = .skinny_u });
    defer res_skinny.deinit();
    const u_skinny = &res_skinny.u;
    const s_skinny = &res_skinny.s;

    // Test full_u mode
    var res_full = try a.svd(allocator, .{ .with_v = true, .mode = .full_u });
    defer res_full.deinit();
    const u_full = &res_full.u;
    const s_full = &res_full.s;

    // Singular values should be the same across modes
    const tol = @sqrt(std.math.floatEps(f64));
    for (0..n) |i| {
        try std.testing.expectApproxEqRel(s_no_u.at(i, 0).*, s_skinny.at(i, 0).*, tol);
        try std.testing.expectApproxEqRel(s_skinny.at(i, 0).*, s_full.at(i, 0).*, tol);
    }

    // Check matrix dimensions
    try std.testing.expect(u_skinny.rows == m and u_skinny.cols == n);
    try std.testing.expect(u_full.rows == m and u_full.cols == m);
}

test "svd identity matrix" {
    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const n: usize = 3;
    var a: Matrix(f64) = try .identity(allocator, n, n);

    var res = try a.svd(allocator, .{ .with_v = true, .mode = .full_u });
    defer res.deinit();
    const s = &res.s;

    // Identity matrix should have all singular values equal to 1
    const tol = @sqrt(std.math.floatEps(f64));
    for (0..n) |i| {
        try std.testing.expectApproxEqRel(s.at(i, 0).*, 1.0, tol);
    }
}

test "svd singular matrix" {
    var arena: std.heap.ArenaAllocator = .init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();

    const m: usize = 3;
    const n: usize = 3;
    var a: Matrix(f64) = try .fromSlice(allocator, m, n, &.{
        1, 2, 3,
        2, 4, 6,
        1, 2, 3,
    });

    var res = try a.svd(allocator, .{ .with_v = true, .mode = .full_u });
    defer res.deinit();
    const s = &res.s;

    // This matrix has rank 1, so should have 2 zero singular values
    const tol = @sqrt(std.math.floatEps(f64));
    var zero_count: usize = 0;
    for (0..n) |i| {
        if (s.at(i, 0).* < tol) {
            zero_count += 1;
        }
    }
    try std.testing.expect(zero_count == 2);
}
