09 · Siemens SOMATOM Force · UFC MC LUT · Dual-Source VMI

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;

phantom_cpu = BS.create_gammex_472(
    n_voxels = 512,
    n_slices = 16,
    fov_cm = 35.0,
    z_cm = 1.0,
);

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

# Tube/detector A geometry — shared by both modeled tubes (see md above).
scanner = BS.EICTScanner(
    source_to_isocenter = 595.0,
    source_to_detector = 1085.6,

    detector_rows = 96,
    detector_cols = 920,
    detector_row_size = 0.6,
    detector_col_size = 0.561,

    focal_spot_width = 0.8,
    focal_spot_length = 1.2,
    target_angle = 8.0,

    flat_filter_material = :aluminum,
    flat_filter_thickness = 3.0,
    bowtie_filter = :large_body,

    detector_material = :ufc,    # free-form tag; η comes from the MC LUT below
    detector_depth = 1.4,
    fill_factor_row = 0.9,
    fill_factor_col = 0.9,

    electronic_noise = 0,
    detection_gain = 10.0,
);

protocol_low = BS.CTProtocol(
    kVp = 100,
    mA = 380.0,
    views = 1160,
    rotation_time = 0.5,
    collimation_mm = 4.8,    # nominal beam width; axial cone guards are automatic
    anode_angle = 8,
    additional_filters = [("Ti", 0.9)],
);

protocol_high = BS.CTProtocol(
    kVp = 140,
    mA = 190.0,
    views = 1160,
    rotation_time = 0.5,
    collimation_mm = 4.8,
    anode_angle = 8,
    additional_filters = [("Ti", 0.9), ("Sn", 0.6)],
);

begin
    VIEW_SAMPLES = 5
    force_opts = BS.SimOptions(projector = :dd_fast, use_heel_effect = false, view_samples = VIEW_SAMPLES)
    "The same physics switches with one field replaced."
    with_option(opts, field, value) = BS.SimOptions(;
        (f => getfield(opts, f) for f in fieldnames(BS.SimOptions) if f !== field)...,
        (field => value,)...,
    )
end;

begin
    # The two exposures: tube A sees the phantom, tube B the phantom through its −0.88 mm z-offset (§6).
    # Draw 0 is the noise-free expectation; draw r ≥ 1 gives each tube the seed `seed + 1000 (r − 1)`,
    # so the tubes stay independent within a draw, and draws of each other.
    EXPOSURES = [
        (label = "100 kVp", protocol = protocol_low, seed = 1234, view_arc = 1.0, tube = :a),
        (label = "Sn140 kVp", protocol = protocol_high, seed = 4321, view_arc = 1.0, tube = :b),
    ]
    DE_DRAWS = (reference = 0, measured = 1, calibration = 9)
    realization_seed(e, r) = e.seed + 1000 * (max(r, 1) - 1)
    opts_for(e, r) = with_option(with_option(with_option(force_opts, :seed, realization_seed(e, r)),
        :use_noise, r > 0), :view_arc, e.view_arc)
end;

# Keep the intended centered 5 × 0.6 mm saved grid. The axial workspace
# automatically adds symmetric detector guard rows so peripheral voxels
# on both terminal slices retain measured cone-beam support.
recon_opts = BS.ReconOptions(
    matrix_size = (512, 512, 5),
    fov_cm = 35.0,
    z_cm = 0.30,
);

"""
    ufc_detected_spectrum(protocol) -> (e, w_eta)

Display helper: tube spectrum × flat filtration × protocol filters
(IPEM, absolute flux) with the src UFC MC η(E) folded in — what the
detector actually integrates (centered ray, no bowtie).
"""
function ufc_detected_spectrum(protocol)
    e, w = BS.resolve_source_spectrum_without_bowtie(
        force_opts, protocol; scanner = scanner,
    )
    return e, Float64.(w) .* BS.get_ufc_mc_efficiency.(e)
end;

# Tube-B view of the phantom: origin shifted by the real −0.88 mm detector
# z-offset (no-op for the z-invariant Gammex; see §6 md).
phantom_b = BS.Phantom(
    phantom.mask,
    phantom.materials,
    phantom.voxel_size,
    (phantom.origin[1], phantom.origin[2], phantom.origin[3] - 0.088),
    phantom.extent,
);

