XCAT UHR → CT: Grid Mapping and the Recon-Affine Round-Trip

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

using Markdown: @md_str, Markdown

using Statistics: mean, std

using Unitful: @u_str

import BasisSimulator as BS

# import CairoMakie as Mke
import WasmMakie 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

const XCAT_DIR = get(
    ENV, "BASISSIM_XCAT_DIR",
    joinpath(@__DIR__, "data", "xcat")
)

const PHANTOM_PATH = joinpath(
    XCAT_DIR,
    "vmale_50_1600x1400x500_8bit_little_endian_act_1.bin"
)

const HAS_XCAT = isfile(PHANTOM_PATH)

function load_xcat_bin(
        filepath::AbstractString;
        cols::Int = 1600, rows::Int = 1400, slices::Int = 500,
    )
    expected = cols * rows * slices * sizeof(UInt8)
    actual = filesize(filepath)
    actual == expected ||
        error("XCAT file size mismatch: expected $(expected) bytes, got $(actual)")

    data = Vector{UInt8}(undef, cols * rows * slices)
    open(filepath, "r") do io
        read!(io, data)
    end

    phantom = reshape(data, (cols, rows, slices))
    return reverse(phantom; dims = (2, 3))
end

"""Nearest-neighbor 3D downsample by an integer factor — preserves labels."""
function downsample_labeled(phantom::AbstractArray{T, 3}, factor::Int) where {T}
    factor == 1 && return phantom
    nx, ny, nz = size(phantom) .÷ factor
    out = similar(phantom, (nx, ny, nz))
    @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
        ii = (i - 1) * factor + factor ÷ 2 + 1
        jj = (j - 1) * factor + factor ÷ 2 + 1
        kk = (k - 1) * factor + factor ÷ 2 + 1
        out[i, j, k] = phantom[ii, jj, kk]
    end
    return out
end

const DOWNSAMPLE_FACTOR = 2

const VOXEL_SIZE_CM = (
    0.02 * DOWNSAMPLE_FACTOR,    # 0.4 mm at DS=2
    0.02 * DOWNSAMPLE_FACTOR,
    0.02 * DOWNSAMPLE_FACTOR,
);

phantom_full_uhr = HAS_XCAT ?
    downsample_labeled(load_xcat_bin(PHANTOM_PATH), DOWNSAMPLE_FACTOR) :
    nothing;

import XLSX

const MATERIAL_XLSX_PATH = joinpath(
    XCAT_DIR, "Material_Spreadsheets",
    "vmale_50_materials_heart_high_contrast.xlsx",
)

const _ATOMIC_MASSES = Dict(
    1 => 1.008, 6 => 12.011, 7 => 14.007, 8 => 15.999, 11 => 22.99, 12 => 24.305,
    15 => 30.974, 16 => 32.06, 17 => 35.45, 19 => 39.098, 20 => 40.078, 26 => 55.845, 53 => 126.904,
)

const _I_VALUES_EV = Dict(
    1 => 19.2, 6 => 81.0, 7 => 82.0, 8 => 95.0, 11 => 149.0, 12 => 156.0,
    15 => 173.0, 16 => 180.0, 17 => 174.0, 19 => 190.0, 20 => 191.0, 26 => 286.0, 53 => 491.0,
)

function compute_ZA_ratio(comp::Dict{Int, Float64})
    Z_sum = sum(w * Z / get(_ATOMIC_MASSES, Z, Float64(Z) * 2) for (Z, w) in comp)
    A_sum = sum(values(comp))
    return Z_sum / A_sum
end

function compute_mean_excitation_energy(comp::Dict{Int, Float64})
    log_I_sum = 0.0
    Z_A_sum = 0.0
    for (Z, w) in comp
        A = get(_ATOMIC_MASSES, Z, Float64(Z) * 2)
        I = get(_I_VALUES_EV, Z, 10.0 * Z)
        Z_A = w * Z / A
        log_I_sum += Z_A * log(I)
        Z_A_sum += Z_A
    end
    return exp(log_I_sum / Z_A_sum) * u"eV"
end

