# # [How to port rules generated by SparseADRules to Zygote](@id lowrank_svd)
#
# In this demo we'll show how to insert SparseADRules's sparse gradient implementation to boost Zygote's gradient in lowrank svd.
using Zygote
using SparseArrays, LinearAlgebra
using SparseADRules:low_rank_svd
import FiniteDifferences
using BenchmarkTools
using ChainRulesCore, ChainRulesTestUtils
include(joinpath(@__DIR__, "..", "test", "testutils.jl"))

function svd_loss(A::AbstractSparseMatrix{T}, z::Vector{T}, r::Int) where T
    U, S, Vt = low_rank_svd(A, r)
    residual_mat = U * Diagonal(S) * Vt - U * Diagonal(z) * Vt
    return tr(residual_mat'*residual_mat) # square of Frobenius norm
end

# generate test data
r = 10
A = sprand(100, r, 0.2) * Diagonal(rand(r)) * sprand(r, 50, 0.2)
z = rand(r)

@btime grad_A = Zygote.gradient(A -> svd_loss(A, z, r), $A) 
grad_A = Zygote.gradient(A -> svd_loss(A, z, r), A)[1] 

grad_fA = FiniteDifferences.grad(FiniteDifferences.central_fdm(5, 1), A -> svd_loss(A, z, r), A)[1]

using Test
@testset "test AD on lowrank svd" begin
    @test grad_A ≈ grad_fA
end







