QRM-Thorax Pure-Material PCCT VMI: Photon-Counting Reference

begin
    import Pkg
    Pkg.activate(joinpath(@__DIR__, ".."))
end

using Markdown: @md_str, Markdown

using Statistics: mean, std, quantile

import BasisSimulator as BS

# Use CairoMakie for faithful build-time rendering. Snapshot can still isolate
# and compile independent browser-safe islands without hoisting this import.
import CairoMakie as Mke

import PlutoUI

begin
    import GPUSelect
    AT = GPUSelect.Storage()     # the backend array type, directly: MtlArray / CuArray / ROCArray
    to_gpu(x) = AT(x)
    GPU_BACKEND = (name = string(nameof(AT)),)
end;

begin
    # The QRM-Thorax mid-slice as 2-D shapes, measured on the labelled 1600 × 1100 (0.2 mm)
    # reference slice. Millimetres from the slice centre, +x to the image right and +y up:
    # the CT display orientation, sternum at the top and spine at the bottom.
    const QRM_GEOMETRY = (
        # |x/a|^n + |(y - yc)/b|^n ≤ 1, cut flat at the front and back (y range)
        body = (a = 150.2, b = 106.75, n = 2.30, yc = 0.0, y = (-100.2, 101.0)),
        lungs = (a = 127.6, b = 92.2, n = 2.10, yc = 0.4, y = (-76.8, 76.8)),
        # soft tissue between the lungs above the heart; its corners at the lung apex are rounded
        mediastinum = (half_width = 16.8, corner_r = 2.0),
        heart = (center = (0.0, 18.4), r = 55.2),     # the cardiac insert, soft tissue (label 3)
        # the four pure-material rods in the heart, as the reference slice has them (labels 9–12)
        heart_rods = (
            (center = (0.0, 45.8), r = 7.5, label = 9),              # north: water
            (center = (24.9, 20.7), r = 7.5, label = 10),            # east: lipid
            (center = (0.0, -4.3), r = 7.5, label = 11),             # south: collagen
            (center = (-25.1, 20.7), r = 7.5, label = 12),           # west: 5 mg/mL iodine
        ),
        lung_rods = (
            (center = (-91.4, 0.0), r = 7.5, label = 6),             # water rod
            (center = (90.2, 0.0), r = 7.5, label = 7),              # air rod
        ),
        # the spine; `cortex` is the thickness of the cortical-bone shell, mm
        vertebra = (center = (0.0, -53.3), r = 17.4, cortex = 0.6),
        canal = (center = (0.0, -69.6), r = 8.0, cortex = 0.6),      # the spinal canal
        laminae = (corners = ((-35.4, -77.2), (35.4, -77.2), (0.0, -62.3)), cortex = 1.0),
        spinous = (corners = ((-2.6, -76.8), (2.6, -76.8), (0.0, -91.6)), cortex = 0.6),
        ribs = (                                                     # (x, y, long-axis angle in °)
            (-112.4, 64.9, 40.0), (108.9, 67.9, -40.0),
            (-138.1, 1.5, 90.0), (139.0, 6.1, 90.0),
            (-109.2, -65.9, -40.0), (109.2, -67.5, 40.0),
        ),
        rib_axes = (8.2, 3.0),                                       # semi-axes, mm
        sternum = (center = (-2.7, 89.4), axes = (7.6, 3.2)),
        marrow_scale = 0.85, # rib and sternum marrow = their cortical ellipse scaled by this
    )

    """
        qrm_thorax_slice(g = QRM_GEOMETRY; n = (1600, 1100), voxel_mm = 0.2, inserts = ())

    Label image of the QRM-Thorax mid-slice (1 air, 2 lung, 3 soft tissue, including the
    cardiac insert, 4 cortical bone, 5 marrow, 6 water rod, 7 air rod). `inserts` are extra
    discs `(center, r, label)` in mm, painted last: the heart rods of `g.heart_rods`.
    """
    function qrm_thorax_slice(g = QRM_GEOMETRY; n = (1600, 1100), voxel_mm = 0.2, inserts = ())
        superellipse(x, y, s) =
            s.y[1] ≤ y ≤ s.y[2] && abs(x / s.a)^s.n + abs((y - s.yc) / s.b)^s.n ≤ 1
        disc(x, y, c, r) = (x - c[1])^2 + (y - c[2])^2 ≤ r^2
        function ellipse(x, y, c, axes, θ)
            s, co = sincosd(θ)
            u, v = co * (x - c[1]) + s * (y - c[2]), -s * (x - c[1]) + co * (y - c[2])
            return (u / axes[1])^2 + (v / axes[2])^2 ≤ 1
        end
        # inside all three edges, each moved `inset` mm toward the interior
        function triangle(x, y, t, inset = 0.0)
            for k in 1:3
                p, q, r = t[k], t[mod1(k + 1, 3)], t[mod1(k + 2, 3)]
                ex, ey = q[1] - p[1], q[2] - p[2]
                side(px, py) = (ex * (py - p[2]) - ey * (px - p[1])) / hypot(ex, ey)
                side(x, y) * sign(side(r...)) ≥ inset || return false
            end
            return true
        end
        function lung(x, y)
            superellipse(x, y, g.lungs) || return false
            disc(x, y, g.heart.center, g.heart.r) && return false
            m = g.mediastinum
            abs(x) < m.half_width && y > g.heart.center[2] && return false
            # round the apex corner where the lung top meets the mediastinum
            cx, cy = m.half_width + m.corner_r, g.lungs.y[2] - m.corner_r
            abs(x) < cx && y > cy && hypot(abs(x) - cx, y - cy) > m.corner_r && return false
            return true
        end

        bones = (
            ((c = (rx, ry), axes = g.rib_axes, θ = θ) for (rx, ry, θ) in g.ribs)...,
            (c = g.sternum.center, axes = g.sternum.axes, θ = 0.0),
        )
        img = ones(UInt8, n)
        for j in 1:n[2], i in 1:n[1]
            x, y = (i - (n[1] + 1) / 2) * voxel_mm, (j - (n[2] + 1) / 2) * voxel_mm
            superellipse(x, y, g.body) || continue
            label = lung(x, y) ? 2 : 3    # the heart is not lung, so it is soft tissue
            for rod in g.lung_rods
                disc(x, y, rod.center, rod.r) && (label = rod.label)
            end
            for b in bones
                ellipse(x, y, b.c, b.axes, b.θ) &&
                    (label = ellipse(x, y, b.c, b.axes .* g.marrow_scale, b.θ) ? 5 : 4)
            end
            for t in (g.laminae, g.spinous)
                triangle(x, y, t.corners) && (label = triangle(x, y, t.corners, t.cortex) ? 5 : 4)
            end
            v, cn = g.vertebra, g.canal
            disc(x, y, v.center, v.r) && (label = disc(x, y, v.center, v.r - v.cortex) ? 5 : 4)
            disc(x, y, cn.center, cn.r) && (label = disc(x, y, cn.center, cn.r - cn.cortex) ? 3 : 4)
            for (c, r, l) in inserts
                disc(x, y, c, r) && (label = l)
            end
            img[i, j] = label
        end
        return img
    end