function load_materials_from_xlsx(xlsx_path::AbstractString)
    sheet = XLSX.readxlsx(xlsx_path)["Sheet1"]
    data = sheet["A2:P34"]
    out = Dict{Int, BS.XA.Material}()

    Z_cols = (1, 6, 7, 8, 11, 12, 15, 16, 17, 19, 20, 26, 53)

    for r in 1:size(data, 1)
        name = data[r, 1]
        oid = data[r, 16]
        ρ = data[r, 15]
        (name === nothing || oid === nothing || ρ === nothing) && continue

        comp = Dict{Int, Float64}()
        for (k, Z) in enumerate(Z_cols)
            v = data[r, k + 1]
            v isa Number && v > 0 && (comp[Z] = Float64(v))
        end
        isempty(comp) && continue

        out[Int(oid)] = BS.XA.Material(
            String(name),
            compute_ZA_ratio(comp),
            compute_mean_excitation_energy(comp),
            Float64(ρ) * u"g/cm^3",
            comp,
        )
    end
    return out
end

materials_full = phantom_full_uhr === nothing ? nothing : let
        base = load_materials_from_xlsx(MATERIAL_XLSX_PATH)
        for l in unique(phantom_full_uhr)
            haskey(base, Int(l)) || (base[Int(l)] = BS.XA.Materials.water)
    end
        base
end;

heart_label_ids = materials_full === nothing ? nothing : let
        pattern = r"heart|atrium|ventric|coronary|aorta"i
        ids = sort(
            UInt8[
                UInt8(oid) for (oid, mat) in materials_full
                if occursin(pattern, mat.name)
            ]
        )
        @info "[cardiac bbox] $(length(ids)) organ labels match: $(ids)"
        for oid in ids
            @info "    label $(Int(oid)) → $(materials_full[Int(oid)].name)"
    end
        ids
end;

heart_bbox = (phantom_full_uhr === nothing || heart_label_ids === nothing) ? nothing : let
        is_heart = falses(256)
        for oid in heart_label_ids
            is_heart[Int(oid) + 1] = true
    end

        nx, ny, nz = size(phantom_full_uhr)
        i_lo, i_hi = nx + 1, 0
        j_lo, j_hi = ny + 1, 0
        k_lo, k_hi = nz + 1, 0
        n_voxels = 0
        @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
            if is_heart[Int(phantom_full_uhr[i, j, k]) + 1]
                i < i_lo && (i_lo = i)
                i > i_hi && (i_hi = i)
                j < j_lo && (j_lo = j)
                j > j_hi && (j_hi = j)
                k < k_lo && (k_lo = k)
                k > k_hi && (k_hi = k)
                n_voxels += 1
        end
    end
        n_voxels == 0 && error("[cardiac bbox] no heart-labeled voxels found in phantom")

        # Pad ~1 cm in every direction.
        pad_vox_x = round(Int, 1.0 / VOXEL_SIZE_CM[1])
        pad_vox_y = round(Int, 1.0 / VOXEL_SIZE_CM[2])
        pad_vox_z = round(Int, 1.0 / VOXEL_SIZE_CM[3])

        i_lo = max(1, i_lo - pad_vox_x);  i_hi = min(nx, i_hi + pad_vox_x)
        j_lo = max(1, j_lo - pad_vox_y);  j_hi = min(ny, j_hi + pad_vox_y)
        k_lo = max(1, k_lo - pad_vox_z);  k_hi = min(nz, k_hi + pad_vox_z)

        @info "[cardiac bbox] tight bbox + 1 cm pad:"
        @info "  voxel range = ($(i_lo):$(i_hi), $(j_lo):$(j_hi), $(k_lo):$(k_hi))"
        @info "  size  = $(i_hi - i_lo + 1) × $(j_hi - j_lo + 1) × $(k_hi - k_lo + 1) voxels"
        @info "  extent = $(round.((i_hi - i_lo + 1, j_hi - j_lo + 1, k_hi - k_lo + 1) .* VOXEL_SIZE_CM, digits = 2)) cm"
        @info "  cardiac voxels (pre-pad)  = $(n_voxels)"

        (i_lo = i_lo, i_hi = i_hi, j_lo = j_lo, j_hi = j_hi, k_lo = k_lo, k_hi = k_hi)
end;

phantom_cropped = (phantom_full_uhr === nothing || heart_bbox === nothing) ? nothing : let
        b = heart_bbox
        out = phantom_full_uhr[b.i_lo:b.i_hi, b.j_lo:b.j_hi, b.k_lo:b.k_hi]
        full_voxels = length(phantom_full_uhr)
        cropped_voxels = length(out)
        @info "[crop] $(round(full_voxels / 1.0e6, digits = 1))M voxels → $(round(cropped_voxels / 1.0e6, digits = 1))M voxels  ($(round(100 * cropped_voxels / full_voxels, digits = 1))% kept)"
        @info "[crop] memory: $(round(sizeof(phantom_full_uhr) / 1024^2, digits = 1)) MB → $(round(sizeof(out) / 1024^2, digits = 1)) MB"
        @info "[crop] forward-projection ray count drops by the same ratio — $(round(full_voxels / cropped_voxels, digits = 1))× faster simulate!"
        out
