using SPECTrecon: plan_psf, psf_gauss, SPECTplan
using SPECTrecon: project, project!, backproject, backproject!
using MIRTjim: jim, prompt
using LinearAlgebra: mul!
using LinearMapsAA: LinearMapAA
using Plots: scatter, plot!, default; default(markerstrokecolor=:auto)
using Plots # @animate, gif
using InteractiveUtils: versioninfo

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

nx,ny,nz = 128,128,80
T = Float32
xtrue = zeros(T, nx,ny,nz)
xtrue[(1nx÷4):(2nx÷3), 1ny÷5:(3ny÷5), 2nz÷6:(3nz÷6)] .= 1
xtrue[(2nx÷5):(3nx÷5), 1ny÷5:(2ny÷5), 4nz÷6:(5nz÷6)] .= 2

average(x) = sum(x) / length(x)
function mid3(x::AbstractArray{T,3}) where {T}
    (nx,ny,nz) = size(x)
    xy = x[:,:,ceil(Int, nz÷2)]
    xz = x[:,ceil(Int,end/2),:]
    zy = x[ceil(Int, nx÷2),:,:]'
    return [xy xz; zy fill(average(x), nz, nz)]
end
jim(mid3(xtrue), "Middle slices of x")

px = 11
psf1 = psf_gauss( ; ny, px, fwhm_end = 6)
jim(psf1, "PSF for each of $ny planes")

nview = 60
psfs = repeat(psf1, 1, 1, 1, nview)
size(psfs)

plan = plan_psf( ; nx, nz, px)

dy = 4 # transaxial pixel size in mm
mumap = zeros(T, size(xtrue)) # μ-map just zero for illustration here
views = project(xtrue, mumap, psfs, dy)
size(views)

jim(views[:,:,1:4:end], "Every 4th of $nview projection views")

tmp = zeros(T, size(views))
tmp[nx÷2, nz÷2, nview÷5] = 1
tmp[nx÷2, nz÷2, 1] = 1
tmp = backproject(tmp, mumap, psfs, dy)
jim(mid3(tmp), "Back-projection of two rays")

back = backproject(views, mumap, psfs, dy)
jim(mid3(back), "Back-projection of ytrue")

#viewangle = (0:(nview-1)) * 2π # default
plan = SPECTplan(mumap, psfs, dy; T)

tmp = Array{T}(undef, nx, nz, nview)
project!(tmp, xtrue, plan)
@assert tmp == views

tmp = Array{T}(undef, nx, ny, nz)
backproject!(tmp, views, plan)
@assert tmp ≈ back

forw! = (y,x) -> project!(y, x, plan)
back! = (x,y) -> backproject!(x, y, plan)
idim = (nx,ny,nz)
odim = (nx,nz,nview)
A = LinearMapAA(forw!, back!, (prod(odim),prod(idim)); T, odim, idim)

@assert A * xtrue == views
@assert A' * views ≈ back

tmp = Array{T}(undef, nx, nz, nview)
mul!(tmp, A, xtrue)
@assert tmp == views
tmp = Array{T}(undef, nx, ny, nz)
mul!(tmp, A', views)
@assert tmp ≈ back

anim = @animate for i in 1:nview
    ymax = maximum(views)
    jim(views[:,:,i],
        "SPECT projection view $i of $nview",
        clim = (0, ymax),
    )
end
gif(anim, "views.gif", fps = 8)