end;

const QRM_NZ = 20;    # 20 × 0.2 mm = 4 mm: a short z-invariant phantom

const QRM_VOXEL_SIZE_CM = (0.02, 0.02, 0.02);   # (x, y, z) cm: 0.2 mm isotropic

# (center in mm, radius in mm, label) of each heart rod
HEART_RODS = [(rod.center, rod.r, rod.label) for rod in QRM_GEOMETRY.heart_rods];

mask_3d = let
    slice = qrm_thorax_slice(; voxel_mm = QRM_VOXEL_SIZE_CM[1] * 10, inserts = HEART_RODS)
    repeat(slice; outer = (1, 1, QRM_NZ))
end;

materials_dict = Dict{Int, BS.XA.Material}(
    # Anatomy
    1 => BS.XA.Materials.air,
    2 => BS.XA.Materials.lung,
    3 => BS.XA.Materials.muscle,           # soft tissue: ICRU-44 muscle
    4 => BS.XA.Materials.corticalbone,
    5 => BS.XA.Materials.marrow_red,
    6 => BS.XA.Materials.water,            # water rod in the lung
    7 => BS.XA.Materials.air,              # air rod
    8 => BS.XA.Materials.muscle,           # unused: the heart insert is soft tissue (label 3)
    # Pure-material rod inserts
    9 => BS.XA.Materials.basis_water,
    10 => BS.XA.Materials.basis_lipid,
    11 => BS.XA.Materials.basis_collagen,
    12 => BS.XA.Materials.gammex_472_i5_0,  # Gammex 472 5 mg/mL iodine insert
);

