src/ImGenerate.jl

###################
# Image generation
###################

module ImGenerate

using ColorTypes: Gray
import TestImages
# We don't really *directly* depend on QuartzImageIO. The import here is
# merely a workaround to suppress warnings when  loading TestImages.
# Something is broken in those packages.
import QuartzImageIO

using AlgTools.Util
using AlgTools.Comms
using ImageTools.Translate

using ..OpticalFlow: Image, DisplacementConstant, DisplacementFull
using ..Radon

# Added for reproducibility
import StableRNGs: StableRNG, Random
const rng = StableRNG(9182737465)

# Added for PET
import PoissonRandom: pois_rand
import Random: shuffle
import Images: center, warp
import CoordinateTransformations: recenter
import Rotations: RotMatrix
import Interpolations: Flat


##############
# Our exports
##############

export ImGen,
       OnlineData,
       imgen_square,
       imgen_shake,
       PetOnlineData,
       imgen_shepplogan_radon

##################
# Data structures
##################

struct ImGen
    f :: Function
    dim :: Tuple{Int64,Int64}
    Λ :: Float64
    dynrange :: Float64
    name :: String
end

struct OnlineData{DisplacementT}
    b_true :: Image
    b_noisy :: Image
    v :: DisplacementT
    v_true :: DisplacementT
    v_cumul_true :: DisplacementT
end

struct PetOnlineData{DisplacementT}
    b_true :: Image
    sinogram_true :: Image
    sinogram_noisy :: Image
    v :: DisplacementT
    v_true :: DisplacementT
    v_cumul_true :: DisplacementT
    theta :: DisplacementT              # theta = thetaknown, theta_cumul
    S :: Image
end

###################
# Shake generation
###################

function make_const_v(displ, sz)
    v = zeros(2, sz...)
    v[1, :, :] .= displ[1]
    v[2, :, :] .= displ[2]
    return v
end

function shake(params)
    if !haskey(params, :shaketype) || params.shaketype == :gaussian
        return () -> params.shake.*randn(rng,2)
    elseif params.shaketype == :disk
        return () -> begin
            θ = 2π*rand(rng,Float64)
            r = params.shake*√(rand(rng,Float64))
            return [r*cos(θ), r*sin(θ)]
        end
    elseif params.shaketype == :circle
        return () -> begin
            θ = 2π*rand(rng,Float64)
            r = params.shake
            return [r*cos(θ), r*sin(θ)]
        end        
    else
        error("Unknown shaketype $(params.shaketype)")
    end 
end

pixelwise = (shakefn, sz) -> () -> make_const_u(shakefn(), sz)


function rotatebytheta(params)
    r = params.rotation_factor*(2*rand(rng)-1)
    return r
end

function generate_radonmask(params)
    imdim = params.radondims
    sino_sparse = params.sino_sparsity
    numone = Int64(round(sino_sparse*imdim[1]*imdim[2]))
    numzero = imdim[1]*imdim[2]-numone
    A = shuffle(rng,reshape([ones(numone); zeros(numzero)],(imdim[1],imdim[2])))
    return A
end

################
# Moving square
################

function generate_square(sz,
                         :: Type{DisplacementT},
                         datachannel :: Channel{OnlineData{DisplacementT}},
                         params) where DisplacementT

    if false
        v₀ = make_const_v(0.1.*(-1, 1), sz)
        nextv = () -> v₀
    elseif DisplacementT == DisplacementFull
        nextv = pixelwise(shake(params), sz)
    elseif DisplacementT == DisplacementConstant
        nextv = shake(params)
    else
        @error "Invalid DisplacementT"
    end

    # Constant linear displacement everywhere has Jacobian determinant one
    # (modulo the boundaries which we ignore here)
    m = round(Int, sz[1]/5)
    b_orig = zeros(sz...)
    b_orig[sz[1].-(2*m:3*m), 2*m:3*m] .= 1

    v_true = nextv()
    v_cumul = copy(v_true)

    while true
        # Flow original data and add noise
        b_true = zeros(sz...)
        translate_image!(b_true, b_orig, v_cumul; threads=true)
        b = b_true .+ params.noise_level.*randn(rng,sz...)
        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
        # Pass true data to iteration routine
        data = OnlineData{DisplacementT}(b_true, b, v, v_true, v_cumul)
        if !put_unless_closed!(datachannel, data)
            return
        end
        # Next step shake
        v_true = nextv()
        v_cumul .+= v_true
    end