"""
    acquire(draws) -> Dict{Int, NamedTuple}

The DE acquisition of the Gammex (`EXPOSURES`, both tubes) for every draw in `draws`: draw 0 the
noise-free expectation, draw `r ≥ 1` a noisy realisation with the seeds `realization_seed(e, r)`.
Each tube is projected once, by the first draw, and every other draw re-uses that projection
(`keep_projection`, `projection`), only its noise drawn. Each entry holds the two log-transmission
channels, the spectral basis (shared by every draw), the geometry and the wall time per tube.
"""
function acquire(draws)
    parts = map(EXPOSURES) do e
        ph = e.tube === :b ? phantom_b : phantom
        projection = nothing
        channels = Dict{Int, Array{Float32, 3}}()
        seconds = Dict{Int, Float64}()
        model = nothing
        for r in draws
            opts = opts_for(e, r)
            t = time()
            ws = BS.create_workspace(scanner, e.protocol, opts, recon_opts, ph)
            try
                result = BS.simulate!(ws, ph, e.protocol, opts; report_dose = false,
                    keep_projection = projection === nothing, projection)
                projection === nothing && (projection = result.projection)
                channels[r] = Array(ws.sinogram)
                if model === nothing
                    air = ws.bowtie_air_reference === nothing ?
                        ones(Float32, ws.geom.n_cols, ws.geom.n_rows) :
                        Float32.(Array(ws.bowtie_air_reference))
                    I0 = BS.compute_detector_I0(ws.geom, e.protocol, sum(ws.weights)) * Float64(ws.η_eff)
                    energies, response = BS.resolve_source_spectrum_full(
                        opts, e.protocol; scanner = scanner, geom = ws.geom)
                    model = (geom = ws.geom, I0_ray = Float32.(I0 .* air),
                        energies = Float64.(energies), response = Float32.(response))
                end
            finally
                BS.release_backend!(ws)
            end
            seconds[r] = time() - t
        end
        projection = nothing
        (; channels, seconds, model...)
    end
    GC.gc(true)
    basis = BS.spectral_basis_from_acquisitions(acquisitions = [
        (energies = p.energies, response = p.response, I0_ray = p.I0_ray) for p in parts])
    Dict(r => (draw = r, channels = [p.channels[r] for p in parts], basis = basis,
               geom = first(parts).geom, seconds = [p.seconds[r] for p in parts])
         for r in draws)
end;

de_acq = acquire((DE_DRAWS.reference, DE_DRAWS.measured, DE_DRAWS.calibration));

# The measured draw: every result below comes from it.
de_measured = de_acq[DE_DRAWS.measured];

"""
    ufc_bhc_calibration(protocol, geom)

η-aware per-tube BHC: bowtie-hardened per-column spectrum × UFC η(E) →
per-column two-material polynomial.  Returns `(model, μ_water, ref_E_keV)`.
"""
function ufc_bhc_calibration(protocol, geom)
    # resolve_source_spectrum_full folds bowtie AND the src UFC η(E)
    # (via the same build_physics_config the forward model used).
    e, ŵ = BS.resolve_source_spectrum_full(
        force_opts, protocol; scanner = scanner, geom = geom,
    )
    e2, w_col = BS.bhc_spectrum_per_column(e, ŵ)          # [n_E, n_col]

    # Single mono-equivalent target = mean energy of the η-folded mean spectrum
    w_mean = vec(sum(w_col; dims = 2)) ./ size(w_col, 2)
    ref_E = sum(e2 .* w_mean) / sum(w_mean)

    # Parameter-free water BHC from the custom UFC-η per-column spectrum — zero
    # segmentation thresholds (the two-material bone pass is deprecated:
    # its 450–600 HU window misclassified dense iodine as bone).
    model = BS.calibrate_bhc_water(
        e2, w_col;
        reference_energy_keV = ref_E,
    )
    return (model = model, μ_water = model.μ_water_ref, ref_E_keV = model.reference_energy_keV)
end;