phantom_cpu = BS.create_phantom_from_mask(
    Array{Int, 3}(mask_3d),
    materials_dict,
    QRM_VOXEL_SIZE_CM,
);

phantom = BS.Phantom(
    to_gpu(phantom_cpu.mask),
    phantom_cpu.materials,
    phantom_cpu.voxel_size,
    phantom_cpu.origin,
    phantom_cpu.extent,
);

QRM_LABEL_NAMES = Dict{Int, String}(
    1 => "1 air",
    2 => "2 lung",
    3 => "3 soft tissue",
    4 => "4 cortical bone",
    5 => "5 marrow",
    6 => "6 water rod",
    7 => "7 air rod",
    8 => "8 (unused)",
    9 => "9 water",
    10 => "10 lipid",
    11 => "11 collagen",
    12 => "12 iodine 5 mg/mL",
);

scanner = let
    native_col_mm, native_row_mm, binning = 0.275, 0.322, 2
    sid, sdd = 610.0, 1113.0
    col_iso = native_col_mm * binning / (sdd / sid)     # binned pixel at isocentre, mm
    row_iso = native_row_mm * binning / (sdd / sid)
    BS.PCCTScanner(
        source_to_isocenter = sid,
        source_to_detector = sdd,
        detector_rows = 144,
        detector_cols = ceil(Int, 500.0 / col_iso),     # covers the 50 cm scan field
        detector_row_size = row_iso,
        detector_col_size = col_iso,
        detector_shape = :arc,
        detector_col_offset = 0.25,                     # quarter-detector offset
        focal_spot_width = 0.4,
        focal_spot_length = 0.5,
        target_angle = 7.0,
        gantry_rotation_time = 0.5,
        scan_diameter = 500.0,
        gantry_aperture = 820.0,
        flat_filter_material = :aluminum,
        flat_filter_thickness = 3.0,
        bowtie_filter = :large_body,
        detector_material = :cdte,
        detector_depth = 1.6,
        fill_factor_row = 0.95,
        fill_factor_col = 0.95,
        energy_thresholds = [20.0, 35.0, 55.0, 70.0],   # four energy windows, keV
        energy_resolution = 10.0,
        charge_sharing_fwhm = 0.08,
        dead_time_ns = 5.0,
        pixel_mode = :standard,
        native_dexel_col_mm = native_col_mm,
        native_dexel_row_mm = native_row_mm,
        binning_factor = binning,
        pileup_correction = true,      # model-based inverse of the pile-up migration
        scatter_correction = true,     # scatter re-estimated and subtracted per bin
        noise_reduction = 0.0,         # exact Poisson counts
    )
end;

"The NAEOTOM Alpha's 140 kVp axial protocol at tube current `mA`."
alpha_protocol(mA) = BS.CTProtocol(
    kVp = 140,
    mA = mA,
    views = 1200,
    rotation_time = 0.5,
    collimation_mm = 5.0,
    additional_filters = [("Ti", 0.9)],
);

# The tube-current scale that makes the simulated CTDIvol equal the physical scan's
DOSE_SCALE = let physical = (mA = 174.0, ctdi_vol_mGy = 10.12)
    p = alpha_protocol(physical.mA)
    physical.ctdi_vol_mGy / BS.compute_dose(BS.dose_source(scanner, p), p).ctdi_vol_mGy
end

protocol = alpha_protocol(174.0 * DOSE_SCALE);

