using ImagePhantoms: rect
using ImagePhantoms: phantom, radon, spectrum
import ImagePhantoms as IP
using ImageGeoms: ImageGeom, axesf
using MIRTjim: jim, prompt
using FFTW: fft, fftshift, ifftshift
using Unitful: mm, unit, °
using Plots: plot, plot!, scatter!, default
default(markerstrokecolor=:auto)

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

width = (20mm, 80mm)
ob = rect((40mm, 30mm), width, π/6, 1.0f0)

dx, dy = 0.8mm, 1.0mm
M, N = (2^8, 2^8+2)
x = (-M÷2:M÷2-1) * dx
y = (-N÷2:N÷2-1) * dy
oversample = 2
img = phantom(x, y, [ob], oversample)
jim(x, y, img, "Rect image")

M, N = (2^8, 2^8+17) # odd
ig = ImageGeom(dims=(M,N), deltas=(dx,dy), offsets=:dsp)
@assert axes(ig)[1] ≈ x
oversample = 2
img = phantom(axes(ig)..., [ob], oversample)
p1 = jim(axes(ig), img, "Rect phantom", xlabel="x", ylabel="y")

area = IP.area(ob)
(sum(img) * prod(ig.deltas), area)

vscale = 1 / area # normalize spectra by area
spectrum_exact = spectrum(axesf(ig)..., [ob]) * vscale
sp = z -> max(log10(abs(z)/oneunit(abs(z))), -6) # log-scale for display
clim = (-6, 0) # colorbar limit for display
(xlabel, ylabel) = ("ν₁", "ν₂")
p2 = jim(axesf(ig), sp.(spectrum_exact), "log10|Spectrum|"; clim, xlabel, ylabel)

function myfft(x)
    u = unit(eltype(x))
    return fftshift(fft(ifftshift(x) / u)) * u
end

spectrum_fft = myfft(img) * (prod(ig.deltas) * vscale)
p3 = jim(axesf(ig), sp.(spectrum_fft), "log10|DFT|"; clim, xlabel, ylabel)

errf = maximum(abs, spectrum_exact - spectrum_fft) / maximum(abs, spectrum_exact)
@assert errf < 2e-2
p4 = jim(axesf(ig), 1e3*abs.(spectrum_fft - spectrum_exact),
    "|Difference| × 10³"; xlabel, ylabel)
jim(p1, p4, p2, p3)

dr = 0.2mm # radial sample spacing
nr = 2^10 # radial sinogram bins
r = (-nr÷2:nr÷2-1) * dr # radial samples
fr = (-nr÷2:nr÷2-1) / nr / dr # corresponding spectral axis
ϕ = deg2rad.(0:180)
sino = radon(ob).(r, ϕ') # sample Radon transform of a single shape object
smax = ob.value * sqrt(sum(abs2, maximum(width)))
p5 = jim(r, rad2deg.(ϕ), sino; title="sinogram",
    xlabel="r", ylabel="ϕ", clim = (0,1) .* smax)

ia = argmin(abs.(ϕ .- 55°))
slice = sino[:,ia]
slice_fft = myfft(slice) * dr
ϕd = round(rad2deg(ϕ[ia]), digits=1)

kx, ky = (fr * cos(ϕ[ia]), fr * sin(ϕ[ia])) # Fourier-slice theorem
slice_ft = spectrum(ob).(kx, ky)
errs = maximum(abs, slice_ft - slice_fft) / maximum(abs, slice_ft)
@assert errs < 3e-5

p3 = plot(r, slice, title="profile at ϕ = $ϕd", label="")
p4 = plot(title="1D spectra")
scatter!(fr, abs.(slice_fft), label="abs fft", color=:blue)
scatter!(fr, real(slice_fft), label="real fft", color=:green)
scatter!(fr, imag(slice_fft), label="imag fft", color=:red,
    xlims=(-1,1) .* (0.1/mm))

plot!(fr, abs.(slice_ft), label="abs", color=:blue)
plot!(fr, real(slice_ft), label="real", color=:green)
plot!(fr, imag(slice_ft), label="imag", color=:red)
plot(p1, p5, p3, p4)