"""
    ufc_poly_recon(sino_cpu, geom, bhc) -> Array{Float32, 3}

Doctrine correction stack for one tube: knobless water sino-BHC → FDK → HU.
"""
function ufc_poly_recon(sino_cpu, geom, bhc)
    matrix_size = recon_opts.matrix_size

    sino_gpu = to_gpu(sino_cpu)
    # Knobless water BHC: one sinogram-domain pass, no recon round-trip.
    sino_bhc = BS.apply_bhc_water(sino_gpu, bhc.model)
    sino_gpu = sino_bhc

    ws_fdk = BS.create_fdk_recon_workspace(sino_gpu, geom, matrix_size)
    recon_μ = BS.reconstruct!(ws_fdk, sino_gpu, geom)

    # (image-domain BHC + applied cupping removed — deprecated; cupping is a
    #  QA metric via measure_radial_cupping)
    hu = Float32.(BS.to_hounsfield(Array(recon_μ); μ_water = bhc.μ_water))

    ws_fdk = nothing; sino_gpu = nothing; recon_μ = nothing
    GC.gc(true)
    return hu
end;

bhc_low = ufc_bhc_calibration(protocol_low, de_measured.geom);

bhc_high = ufc_bhc_calibration(protocol_high, de_measured.geom);

hu_tube = (
    low = ufc_poly_recon(de_measured.channels[1], de_measured.geom, bhc_low),
    high = ufc_poly_recon(de_measured.channels[2], de_measured.geom, bhc_high),
);

# Siemens linear mixed image: M = w·I_low + (1−w)·I_high (image domain,
# Yu 2009).  w = 0.5 is the common Sn150-pair default (0.3–0.7 clinical).
MIX_W_LOW = 0.5f0;

hu_mixed = MIX_W_LOW .* hu_tube.low .+ (1.0f0 - MIX_W_LOW) .* hu_tube.high;

poly_water_stats = let
    ERODE_PX = 12.0
    mask_2d_raw = phantom_cpu.mask[:, :, size(phantom_cpu.mask, 3) ÷ 2]
    sw_bool = BS.erode_mask_2d(
        mask_2d_raw .== UInt8(BS.REGION_SOLID_WATER); erode_px = ERODE_PX,
    )
    sw_idx = findall(sw_bool)
    n_z = size(hu_mixed, 3)

    function _stats(vol)
        vals = Float64[Float64(vol[ci, z]) for z in 1:n_z, ci in sw_idx]
        (mean = mean(vals), std = std(vals), n = length(vals))
    end
    stats = (
        low = _stats(hu_tube.low),
        high = _stats(hu_tube.high),
        mixed = _stats(hu_mixed),
    )
    for (tag, s) in pairs(stats)
        @info "[poly · UFC] $(tag) SW ROI: ⟨HU⟩ = $(round(s.mean, digits = 2)), σ = $(round(s.std, digits = 2)) HU (n = $(s.n))"
    end
    (stats..., mask_2d = collect(sw_bool))
end;

basis = de_measured.basis;

# basis-spectral-denoising's HYPR_CHAIN: ProjectionHYPR(3 × 3 Box, view_stride = 2) and ImageHYPR()
HYPR_CHAIN = BS.SpectralHYPR();

VMI_CHAIN = (method = :nchannel, controls = BS.NChannelControls(), use_tlbf = false, antialias = true);

# One window for the composite and its complement alike: no fitted Force pair exists (see above).
FORCE_FILTER = BS.SoftFilter();

# `pair_basis` and `composite_energy` (basis-spectral-denoising's `PAIR_BASIS`), from the calibration draw
calibration = let a = de_acq[DE_DRAWS.calibration]
    decompose(denoiser) = BS.vmi_pipeline(; channels = a.channels, basis = a.basis, geom = a.geom,
        to_backend = to_gpu, matrix_size = recon_opts.matrix_size, vmi_energies = Tuple(de_vmi_energies),
        denoiser, keep_sinograms = true, VMI_CHAIN..., use_acnr = false).sinograms
    sp(d; kw...) = BS.spectral_pair(d.water, d.iodine, a.geom, recon_opts.matrix_size;
        filter = BS.SoftFilter(), antialias = VMI_CHAIN.antialias, to_backend = to_gpu, kw...)
    x = sp(decompose(nothing))                                   # the plain decomposition
    c = sp(decompose(BS.SpectralHYPR(image = nothing));          # after the projection HYPR
        basis = (Estar = x.Estar, β = x.β))
    (pair_basis = (Estar = x.Estar, β = x.β), composite_energy = c.Estar)
end;

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