"The Alpha model's physics switches for one draw."
draw_options(; seed, use_noise) = BS.SimOptions(
    seed = seed,
    use_noise = use_noise,
    projector = :dd_fast,
    use_fill_factor = false,
    use_detector_efficiency = false,
    use_optical_crosstalk = false,
    use_focal_spot = true,
    use_heel_effect = false,
    use_lag = false,
    view_samples = 5,     # sub-views per view: the arc the gantry turns while a view is read
    view_arc = 1.0,       # a photon-counting detector integrates for the whole view period
);

# basis-spectral-denoising's seeds: draw r has seed 1234 + 1000 (r − 1); the calibration draw is draw 9
DRAWS = (
    reference = (seed = 1234, use_noise = false),
    calibration = (seed = 1234 + 1000 * (9 - 1), use_noise = true),
    measured = (seed = 1234, use_noise = true),
);

# The saved recon grid of notebook 07: 512 × 512 at 0.625 mm isotropic, 3 slices.
recon_opts = BS.ReconOptions(
    matrix_size = (512, 512, 3),
    fov_cm = 32.0,
    z_cm = 0.1875,
);

acquisition = let
    projection = nothing
    model = nothing
    channels = Dict{Symbol, Vector{Array{Float32, 3}}}()
    seconds = Dict{Symbol, Float64}()
    for name in (:reference, :calibration, :measured)     # the noise-free draw first: it projects
        opts = draw_options(; DRAWS[name]...)
        t0 = time()
        ws = BS.create_workspace(scanner, protocol, opts, recon_opts, phantom)
        try
            # the dose and the projection do not depend on the noise: both come from the first draw
            sim = BS.simulate!(ws, phantom, protocol, opts; capture_raw_counts = false,
                report_dose = projection === nothing,
                keep_projection = projection === nothing, projection = projection)
            projection === nothing && (projection = sim.projection)
            channels[name] = [Array(b) for b in sim.pcct_sino.bins]   # the four corrected bins
            if model === nothing
                I0 = Float32.(Array(sim.I0_bins))                     # [n_cols, n_rows, n_bins]
                model = (
                    I0 = I0,
                    basis = BS.spectral_basis(ws; I0 = I0),           # ray-resolved spectral model
                    geom = ws.geom,
                    dose = sim.dose,
                )
            end
        finally
            BS.release_backend!(ws)
        end
        seconds[name] = time() - t0
    end
    (; channels, seconds, model...)
end;

# The phantom labels resampled onto the reconstruction grid (nearest neighbour); every ROI
# below is built from these labels in recon-voxel coordinates.
phantom_in_recon = BS.resample_to_recon(
    phantom_cpu, acquisition.geom, recon_opts.matrix_size; method = :nearest,
);

HYPR_CHAIN = BS.SpectralHYPR();

# The decomposition and reconstruction settings of the chain; ACNR is added on the measured draw
VMI_CHAIN = (method = :nchannel, controls = BS.NChannelControls(), use_tlbf = false, antialias = true);

pcct_vmi_energies = [50.0, 70.0, 100.0, 140.0];

ALPHA_WINDOWS = (composite = (fc = 0.5, p = 2.0), complement = (fc = 0.45, p = 6.0));

PAIR_FILTER = let knots = Tuple(range(0.0, 1.0, length = 11))
    # the grid the windows were fitted on: the Alpha's clinical series, 512² over 35 cm
    fitted = BS.CTGeometry(scanner; n_angles = protocol.views, fov_cm = 35.0, z_cm = 0.48,
        collimation_mm = protocol.collimation_mm)
    r = BS.grid_bandlimit(acquisition.geom, recon_opts.matrix_size) /
        BS.grid_bandlimit(fitted, (512, 512, 12))
    window(q) = BS.CustomFilter(knots, Tuple(round(exp(-(x * r / q.fc)^q.p), digits = 5) for x in knots))
    (filter = BS.PairFilter(window(ALPHA_WINDOWS.composite), window(ALPHA_WINDOWS.complement)), r = r)
end;