end;

materials_cropped = (phantom_cropped === nothing || materials_full === nothing) ? nothing : let
        base = copy(materials_full)
        for l in unique(phantom_cropped)
            haskey(base, Int(l)) || (base[Int(l)] = BS.XA.Materials.water)
    end
        base
end;

phantom = (phantom_cropped === nothing || materials_cropped === nothing) ? nothing :
    BS.Phantom(to_gpu(phantom_cropped), materials_cropped, VOXEL_SIZE_CM);

phantom_cpu = (phantom_cropped === nothing || materials_cropped === nothing) ? nothing :
    BS.Phantom(phantom_cropped, materials_cropped, VOXEL_SIZE_CM);

scanner = BS.Scanner(
    source_to_isocenter = 625.6,
    source_to_detector = 1100.0,
    detector_rows = 256,
    detector_cols = 834,
    detector_row_size = 0.625,
    detector_col_size = 0.6,
    focal_spot_width = 1.0,
    focal_spot_length = 1.0,
    target_angle = 10.0,
    flat_filter_material = :aluminum,
    flat_filter_thickness = 2.5,
    bowtie_filter = :ge_revolution_large,
    detector_material = :lumex,
    detector_depth = 3.0,
    fill_factor_row = 0.9,
    fill_factor_col = 0.9,
    # DAS/electronic noise enters the counts before the log transform, so it
    # propagates through reconstruction (see `add_system_noise_floor!` docstring).
    electronic_noise = 3500.0,   # e⁻ — clinical GE Apex Elite DAS readout noise
    detection_gain = 10.0,
);

protocol = BS.CTProtocol(
    kVp = 120,
    mA = 250.0,
    views = 500,
    rotation_time = 1.0,
    collimation_mm = 5.0,
    additional_filters = [("Al", 4.5)],
);

sim_opts = BS.SimOptions(fidelity = :eict, seed = 1234, projector = :dd_fast);

recon_opts = let
    slice_thickness_mm = 0.625
    n_z = max(1, round(Int, protocol.collimation_mm / slice_thickness_mm))
    BS.ReconOptions(
        matrix_size = (384, 384, n_z),
        fov_cm = 14.0,
        z_cm = protocol.collimation_mm / 10.0,
    )
end;

geom_inspect = BS.CTGeometry(
    scanner;
    n_angles = protocol.views,
    fov_cm = recon_opts.fov_cm,
    z_cm = recon_opts.z_cm,
    collimation_mm = protocol.collimation_mm,
);

sim = phantom === nothing ? nothing : let
        @info "Simulating cardiac CTA: 120 kVp / 250 mA / cropped UHR phantom…"
        ws = BS.create_eict_workspace(scanner, protocol, sim_opts, recon_opts, phantom)
        BS.simulate!(ws, phantom, protocol, sim_opts)

        result = (sino = Array(ws.sinogram), geom = ws.geom)
        ws = nothing
        GC.gc(true)
        result
end;

μ_water_120 = phantom_cpu === nothing ? nothing : let
        # Phantom-aware water_path: pull the body chord straight off the
        # cropped XCAT mask, so the calibration tracks any change to the
        # crop bbox / DOWNSAMPLE_FACTOR without hardcoded cm.
        body_diameter_cm = BS.estimate_phantom_diameter_cm(
            phantom_cpu.mask, phantom_cpu.voxel_size .* 10.0,
        )
        μ = BS.compute_polychromatic_μ_water(
            sim_opts, protocol;
            scanner = scanner,
            geom = geom_inspect,
            water_path_cm = body_diameter_cm,
        )
        @info "[analytic μ_water]  120 kVp + 4.5 mm Al + $(round(body_diameter_cm, digits = 1)) cm body hardening: $(round(μ, digits = 5)) cm⁻¹"
        μ
end;

recon_HU = sim === nothing ? nothing : let
        sino_gpu = to_gpu(Float32.(sim.sino))
        ws = BS.create_fdk_recon_workspace(sino_gpu, sim.geom, recon_opts.matrix_size; filter = :standard)
        recon_μ = Array(BS.reconstruct!(ws, sino_gpu, sim.geom))
        ws = nothing; sino_gpu = nothing; GC.gc(true)

        # quantum + DAS noise already carried by the sinogram (counts domain)
        HU = Float32.(BS.to_hounsfield(recon_μ; μ_water = μ_water_120))
        HU
