using SPECTrecon: psf_gauss, plan_psf, fft_conv!, fft_conv_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
nx = 32
nz = 30
image = zeros(T, nx, nx, nz) # ny = nx required
image[1nx÷4, 1nx÷4, 3nz÷4] = 1
image[2nx÷4, 2nx÷4, 2nz÷4] = 1
image[3nx÷4, 3nx÷4, 1nz÷4] = 1
jim(image, "Original image")

px = 11
nview = 1 # for simplicity in this illustration
psf = repeat(psf_gauss( ; ny=nx, px), 1, 1, 1, nview)
jim(psf, "PSF for each of $nx planes")

plan = plan_psf( ; nx, nz, px, T)

plan[1]

result = similar(image) # allocate memory for the result
fft_conv!(result, image, psf[:,:,:,1], plan) # mutates the first argument
jim(result, "After applying PSF")

adj = similar(result)
fft_conv_adj!(adj, result, psf[:,:,:,1], plan)
jim(adj, "Adjoint of PSF modeling")

using LinearMapsAA: LinearMapAA

nx, nz, px = 10, 7, 5 # small size for illustration
psf3 = psf_gauss( ; ny=nx, px)
plan = plan_psf( ; nx, nz, px, T)
idim = (nx,nx,nz)
odim = (nx,nx,nz)
forw! = (y,x) -> fft_conv!(y, x, psf3, plan)
back! = (x,y) -> fft_conv_adj!(x, y, psf3, plan)
A = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)

Afull = Matrix(A)
Aadj = Matrix(A')
jim(cat(dims=3, Afull, Aadj'), "Linear map for PSF modeling and its adjoint")

@assert Afull ≈ Aadj'