PAIR_BASIS = let ch = acquisition.channels[:calibration], I0 = acquisition.I0
    decompose(channels) = BS.vmi_pipeline(; channels, basis = acquisition.basis,
        geom = acquisition.geom, to_backend = to_gpu, matrix_size = recon_opts.matrix_size,
        vmi_energies = Tuple(pcct_vmi_energies), keep_sinograms = true, use_acnr = false,
        VMI_CHAIN...).sinograms
    p = HYPR_CHAIN.projection
    d_none = decompose(ch)
    d_proj = decompose(BS.hypr_lr(ch, I0; kernel = p.kernel, dispersion = BS.estimate_dispersion(ch, I0),
        view_stride = p.view_stride, to_backend = to_gpu))
    sp(d; kw...) = BS.spectral_pair(d.water, d.iodine, acquisition.geom, recon_opts.matrix_size;
        filter = BS.SoftFilter(), antialias = VMI_CHAIN.antialias, to_backend = to_gpu, kw...)
    x = sp(d_none)
    c = sp(d_proj; basis = (Estar = x.Estar, β = x.β))
    (Estar = x.Estar, β = x.β, composite = c.Estar)
end

pcct_vmi = BS.vmi_pipeline(;
    channels = acquisition.channels[:measured],
    basis = acquisition.basis,
    geom = acquisition.geom,
    to_backend = to_gpu,
    matrix_size = recon_opts.matrix_size,
    vmi_energies = Tuple(pcct_vmi_energies),
    denoiser = HYPR_CHAIN,
    fbp_filter = PAIR_FILTER.filter,
    pair_basis = (Estar = PAIR_BASIS.Estar, β = PAIR_BASIS.β),
    composite_energy = PAIR_BASIS.composite,
    keep_sinograms = true,
    VMI_CHAIN...,
);

reference_vmi = BS.vmi_pipeline(;
    channels = acquisition.channels[:reference],
    basis = acquisition.basis,
    geom = acquisition.geom,
    to_backend = to_gpu,
    matrix_size = recon_opts.matrix_size,
    vmi_energies = Tuple(pcct_vmi_energies),
    fbp_filter = PAIR_FILTER.filter,
    pair_basis = (Estar = PAIR_BASIS.Estar, β = PAIR_BASIS.β),
    composite_energy = PAIR_BASIS.composite,
    use_acnr = false,
    VMI_CHAIN...,
);

# The VMI stack keyed by energy, (512, 512, 3) HU volumes.
vmi_HU_final = Dict(
    Float64(E) => pcct_vmi.vmis[:, :, :, k] for (k, E) in enumerate(pcct_vmi.energies)
);

# The noise-free reference's VMI stack, keyed alike.
vmi_HU_reference = Dict(
    Float64(E) => reference_vmi.vmis[:, :, :, k] for (k, E) in enumerate(reference_vmi.energies)
);

const ROD_LABELS = (UInt8(9), UInt8(10), UInt8(11), UInt8(12));

const ROD_NAMES = ("Water", "Lipid", "Collagen", "Iodine 5 mg/mL");

const ROD_MATERIALS = (
    BS.XA.Materials.basis_water,
    BS.XA.Materials.basis_lipid,
    BS.XA.Materials.basis_collagen,
    BS.XA.Materials.gammex_472_i5_0,
);

const ROD_ROI_RADIUS_PX = 8;    # 5 mm at 0.625 mm/voxel, inside the 7.5 mm rods

const HEART_NOISE_ROI_RADIUS_PX = 12;    # 7.5 mm: fits between the four rods