de_vmi = BS.vmi_pipeline(;
    channels = de_measured.channels,
    basis,
    geom = de_measured.geom,
    to_backend = to_gpu,
    matrix_size = recon_opts.matrix_size,
    vmi_energies = Tuple(de_vmi_energies),
    denoiser = HYPR_CHAIN,
    fbp_filter = FORCE_FILTER,
    pair_basis = calibration.pair_basis,
    composite_energy = calibration.composite_energy,
    keep_sinograms = true,
    VMI_CHAIN...,
);

# The same draw and reconstruction without the denoising (no HYPR, no ACNR), and the noise-free
# reference (draw 0, the same projection), reconstructed alike.
de_none, de_ref = (BS.vmi_pipeline(;
        channels = de_acq[r].channels, basis, geom = de_measured.geom, to_backend = to_gpu,
        matrix_size = recon_opts.matrix_size, vmi_energies = Tuple(de_vmi_energies),
        denoiser = nothing, use_acnr = false, fbp_filter = FORCE_FILTER,
        pair_basis = calibration.pair_basis, composite_energy = calibration.composite_energy,
        VMI_CHAIN...,
    ) for r in (DE_DRAWS.measured, DE_DRAWS.reference));

solid_water_basis = let
    ERODE_PX = 12.0

    mask_2d_raw = phantom_cpu.mask[:, :, size(phantom_cpu.mask, 3) ÷ 2]
    sw_bool_raw = (mask_2d_raw .== UInt8(BS.REGION_SOLID_WATER))
    sw_bool = BS.erode_mask_2d(sw_bool_raw; erode_px = ERODE_PX)

    n_raw = count(sw_bool_raw); n_eroded = count(sw_bool)
    n_eroded == 0 && error(
        "solid_water_basis: deep erosion (σ = $(ERODE_PX) px) wiped out the SW " *
            "ROI (raw count = $(n_raw)).  Reduce erode_px or check phantom mask."
    )

    sw_idx = findall(sw_bool)
    _mean(vol) = mean(vol[ci, z] for z in axes(vol, 3) for ci in sw_idx)

    c_w = Float64(_mean(de_vmi.images.water))
    c_i = Float64(_mean(de_vmi.images.iodine))
    @info "solid_water_basis: ⟨c_water⟩_SW = $(round(c_w, digits = 4)) g/cm³, " *
        "⟨c_iodine⟩_SW = $(round(1000c_i, digits = 3)) mg/mL"

    (
        c_water = c_w, c_iodine = c_i, n_voxels = length(sw_idx) * size(de_vmi.images.water, 3),
        mask_2d = collect(sw_bool),
    )
end;

# The VMI stack as one volume per energy (keV → (nx, ny, nz) HU).
vmi_HU_final = Dict(
    Float64(E) => de_vmi.vmis[:, :, :, k] for (k, E) in pairs(de_vmi.energies)
);

ROD_LABELS = (
    Ca = (UInt8(10), UInt8(11), UInt8(12), UInt8(13), UInt8(14), UInt8(15), UInt8(16)),
    I = (UInt8(20), UInt8(21), UInt8(22), UInt8(23), UInt8(24), UInt8(25), UInt8(26)),
);

ROD_NAMES = (
    Ca = ("50 mg/mL", "100 mg/mL", "200 mg/mL", "300 mg/mL", "400 mg/mL", "500 mg/mL", "600 mg/mL"),
    I = ("2.0 mg/mL", "2.5 mg/mL", "5.0 mg/mL", "7.5 mg/mL", "10.0 mg/mL", "15.0 mg/mL", "20.0 mg/mL"),
);

