using SPECTrecon: plan_rotate, imrotate!, imrotate_adj!
using MIRTjim: jim, prompt
using Plots: scatter, scatter!, plot!, default
default(markerstrokecolor=:auto, markersize=3)

isinteractive() ? jim(:prompt, true) : prompt(:draw);

T = Float32 # work with single precision to save memory
image = zeros(T, 64, 64, 2)
image[30:50,20:30,1] .= 1
image[25:28,20:40,2] .= 1
jim(image, "Original image")

plan2 = plan_rotate(size(image, 1); T)

plan2[1]

result2 = similar(image) # allocate memory for the result
imrotate!(result2, image, π/6, plan2) # mutates the first argument
jim(result2, "Rotated image by π/6 (2D bilinear)")

plan1 = plan_rotate(size(image, 1); T, method=:one)

plan1[1]

result1 = similar(image)
imrotate!(result1, image, π/6, plan1)
jim(result1, "Rotated image by π/6 (3-pass 1D)")

jim(result1 - result2, "Difference images")

adj2 = similar(result2)
imrotate_adj!(adj2, result2, π/6, plan2)
jim(adj2, "Adjoint image rotation (2D)")

adj1 = similar(result1)
imrotate_adj!(adj1, result1, π/6, plan1)
jim(adj1, "Adjoint image rotation (3-pass 1D)")

using LinearMapsAA: LinearMapAA

nx = 20 # small size for illustration
r1 = plan_rotate(nx; T, nthread = 1, method=:two)[1]
r2 = plan_rotate(nx; T, nthread = 1, method=:one)[1]
idim = (nx,nx)
odim = (nx,nx)
forw! = (y,x) -> imrotate!(y, x, π/6, r1)
back! = (x,y) -> imrotate_adj!(x, y, π/6, r1)
A1 = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)
forw! = (y,x) -> imrotate!(y, x, π/6, r2)
back! = (x,y) -> imrotate_adj!(x, y, π/6, r2)
A2 = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)

Afull1 = Matrix(A1)
Aadj1 = Matrix(A1')
Afull2 = Matrix(A2)
Aadj2 = Matrix(A2')
jim(cat(dims=3, Afull1', Aadj1', Afull2', Aadj2'), "Linear map for 2D rotation and its adjoint")

@assert Afull1' ≈ Aadj1
@assert Afull2' ≈ Aadj2

image2 = zeros(nx,nx); image2[4:6, 5:13] .= 1
jim(cat(dims=3, image2, A2 * image2), "Rotation via linear map: 2D")

image3 = cat(dims=3, image2, image2')
jim(cat(dims=4, image3, A2 * image3), "Rotation via linear map: 3D")

scatter(xlabel="pixel index", ylabel="row or col sum")
scatter!(vec(sum(Afull1, dims=1)), label="dim1 sum1", marker=:x)
scatter!(vec(sum(Afull1, dims=2)), label="dim2 sum1", marker=:square)
scatter!(vec(sum(Afull2, dims=1)), label="dim1 sum2")
scatter!(vec(sum(Afull2, dims=2)), label="dim2 sum2")