end;

gt_resampled_nn = (phantom_cpu === nothing || sim === nothing) ? nothing :
    BS.resample_to_recon(phantom_cpu, sim.geom, recon_opts.matrix_size; method = :nearest);

gt_resampled_lin = (phantom_cpu === nothing || sim === nothing) ? nothing :
    BS.resample_to_recon(phantom_cpu, sim.geom, recon_opts.matrix_size; method = :linear);

# Fractional cardiac coverage: build a *binary* cardiac mask, then resample
# `:linear`.  Trilinear on a multi-label integer mask (gt_resampled_lin
# above) arithmetically-mixes label IDs and isn't physically meaningful;
# trilinear on a 0/1 mask gives true partial-volume fractions ∈ [0, 1].
cardiac_coverage_lin = (
        phantom_cpu === nothing || sim === nothing ||
        heart_label_ids === nothing
    ) ? nothing : let
        binary_mask = zeros(UInt8, size(phantom_cropped))
        for oid in heart_label_ids
            binary_mask[phantom_cropped .== oid] .= 0x01
    end
        binary_phantom = BS.Phantom(
            binary_mask,
            Dict(0 => BS.XA.Materials.water, 1 => BS.XA.Materials.water),
            VOXEL_SIZE_CM,
        )
        BS.resample_to_recon(binary_phantom, sim.geom, recon_opts.matrix_size; method = :linear)
end;

heart_bbox_tall = (phantom_full_uhr === nothing || heart_bbox === nothing) ? nothing : let
    nz = size(phantom_full_uhr, 3)
    extra_z = round(Int, 4.0 / VOXEL_SIZE_CM[3])   # +4 cm each side beyond the cardiac bbox
    b = heart_bbox
    tall = (i_lo = b.i_lo, i_hi = b.i_hi, j_lo = b.j_lo, j_hi = b.j_hi,
            k_lo = max(1, b.k_lo - extra_z), k_hi = min(nz, b.k_hi + extra_z))
    @info "[Scan B tall bbox] z range $(b.k_lo):$(b.k_hi) → $(tall.k_lo):$(tall.k_hi)  ($(round((tall.k_hi - tall.k_lo + 1) * VOXEL_SIZE_CM[3], digits = 1)) cm z extent)"
    tall
end;

phantom_cropped_tall = (phantom_full_uhr === nothing || heart_bbox_tall === nothing) ? nothing : let
    b = heart_bbox_tall
    phantom_full_uhr[b.i_lo:b.i_hi, b.j_lo:b.j_hi, b.k_lo:b.k_hi]
end;

materials_tall = (phantom_cropped_tall === nothing || materials_full === nothing) ? nothing : let
    base = copy(materials_full)
    for l in unique(phantom_cropped_tall)
        haskey(base, Int(l)) || (base[Int(l)] = BS.XA.Materials.water)
    end
    base
end;

phantom_helical = (phantom_cropped_tall === nothing || materials_tall === nothing) ? nothing :
    BS.Phantom(to_gpu(phantom_cropped_tall), materials_tall, VOXEL_SIZE_CM);

phantom_helical_cpu = (phantom_cropped_tall === nothing || materials_tall === nothing) ? nothing :
    BS.Phantom(phantom_cropped_tall, materials_tall, VOXEL_SIZE_CM);

recon_opts_helical = let
    slice_thickness_mm = 0.625
    z_cm = 4.0                        # taller recon slab than Scan A — the helix supplies the z coverage
    n_z = max(1, round(Int, z_cm * 10 / slice_thickness_mm))   # ~64 slices to scroll
    BS.ReconOptions(matrix_size = (384, 384, n_z), fov_cm = 14.0, z_cm = z_cm)
end;

sim_helical = phantom_helical === nothing ? nothing : let
    protocol_hel = BS.CTProtocol(
        kVp = 120, mA = 250.0, views = 500, rotation_time = 1.0,
        collimation_mm = 10.0, additional_filters = [("Al", 4.5)],
        pitch = 1.0, n_rotations = 8.0,
    )
    @info "Simulating HELICAL cardiac CTA: pitch 1.0 × 8 rotations, 10 mm collimation…"
    ws = BS.create_eict_workspace(scanner, protocol_hel, sim_opts, recon_opts_helical, phantom_helical)
    BS.simulate!(ws, phantom_helical, protocol_hel, sim_opts)
    result = (sino = Array(ws.sinogram), geom = ws.geom)
    ws = nothing
    GC.gc(true)
    result
end;