rod_data = let
    materials = phantom_cpu.materials
    mask_2d = phantom_cpu.mask[:, :, size(phantom_cpu.mask, 3) ÷ 2]
    nx, ny = size(mask_2d)
    ROI_RADIUS_PX = 8

    function rod_centroid(label::UInt8)
        idx = findall(==(label), mask_2d)
        isempty(idx) && error("rod_centroid: no voxels with label $label")
        cx = sum(ci -> Float64(ci[1]), idx) / length(idx)
        cy = sum(ci -> Float64(ci[2]), idx) / length(idx)
        return (cx, cy)
    end

    function rod_roi_mask(label::UInt8)
        cx, cy = rod_centroid(label)
        i_lo = max(1, floor(Int, cx - ROI_RADIUS_PX))
        i_hi = min(nx, ceil(Int, cx + ROI_RADIUS_PX))
        j_lo = max(1, floor(Int, cy - ROI_RADIUS_PX))
        j_hi = min(ny, ceil(Int, cy + ROI_RADIUS_PX))
        roi = CartesianIndex{2}[]
        r² = Float64(ROI_RADIUS_PX)^2
        for j in j_lo:j_hi, i in i_lo:i_hi
            ((i - cx)^2 + (j - cy)^2) ≤ r² && push!(roi, CartesianIndex(i, j))
        end
        return roi
    end

    rod_rois = Dict(
        lab => rod_roi_mask(lab)
            for lab in vcat(collect(ROD_LABELS.Ca), collect(ROD_LABELS.I))
    )

    μ_water_E = Dict(
        E => BS.compute_μ_at_energy(BS.XA.Materials.water, E)
            for E in de_vmi_energies
    )

    function theoretical_hu(material, E::Float64)
        μ = BS.compute_μ_at_energy(material, E)
        return 1000.0 * (μ - μ_water_E[E]) / μ_water_E[E]
    end

    function measured_hu(vmi_vol, label::UInt8)
        roi = rod_rois[label]
        s = 0.0; n = 0
        for z in 1:size(vmi_vol, 3), ci in roi
            s += vmi_vol[ci, z]; n += 1
        end
        return s / n
    end

    out = Dict{Symbol, NamedTuple}()
    for group in (:Ca, :I)
        labels = ROD_LABELS[group]
        n_rods = length(labels)
        n_E = length(de_vmi_energies)
        meas = zeros(Float64, n_rods, n_E)
        theo = zeros(Float64, n_rods, n_E)
        for (i, lab) in pairs(labels)
            mat = materials[Int(lab) + 1]   # mask_value + 1
            for (j, E) in pairs(de_vmi_energies)
                meas[i, j] = measured_hu(vmi_HU_final[E], lab)
                theo[i, j] = theoretical_hu(mat, E)
            end
        end
        out[group] = (
            labels = labels, names = ROD_NAMES[group],
            measured = meas, theoretical = theo,
        )
    end
    out
end;

const WATER_NOISE_ROI_RADIUS_PX = 12;   # ≈8.2 mm at 0.683 mm/px

water_noise_roi = let
    nx_r, ny_r, nz_r = size(de_vmi.images.water)
    cx = nx_r ÷ 2 + 1
    cy = ny_r ÷ 2 + 1

    roi_bool = falses(nx_r, ny_r)
    r² = Float64(WATER_NOISE_ROI_RADIUS_PX)^2
    @inbounds for j in 1:ny_r, i in 1:nx_r
        ((i - cx)^2 + (j - cy)^2) ≤ r² && (roi_bool[i, j] = true)
    end

    n_vox = count(roi_bool)
    @info "water_noise_roi: center = ($(cx), $(cy)), radius = $(WATER_NOISE_ROI_RADIUS_PX) px, " *
        "$(n_vox) vx × $(nz_r) z = $(n_vox * nz_r) total"

    (
        center_xy = (Float64(cx), Float64(cy)), mask_2d = roi_bool,
        n_voxels = n_vox, n_total = n_vox * nz_r,
    )
end;

vmi_noise_by_keV = let
    roi_idx = findall(water_noise_roi.mask_2d)
    nz_r = size(de_vmi.vmis, 3)
    vals(run, k) = Float64[Float64(run.vmis[ci, z, k]) for z in 1:nz_r, ci in roi_idx]

    out = Dict{Float64, NamedTuple}()
    for (k, E) in enumerate(de_vmi_energies)
        p, n, r = vals(de_vmi, k), vals(de_none, k), vals(de_ref, k)
        out[E] = (mean = mean(p), reference = mean(r), std = std(p .- r), std_none = std(n .- r),
            n = length(p))
        @info "water ROI @ $(Int(E)) keV: ⟨HU⟩ = $(round(mean(p), digits = 2)) " *
            "(reference $(round(mean(r), digits = 2))), σ = $(round(std(p .- r), digits = 2)) HU " *
            "(no denoising $(round(std(n .- r), digits = 2)) HU, n = $(length(p)))"
    end
    out
end;