#!/usr/bin/env Rscript --no-save --no-restore

# check if rpyANTs is configured
if (!isTRUE(rpyANTs::ants_available())) {
  rpyANTs::install_ants()
}

work_path <- normalizePath("{{ work_path }}", winslash = "/")
ct_path0 <- normalizePath("{{ ct_path }}", mustWork = TRUE, winslash = "/")
mri_path0 <- normalizePath("{{ mri_path }}", mustWork = TRUE, winslash = "/")
aff_metric <- "{{ aff_metric }}"
syn_metric <- "{{ syn_metric }}"
reg_type <- "{{ reg_type }}"
verbose <- TRUE

# DIPSAUS DEBUG START
# work_path <- "/Users/dipterix/Dropbox (PENN Neurotrauma)/RAVE/Samples/raw/PAV006/rave-imaging"
# ct_path0 = "/Users/dipterix/Dropbox (PENN Neurotrauma)/RAVE/Samples/raw/PAV006/rave-imaging/derivative/CT_RAW.nii.gz"
# mri_path0 = "/Users/dipterix/Dropbox (PENN Neurotrauma)/RAVE/Samples/raw/PAV006/rave-imaging/derivative/MRI_RAW.nii"
# reg_type <- "rigid"
# interp <- "cubic"
# verbose <- TRUE


coreg_path <- file.path(work_path, "coregistration")

deriv_path <- file.path(work_path, "derivative")
if (!dir.exists(deriv_path)) {
  dir.create(deriv_path, showWarnings = FALSE, recursive = TRUE)
}

# Write down configuration files
outputs <- list(
  ct_in_t1 = list(
    type = "nifti",
    path = "./coregistration/ct_in_t1.nii.gz",
    backup = "./derivative/ct_in_t1.nii.gz",
    comment = "re-sampled CT; the resolution is the same as reference MRI"
  )
)

is_linear <- grepl("(rigid|affine)", reg_type, ignore.case = TRUE)
if ( is_linear ) {
  outputs$CT_IJK_to_MR_RAS <- list(
    type = "transform",
    dimension = "4x4",
    path = "./coregistration/CT_IJK_to_MR_RAS.txt",
    backup = "./coregistration/transform-ctIJK2mrRAS.txt",
    transform_from = list(
      volume = "input_image",
      coordinate_system = "IJK"
    ),
    transform_to = list(
      volume = "reference_image",
      space = "scanner",
      coordinate_system = "RAS"
    ),
    comment = "From voxel IJK coordinate to MRI scanner RAS coordinate"
  )
  outputs$CT_RAS_to_MR_RAS <- list(
    type = "transform",
    dimension = "4x4",
    path = "./coregistration/CT_RAS_to_MR_RAS.txt",
    backup = "./coregistration/transform-ctRAS2mrRAS.txt",
    transform_from = list(
      volume = "input_image",
      space = "scanner (CT)",
      coordinate_system = "RAS"
    ),
    transform_to = list(
      volume = "reference_image",
      space = "scanner",
      coordinate_system = "RAS"
    ),
    comment = "From CT scanner RAS coordinate to MRI scanner RAS coordinate"
  )
} else {
  outputs$forward_control_points <- list(
    type = "nifti",
    path = "./coregistration/ANTs_1Warp.nii.gz",
    comment = "Non-linear transformation control points from CT to MRI"
  )
  outputs$forward_control_points <- list(
    type = "nifti",
    path = "./coregistration/ANTs_1InverseWarp.nii.gz",
    comment = "Non-linear transformation control points from MRI to CT"
  )
}

ieegio::io_write_yaml(
  con = file.path(deriv_path, "conf-coregistration.yaml"),
  list(
    `Heads up` = "Do NOT edit this file",
    profile = "CT coregister to MRI",
    work_path = work_path,
    timestamp = strftime(Sys.time(), "%a %b %d %H:%M:%S %Z %Y", tz = "UTC"),
    command = list(
      comment = "This is R function call",
      execute = "ravecore::cmd_run_ants_coreg",
      arguments = list(
        reg_type = reg_type,
        aff_metric = aff_metric,
        syn_metric = syn_metric
      )
    ),
    input_image = list(
      type = "nifti",
      path = ct_path0,
      backup = c(
        "./derivative/CT_RAW.nii.gz",
        "./coregistration/CT_RAW.nii.gz"
      ),
      comment = "original CT image file"
    ),
    reference_image = list(
      type = "nifti",
      path = mri_path0,
      backup = c(
        "./coregistration/MRI_reference.nii.gz"
      ),
      comment = "Reference MR image file, the CT is aligned to this reference image"
    ),
    outputs = outputs
  )
)


ravecore::ants_coreg(
  ct_path = ct_path0,
  mri_path = mri_path0,
  coreg_path = coreg_path,
  reg_type = reg_type,
  aff_metric = aff_metric,
  syn_metric = syn_metric,
  verbose = verbose
)

# Finalize
file.copy(
  from = file.path(coreg_path, "ct_in_t1.nii.gz"),
  to = file.path(deriv_path, "ct_in_t1.nii.gz"),
  overwrite = TRUE,
  recursive = FALSE,
  copy.mode = TRUE,
  copy.date = TRUE
)
file.copy(
  from = file.path(coreg_path, "CT_RAW.nii.gz"),
  to = file.path(deriv_path, "CT_RAW.nii.gz"),
  overwrite = TRUE,
  recursive = FALSE,
  copy.mode = TRUE,
  copy.date = TRUE
)
file.copy(
  from = file.path(coreg_path, "MRI_reference.nii.gz"),
  to = file.path(deriv_path, "MRI_reference.nii.gz"),
  overwrite = TRUE,
  recursive = FALSE,
  copy.mode = TRUE,
  copy.date = TRUE
)

if ( is_linear ) {
  file.copy(
    from = file.path(coreg_path, "CT_IJK_to_MR_RAS.txt"),
    to = file.path(deriv_path, "transform-ctIJK2mrRAS.txt"),
    overwrite = TRUE,
    recursive = FALSE,
    copy.mode = TRUE,
    copy.date = TRUE
  )

  file.copy(
    from = file.path(coreg_path, "CT_RAS_to_MR_RAS.txt"),
    to = file.path(deriv_path, "transform-ctRAS2mrRAS.txt"),
    overwrite = TRUE,
    recursive = FALSE,
    copy.mode = TRUE,
    copy.date = TRUE
  )
}



## END OF RAVE Script: CT MR coregistration via rpyANTs (essentially ANTs)