recon_HU_helical = sim_helical === nothing ? nothing : let
    sino_gpu = to_gpu(Float32.(sim_helical.sino))
    # is_helical(geom) routes reconstruct! to the rebinned-WFBP path
    ws = BS.create_fdk_recon_workspace(sino_gpu, sim_helical.geom, recon_opts_helical.matrix_size; filter = :standard)
    recon_μ = Array(BS.reconstruct!(ws, sino_gpu, sim_helical.geom))
    ws = nothing; sino_gpu = nothing; GC.gc(true)
    Float32.(BS.to_hounsfield(recon_μ; μ_water = μ_water_120))
end;

gt_helical_nn = (phantom_helical_cpu === nothing || sim_helical === nothing) ? nothing :
    BS.resample_to_recon(phantom_helical_cpu, sim_helical.geom, recon_opts_helical.matrix_size; method = :nearest);

# Fractional cardiac coverage on the HELICAL grid: binary mask → `:linear`
# resample (the multi-label `:nearest` map mixes IDs under trilinear, so we
# resample a 0/1 mask to get physical partial-volume fractions).
cardiac_coverage_lin_helical = (
        phantom_cropped_tall === nothing || sim_helical === nothing ||
        heart_label_ids === nothing
    ) ? nothing : let
        binary_mask = zeros(UInt8, size(phantom_cropped_tall))
        for oid in heart_label_ids
            binary_mask[phantom_cropped_tall .== oid] .= 0x01
        end
        binary_phantom = BS.Phantom(
            binary_mask,
            Dict(0 => BS.XA.Materials.water, 1 => BS.XA.Materials.water),
            VOXEL_SIZE_CM,
        )
        BS.resample_to_recon(binary_phantom, sim_helical.geom, recon_opts_helical.matrix_size; method = :linear)
end;

# The affine audit proves the MAP is exact; this proves the RECON is registered
# TO that map — i.e. the forward projector images the phantom at the same world
# position the backprojector reconstructs it (a half-voxel or rebinning-column
# mismatch would displace the recon and stay invisible to the round-trip audit).
# Correlate recon edge-magnitude against the label-boundary map over integer
# (dx,dy) shifts; the peak shift is the real offset.  (0,0) = registered.
# (Integer resolution: a systematic ≤½-voxel shift also reads (0,0) — but that is
#  the same scale as :nearest quantization / recon PSF, not a registration bug.)
let
    reg_offset = function (recon, gt; R = 5)
        mid = size(recon, 3) ÷ 2
        hu = Float32.(recon[:, :, mid]); lab = gt[:, :, mid]
        nx, ny = size(hu)
        gr = zeros(Float32, nx, ny)          # recon edge strength
        lb = zeros(Float32, nx, ny)          # gt label-boundary map
        for j in 2:(ny - 1), i in 2:(nx - 1)
            gx = hu[i + 1, j] - hu[i - 1, j]; gy = hu[i, j + 1] - hu[i, j - 1]
            gr[i, j] = sqrt(gx * gx + gy * gy)
            (lab[i, j] != lab[i + 1, j] || lab[i, j] != lab[i, j + 1]) && (lb[i, j] = 1f0)
        end
        best = (0, 0); bs = -Inf; sc = Dict{Tuple{Int,Int},Float64}()
        for dx in -R:R, dy in -R:R
            s = 0.0
            for j in (1 + R):(ny - R), i in (1 + R):(nx - R)
                s += gr[i, j] * lb[i + dx, j + dy]
            end
            sc[(dx, dy)] = s
            s > bs && (bs = s; best = (dx, dy))
        end
        (best, sc[(0, 0)] / bs)              # best shift + how good (0,0) is vs peak
    end
    ax = (sim === nothing) ? "—" : reg_offset(recon_HU, gt_resampled_nn)
    hx = (sim_helical === nothing) ? "—" : reg_offset(recon_HU_helical, gt_helical_nn)
    Markdown.parse("""
    **Registration offset (recon edges vs label boundaries, mid-slice):**

    - axial   : best shift = $(ax isa String ? ax : ax[1]) voxels · (0,0) score = $(ax isa String ? ax : round(ax[2], digits = 3)) of peak
    - helical : best shift = $(hx isa String ? hx : hx[1]) voxels · (0,0) score = $(hx isa String ? hx : round(hx[2], digits = 3)) of peak

    `(0,0)` best shift ⇒ registered; a consistent nonzero shift ⇒ a real
    forward↔backprojector offset.  Score near 1.0 ⇒ `(0,0)` is already the peak.
    """)
end