begin
    "Every value of `vol` inside the in-plane `roi`, over all slices."
    roi_values(vol, roi) = Float64[vol[c, z] for z in axes(vol, 3) for c in roi]

    # Rod centroids from the phantom labels on the recon grid; the noise ROI sits at the
    # centre of the rod cross, the mean of the four rod centroids.
    rod_rois = let
        mask_2d = phantom_in_recon[:, :, size(phantom_in_recon, 3) ÷ 2 + 1]
        nx, ny = size(mask_2d)
        function centroid(label)
            idx = findall(==(label), mask_2d)
            isempty(idx) && error("no recon voxel carries label $(label)")
            (sum(c -> c[1], idx) / length(idx), sum(c -> c[2], idx) / length(idx))
        end
        disc(c, r) = [CartesianIndex(i, j) for j in 1:ny for i in 1:nx
                      if (i - c[1])^2 + (j - c[2])^2 ≤ r^2]
        centers = Dict(lab => centroid(lab) for lab in ROD_LABELS)
        heart_center = (
            mean(first(centers[l]) for l in ROD_LABELS),
            mean(last(centers[l]) for l in ROD_LABELS),
        )
        (
            centers = centers,
            rods = Dict(lab => disc(centers[lab], ROD_ROI_RADIUS_PX) for lab in ROD_LABELS),
            heart_center = heart_center,
            heart = disc(heart_center, HEART_NOISE_ROI_RADIUS_PX),
        )
    end
end;

water_rod_hu = [mean(roi_values(vmi_HU_final[E], rod_rois.rods[UInt8(9)])) for E in pcct_vmi_energies];

water_rod_hu_reference = [mean(roi_values(vmi_HU_reference[E], rod_rois.rods[UInt8(9)])) for E in pcct_vmi_energies];

vmi_noise_by_keV = Dict(
    E => let v = roi_values(vmi_HU_final[E], rod_rois.heart),
             d = roi_values(vmi_HU_final[E] .- vmi_HU_reference[E], rod_rois.heart)
        (mean = mean(v), std = std(v), std_vs_reference = std(d), n = length(v))
    end
    for E in pcct_vmi_energies
);

# The basis-pair covariance in the heart ROI, which sets the noise of every VMI:
#   σ_HU(E)² = 10⁶ V_w + α(E)² V_i + 2·10³ α(E) C_iw,  α(E) = (μ/ρ)_I(E) / (μ/ρ)_w(E),
# with water in g/mL and iodine in mg/mL; lowest at α* = -10³ C_iw / V_i. Anti-correlated basis noise (ρ < 0) is what ACNR removes.
basis_covariance = let
    w = roi_values(pcct_vmi.images.water, rod_rois.heart)
    i = roi_values(pcct_vmi.images.iodine, rod_rois.heart) .* 1000.0   # mg/mL
    Vw, Vi = mean(abs2, w .- mean(w)), mean(abs2, i .- mean(i))
    Ciw = mean((w .- mean(w)) .* (i .- mean(i)))
    α(E) = BS.compute_mass_μ_at_energy(BS.XA.Elements.Iodine, E) /
        BS.compute_mass_μ_at_energy(BS.XA.Materials.water, E)
    α_star = -1000.0 * Ciw / Vi
    E_star = let grid = 40.0:1.0:140.0
        grid[argmin(abs.(α.(grid) .- α_star))]
    end
    (σ_water = sqrt(Vw), σ_iodine = sqrt(Vi), ρ = Ciw / sqrt(Vw * Vi), α_star, E_star,
     α_range = (α(140.0), α(40.0)))
end;

rod_data = let
    μ_water(E) = BS.compute_μ_at_energy(BS.XA.Materials.water, E)
    theoretical_hu(mat, E) = 1000.0 * (BS.compute_μ_at_energy(mat, E) - μ_water(E)) / μ_water(E)
    meas = [mean(roi_values(vmi_HU_final[E], rod_rois.rods[lab])) for lab in ROD_LABELS, E in pcct_vmi_energies]
    noise_free = [mean(roi_values(vmi_HU_reference[E], rod_rois.rods[lab])) for lab in ROD_LABELS, E in pcct_vmi_energies]
    theo = [theoretical_hu(mat, E) for mat in ROD_MATERIALS, E in pcct_vmi_energies]
    rmse = [sqrt(mean(abs2, meas[i, :] .- theo[i, :])) for i in eachindex(ROD_LABELS)]
    rmse_noise_free = [sqrt(mean(abs2, noise_free[i, :] .- theo[i, :])) for i in eachindex(ROD_LABELS)]
    (labels = ROD_LABELS, names = ROD_NAMES, measured = meas, noise_free = noise_free, theoretical = theo,
     rmse = rmse, rmse_noise_free = rmse_noise_free)
end;