end

function imgen_square(sz)
    return ImGen(curry(generate_square, sz), sz, 1, 1, "square$(sz[1])x$(sz[2])")
end

################
# Shake a photo
################

function generate_shake_image(im, sz,
                              :: Type{DisplacementConstant},
                              datachannel :: Channel{OnlineData{DisplacementConstant}},
                              params :: NamedTuple)
    
    # Restart the seed to enable comparison across predictors                       
    Random.seed!(rng,9182737465)   

    nextv = shake(params)
    v_true = nextv()
    v_cumul = copy(v_true)

    while true
        # Extract subwindow of original image and add noise
        b_true = zeros(sz...)
        extract_subimage!(b_true, im, v_cumul; threads=true)
        b = b_true .+ params.noise_level.*randn(rng,sz...)
        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
        # Pass data to iteration routine
        data = OnlineData{DisplacementConstant}(b_true, b, v, v_true, v_cumul)
        if !put_unless_closed!(datachannel, data)
            return
        end
        # Next step shake
        v_true = nextv()
        v_cumul .+= v_true
    end
end

function imgen_shake(imname, sz)
    im = Float64.(Gray.(TestImages.testimage(imname)))
    dynrange = maximum(im)
    return ImGen(curry(generate_shake_image, im, sz), sz, 1, dynrange,
                 "$(imname)$(sz[1])x$(sz[2])")
end


########################################################################
# PETscan
########################################################################
function generate_radon(im, sz,
    :: Type{DisplacementConstant},
    datachannel :: Channel{PetOnlineData{DisplacementConstant}},
    params :: NamedTuple)

    # Restart the seed to enable comparison across predictors                       
    Random.seed!(rng,9182737465)  

    nextv = shake(params)
    v_true = nextv()
    v_cumul = copy(v_true)

    S_true = generate_radonmask(params)

    theta_true = rotatebytheta(params)
    theta_cumul = copy(theta_true)

    while true
        # Define the transformation matrix
        center_point = (sz[1]/2 + v_true[1], sz[2]/2 + v_true[2])
        tform = recenter(RotMatrix(theta_cumul), center_point)

        # Apply the transformation to the image using warp
        b_true = copy(warp(im, tform, axes(im), fillvalue=Flat()))

        v = v_true.*(1.0 .+ params.shake_noise_level.*randn(rng,size(v_true)...))
        theta = theta_true*(1.0 + params.rotation_noise_level.*randn(rng))

        # Generate the true and noisy sinograms
        sinogram_true = zeros(params.radondims...)
        sinogram_true .*= params.scale
        radon!(sinogram_true, b_true)
        sinogram_noisy = copy(sinogram_true)

        for i=1:params.radondims[1], j=1:params.radondims[2]
            sinogram_noisy[i, j] += pois_rand(rng,params.noise_level)
        end

        # Pass data to iteration routine
        data = PetOnlineData{DisplacementConstant}(b_true, sinogram_true, sinogram_noisy, v, v_true, v_cumul, [theta, theta_cumul], S_true)
        if !put_unless_closed!(datachannel, data)
        return
        end

        # Next step shake
        v_true = nextv()
        v_cumul .+= v_true
        # Next theta
        theta_true = rotatebytheta(params)
        theta_cumul += theta_true
        # Next sinogram mask 
        S_true = generate_radonmask(params)
    end
end


function imgen_shepplogan_radon(origsize,sz)
    im = convert(Array{Float64},TestImages.shepp_logan(origsize, highContrast=true))
    dynrange = maximum(im)
    return ImGen(curry(generate_radon, im, sz), sz, 1, dynrange, "shepplogan$(sz[1])x$(sz[2])")
end


end # Module