using SPECTrecon: SPECTplan, psf_gauss
using SPECTrecon: project, project!, backproject, backproject!
using MIRTjim: jim, prompt
using ImagePhantoms: shepp_logan, SheppLoganEmis
using LinearAlgebra: mul!
using LinearMapsAA: LinearMapAA
using Plots: plot, default; default(markerstrokecolor=:auto)

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

T = Float32
nx,ny,nz = 128,128,1
xtrue = T.(shepp_logan(nx, SheppLoganEmis()))
xtrue = reshape(xtrue, nx, ny, 1) # 3D array with nz=1
jim(xtrue, "xtrue: SheppLoganEmis with size $(size(xtrue))")

px,pz = 11,1 # pz=1 is crucial for 2D work
psf1 = psf_gauss( ; ny, px, pz, fwhm_start = 1, fwhm_end = 4) # (px,pz,ny)
tmp = reshape(psf1, px, ny) / maximum(psf1) # (px,ny)
hx = (px-1)÷2
plot(-hx:hx, tmp[:,[1:9:end-10;end]], markershape=:o, label="",
    title = "Depth-dependent PSF profiles",
    xtick = [-hx, -2, 0, 2, hx], # (-1:1) .* ((px-1)÷2),
    ytick = [0; round.(tmp[hx+1,end] * [0.5,1], digits=2); 0.5; 1],
)
prompt()

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

dy = 4 # transaxial pixel size in mm
mumap = zeros(T, size(xtrue)) # μ-map just zero for illustration here
views = project(xtrue, mumap, psfs, dy) # [nx,1,nview]
sino = reshape(views, nx, nview)
size(sino)

jim(sino, "Sinogram")

sino1 = zeros(T, nx, nview)
sino1[nx÷2, nview÷5] = 1
sino1[nx÷2, 1] = 1
sino1 = reshape(sino1, nx, nz, nview)
back1 = backproject(sino1, mumap, psfs, dy)
jim(back1, "Back-projection of two rays")

back = backproject(views, mumap, psfs, dy)
jim(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

points = zeros(T, nx, ny, nz)
points[nx÷2,ny÷2,1] = 1
points[3nx÷4,ny÷4,1] = 1

impulse = A' * (A * points)
jim(impulse, "Impulse response of A'A")