/***
 * See https://github.com/MarkZaidi/QuPath-Image-Alignment/blob/main/Calculate-Transforms.groovy for most up-to-date version.
 * Please link to github repo when referencing code in forums, as the code will be continually updated.
 *
 * Script to align 2 or more images in the project with different pixel sizes, using either intensities or annotations.
 * Run from any image in a project containing all sets of images that require alignment
 * Writes the affine transform to an object inside the Affine subfolder of your project folder.
 * Also grabs the detection objects from the template image. Can change this to annotation objects.
 * Usage:
 * - Load in all sets of images to be aligned. Rename file names such that the only underscore (_) in the image name
 *   separates the SlideID from stain. Example: N19-1107 30Gy M5_PANEL2.vsi
 * - Adjust the inputs specified under "Needed inputs", and run script (can run on any image, iterates over entire project)
 *  - If script errors due to alignment failing to converge, set 'align_specific' to the SlideID of the image it failed on
 *  - Set 'skip_image' to 1, rerun script to skip over the error-causing image
 *  - Set 'skip_image' to 0, and either adjust 'AutoAlignPixelSize' or draw tissue annotations on all stains of images in list
 *  - run script, verify all moving images contain a transform file located in the 'Affine' folder
 *
 * Needed inputs:
 * - registrationType : Set as "AFFINE" for translations, rotations, scaling, and sheering. Set as "RIGID" for only translations and rotations.
 * - refStain : Set to stain name of image to align all subsequent images to
 * - wsiExt : file name extension
 * - align_specific : see above, set to null for first run through
 * - AutoAlignPixelSize : downsample factor when calculating the transform. Greater values result in faster calculation, but may impact quality
 * - skip_image see above, value doesn't matter if align_specific is null
 *
 *
 * Script largely adapted from Sara McArdle's callable implementation of QuPath's Interactive Image Alignment, and Yau Mun Lim's method
 * of matching reference (static) and overlay (moving) images based on file names.
 *
 */

import qupath.lib.objects.PathCellObject
import qupath.lib.objects.PathDetectionObject
import qupath.lib.objects.PathObject
import qupath.lib.objects.PathObjects
import qupath.lib.objects.PathTileObject
import qupath.lib.objects.classes.PathClassFactory
import qupath.lib.roi.RoiTools
import qupath.lib.roi.interfaces.ROI

import java.awt.geom.AffineTransform
import javafx.scene.transform.Affine
import qupath.lib.images.servers.ImageServer

import java.awt.Graphics2D
import java.awt.Transparency
import java.awt.color.ColorSpace
import java.awt.image.BufferedImage

import org.bytedeco.opencv.global.opencv_core;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.TermCriteria;
import org.bytedeco.opencv.global.opencv_video;
import org.bytedeco.javacpp.indexer.FloatIndexer;
import org.bytedeco.javacpp.indexer.Indexer;

import qupath.lib.gui.dialogs.Dialogs;
import qupath.lib.images.servers.PixelCalibration;

import qupath.lib.regions.RegionRequest;
import qupath.opencv.tools.OpenCVTools

import java.awt.image.ComponentColorModel
import java.awt.image.DataBuffer

import static qupath.lib.gui.scripting.QPEx.*;

// Variables to set
//////////////////////////////////
String registrationType="AFFINE" //Specify as "RIGID" or "AFFINE"
String refStain = "PTEN" //stain to use as reference image (all images will be aligned to this)
String wsiExt = ".ndpi" //image name extension
//def align_specific=['N19-1107 30Gy M5']//If auto-align on intensity fails, put the image(s) that it fails on here
def AutoAlignPixelSize = 30 //downsample factor for calculating transform (tform). Does not affect scaling of output image
align_specific=null //When referencing an image, just include the slide name (stuff before _)
skip_image=0 // If 1, skips the images defined by 'align_specific'. If 0, skips all but image(s) in 'align_specific'
//Experimental features
use_single_channel=0 // Use a single channel from each image for alignment (set to channel number to use). Set to 0 to use all channels.
iterations=5 // Number of times to iteratively calculate the transformation
mov_rotation=180 // rotation to apply to ALL moving images before calculating alignment. Strongly recommend ensuring proper orientation before loading into QuPath.
decrement_factor=1.1 // if iterations>1, by what factor to decrease AutoAlignPixelSize (increasing resolution of alignment). Set to 1 to leave AutoAlignPixelSize unchanged across iterations.

/////////////////////////////////


//Lim's code for file name matching
// Get list of all images in project
def projectImageList = getProject().getImageList()

// Create empty lists
def imageNameList = []
def slideIDList = []
def stainList = []
def missingList = []

// Split image file names to desired variables and add to previously created lists
for (entry in projectImageList) {
    def name = entry.getImageName()
    def (imageName, imageExt) = name.split('\\.')
    def (slideID, stain) = imageName.split('_')
    imageNameList << imageName
    slideIDList << slideID
    stainList << stain
}

// Remove duplicate entries from lists
slideIDList = slideIDList.unique()
stainList = stainList.unique()
print (slideIDList)
print (align_specific)
// Remove specific entries if causing alignment to not converge
if (align_specific != null)
    if (skip_image == 1)
        slideIDList.removeAll(align_specific)
    else
        slideIDList.retainAll(align_specific)


if (stainList.size() == 1) {
    print 'Only one stain detected. Target slides may not be loaded.'
    return
}

// Create Affine folder to put transformation matrix files
path = buildFilePath(PROJECT_BASE_DIR, 'Affine')
mkdirs(path)

// Process all combinations of slide IDs, tissue blocks, and stains based on reference stain slide onto target slides
for (slide in slideIDList) {
    for (stain in stainList) {
        if (stain != refStain) {
            refFileName = slide + "_" + refStain + wsiExt
            targetFileName = slide + "_" + stain + wsiExt
            path = buildFilePath(PROJECT_BASE_DIR, 'Affine', targetFileName)
            
            def refImage = projectImageList.find {it.getImageName() == refFileName}
            def targetImage = projectImageList.find {it.getImageName() == targetFileName}
            if (refImage == null) {
                print 'Reference slide ' + refFileName + ' missing!'
                missingList << refFileName
                continue
            }
            if (targetImage == null) {
                print 'Target slide ' + targetFileName + ' missing!'
                missingList << targetFileName
                continue
            }
            println("Aligning reference " + refFileName + " to target " + targetFileName)
            //McArdle's code for image alignment
            ImageServer<BufferedImage> serverBase = refImage.readImageData().getServer()
            ImageServer<BufferedImage> serverOverlay = targetImage.readImageData().getServer()
            def static_img_name = refFileName
            def moving_img_name = targetFileName
            def project_name = getProject()
            def entry_name_static = project_name.getImageList().find { it.getImageName() == static_img_name }
            def entry_name_moving = project_name.getImageList().find { it.getImageName() == moving_img_name }

            def serverBaseMark = entry_name_static.readImageData()
            def serverOverlayMark = entry_name_moving.readImageData()
            Affine affine=[]


            mov_width=serverOverlayMark.getServer().getWidth()
            mov_height=serverOverlayMark.getServer().getHeight()
            affine.prependRotation(mov_rotation,mov_width/2,mov_height/2)



            for(int i = 0;i<iterations;i++) {
                //Perform the alignment. If no annotations present, use intensity. If annotations present, use area
                print("autoalignpixelsize:" + AutoAlignPixelSize)
                if (serverBaseMark.hierarchy.nObjects() > 0 || serverOverlayMark.hierarchy.nObjects() > 0)
                    autoAlignPrep(AutoAlignPixelSize, "AREA", serverBaseMark, serverOverlayMark, affine, registrationType, use_single_channel)
                else
                    autoAlignPrep(AutoAlignPixelSize, "notAREA", serverBaseMark, serverOverlayMark, affine, registrationType, use_single_channel)
                AutoAlignPixelSize/=decrement_factor
                if (AutoAlignPixelSize<1){
                    AutoAlignPixelSize=1
                }
            }

            def matrix = []


            matrix << affine.getMxx()
            matrix << affine.getMxy()
            matrix << affine.getTx()
            matrix << affine.getMyx()
            matrix << affine.getMyy()
            matrix << affine.getTy()

            new File(path).withObjectOutputStream {
                it.writeObject(matrix)
            }
        }
    }
}

if (missingList.isEmpty() == true) {
    print 'Done!'
} else {
    missingList = missingList.unique()
    print 'Done! Missing slides: ' + missingList
}


/*Subfunctions taken from here:
https://github.com/qupath/qupath/blob/a1465014c458d510336993802efb08f440b50cc1/qupath-experimental/src/main/java/qupath/lib/gui/align/ImageAlignmentPane.java
 */

//creates an image server using the actual images (for intensity-based alignment) or a labeled image server (for annotation-based).
double autoAlignPrep(double requestedPixelSizeMicrons, String alignmentMethod, ImageData<BufferedImage> imageDataBase, ImageData<BufferedImage> imageDataSelected, Affine affine,String registrationType, int use_single_channel) throws IOException {
    ImageServer<BufferedImage> serverBase, serverSelected;

    if (alignmentMethod == 'AREA') {
        logger.debug("Image alignment using area annotations");
        Map<PathClass, Integer> labels = new LinkedHashMap<>();
        int label = 1;
        labels.put(PathClassFactory.getPathClassUnclassified(), label++);
        for (def annotation : imageDataBase.getHierarchy().getAnnotationObjects()) {
            def pathClass = annotation.getPathClass();
            if (pathClass != null && !labels.containsKey(pathClass))
                labels.put(pathClass, label++);
        }
        for (def annotation : imageDataSelected.getHierarchy().getAnnotationObjects()) {
            def pathClass = annotation.getPathClass();
            if (pathClass != null && !labels.containsKey(pathClass))
                labels.put(pathClass, label++);
        }

        double downsampleBase = requestedPixelSizeMicrons / imageDataBase.getServer().getPixelCalibration().getAveragedPixelSize().doubleValue();
        serverBase = new LabeledImageServer.Builder(imageDataBase)
                .backgroundLabel(0)
                .addLabels(labels)
                .downsample(downsampleBase)
                .build();

        double downsampleSelected = requestedPixelSizeMicrons / imageDataSelected.getServer().getPixelCalibration().getAveragedPixelSize().doubleValue();
        serverSelected = new LabeledImageServer.Builder(imageDataSelected)
                .backgroundLabel(0)
                .addLabels(labels)
                .downsample(downsampleSelected)
                .build();
        //disable single channel alignment when working with Area annotations, unsure what bugs it can cause
        use_single_channel=0
    } else {
        // Default - just use intensities
        logger.debug("Image alignment using intensities");
        serverBase = imageDataBase.getServer();
        serverSelected = imageDataSelected.getServer();
    }

    scaleFactor=autoAlign(serverBase, serverSelected, registrationType, affine, requestedPixelSizeMicrons,use_single_channel);
    return scaleFactor
}

double autoAlign(ImageServer<BufferedImage> serverBase, ImageServer<BufferedImage> serverOverlay, String regionstrationType, Affine affine, double requestedPixelSizeMicrons, use_single_channel) {
    PixelCalibration calBase = serverBase.getPixelCalibration()
    double pixelSizeBase = calBase.getAveragedPixelSizeMicrons()
    double downsampleBase = 1
    if (!Double.isFinite(pixelSizeBase)) {
      //  while (serverBase.getWidth() / downsampleBase > 2000)
       //     downsampleBase++;
       // logger.warn("Pixel size is unavailable! Default downsample value of {} will be used", downsampleBase)
        pixelSizeBase=50
        downsampleBase = requestedPixelSizeMicrons / pixelSizeBase
    } else {
        downsampleBase = requestedPixelSizeMicrons / pixelSizeBase
    }

    PixelCalibration calOverlay = serverOverlay.getPixelCalibration()
    double pixelSizeOverlay = calOverlay.getAveragedPixelSizeMicrons()
    double downsampleOverlay = 1
    if (!Double.isFinite(pixelSizeOverlay)) {
    //    while (serverBase.getWidth() / downsampleOverlay > 2000)
    //        downsampleOverlay++;
     //   logger.warn("Pixel size is unavailable! Default downsample value of {} will be used", downsampleOverlay)
        pixelSizeOverlay=50
        downsampleOverlay = requestedPixelSizeMicrons / pixelSizeOverlay
    } else {
        downsampleOverlay = requestedPixelSizeMicrons / pixelSizeOverlay
    }

    double scaleFactor=downsampleBase/downsampleOverlay

    BufferedImage imgBase = serverBase.readBufferedImage(RegionRequest.createInstance(serverBase.getPath(), downsampleBase, 0, 0, serverBase.getWidth(), serverBase.getHeight()))
    BufferedImage imgOverlay = serverOverlay.readBufferedImage(RegionRequest.createInstance(serverOverlay.getPath(), downsampleOverlay, 0, 0, serverOverlay.getWidth(), serverOverlay.getHeight()))

    //Determine whether to calculate intensity-based alignment using all channels or a single channel
    Mat matBase
    Mat matOverlay
    if (use_single_channel==0) {
        //print 'using all channels'
        imgBase = ensureGrayScale(imgBase)
        imgOverlay = ensureGrayScale(imgOverlay)
        matBase = OpenCVTools.imageToMat(imgBase)
        matOverlay = OpenCVTools.imageToMat(imgOverlay)

    } else {

        matBase = OpenCVTools.imageToMat(imgBase)
        matOverlay = OpenCVTools.imageToMat(imgOverlay)
        int channel = use_single_channel-1
        //print ('using channel ' + channel)
        matBase = OpenCVTools.splitChannels(matBase)[channel]
        matOverlay = OpenCVTools.splitChannels(matOverlay)[channel]
        //use this to preview how the channel looks
        //OpenCVTools.matToImagePlus('Channel:' + channel.toString(), matBase).show()
    }


    /////pete code block/////

//// New bit
//    int channel = 2
//    matBase = OpenCVTools.splitChannels(matBase)[channel]
//    matOverlay = OpenCVTools.splitChannels(matOverlay)[channel]
//  ///end pete code block///

    Mat matTransform = Mat.eye(2, 3, opencv_core.CV_32F).asMat()
// Initialize using existing transform
//		affine.setToTransform(mxx, mxy, tx, myx, myy, ty)
    try {
        FloatIndexer indexer = matTransform.createIndexer()
        indexer.put(0, 0, (float)affine.getMxx())
        indexer.put(0, 1, (float)affine.getMxy())
        indexer.put(0, 2, (float)(affine.getTx() / downsampleBase))
        indexer.put(1, 0, (float)affine.getMyx())
        indexer.put(1, 1, (float)affine.getMyy())
        indexer.put(1, 2, (float)(affine.getTy() / downsampleBase))
//			System.err.println(indexer)
    } catch (Exception e) {
        logger.error("Error closing indexer", e)
    }

    TermCriteria termCrit = new TermCriteria(TermCriteria.COUNT, 100, 0.0001)

    try {
        int motion
        switch (regionstrationType) {
            case "AFFINE":
                motion = opencv_video.MOTION_AFFINE
                break
            case "RIGID":
                motion = opencv_video.MOTION_EUCLIDEAN
                break
            default:
                logger.warn("Unknown registraton type {} - will use {}", regionstrationType, RegistrationType.AFFINE)
                motion = opencv_video.MOTION_AFFINE
                break
        }
        double result = opencv_video.findTransformECC(matBase, matOverlay, matTransform, motion, termCrit, null)
        logger.info("Transformation result: {}", result)
    } catch (Exception e) {
        Dialogs.showErrorNotification("Estimate transform", "Unable to estimated transform - result did not converge")
        logger.error("Unable to estimate transform", e)
        return
    }

// To use the following function, images need to be the same size
//		def matTransform = opencv_video.estimateRigidTransform(matBase, matOverlay, false);
    Indexer indexer = matTransform.createIndexer()
    affine.setToTransform(
            indexer.getDouble(0, 0),
            indexer.getDouble(0, 1),
            indexer.getDouble(0, 2) * downsampleBase,
            indexer.getDouble(1, 0),
            indexer.getDouble(1, 1),
            indexer.getDouble(1, 2) * downsampleBase
    )
    indexer.release()

    matBase.release()
    matOverlay.release()
    matTransform.release()

    return scaleFactor
}

//to gather detection objects instead of annotation, change line ~250 to def pathObjects = otherHierarchy.getDetectionObjects()
def GatherObjects(boolean deleteExisting, boolean createInverse, File f){
    f.withObjectInputStream {
        matrix = it.readObject()

        // Get the project & the requested image name
        def project = getProject()
        def entry = project.getImageList().find {it.getImageName()+".aff" == f.getName()}
        if (entry == null) {
            print 'Could not find image with name ' + f.getName()
            return
        }

        def otherHierarchy = entry.readHierarchy()
        def pathObjects = otherHierarchy.getDetectionObjects() //OR getAnnotationObjects()

        // Define the transformation matrix
        def transform = new AffineTransform(
                matrix[0], matrix[3], matrix[1],
                matrix[4], matrix[2], matrix[5]
        )
        if (createInverse)
            transform = transform.createInverse()

        if (deleteExisting)
            clearAllObjects()

        def newObjects = []
        for (pathObject in pathObjects) {
            newObjects << transformObject(pathObject, transform)
        }
        addObjects(newObjects)
    }
}

//other subfunctions

PathObject transformObject(PathObject pathObject, AffineTransform transform) {
    // Create a new object with the converted ROI
    def roi = pathObject.getROI()
    def roi2 = transformROI(roi, transform)
    def newObject = null
    if (pathObject instanceof PathCellObject) {
        def nucleusROI = pathObject.getNucleusROI()
        if (nucleusROI == null)
            newObject = PathObjects.createCellObject(roi2, pathObject.getPathClass(), pathObject.getMeasurementList())
        else
            newObject = PathObjects.createCellObject(roi2, transformROI(nucleusROI, transform), pathObject.getPathClass(), pathObject.getMeasurementList())
    } else if (pathObject instanceof PathTileObject) {
        newObject = PathObjects.createTileObject(roi2, pathObject.getPathClass(), pathObject.getMeasurementList())
    } else if (pathObject instanceof PathDetectionObject) {
        newObject = PathObjects.createDetectionObject(roi2, pathObject.getPathClass(), pathObject.getMeasurementList())
        newObject.setName(pathObject.getName())
    } else {
        newObject = PathObjects.createAnnotationObject(roi2, pathObject.getPathClass(), pathObject.getMeasurementList())
        newObject.setName(pathObject.getName())
    }
    // Handle child objects
    if (pathObject.hasChildren()) {
        newObject.addPathObjects(pathObject.getChildObjects().collect({transformObject(it, transform)}))
    }
    return newObject
}

ROI transformROI(ROI roi, AffineTransform transform) {
    def shape = RoiTools.getShape(roi) // Should be able to use roi.getShape() - but there's currently a bug in it for rectangles/ellipses!
    shape2 = transform.createTransformedShape(shape)
    return RoiTools.getShapeROI(shape2, roi.getImagePlane(), 0.5)
}

static BufferedImage ensureGrayScale(BufferedImage img) {
    if (img.getType() == BufferedImage.TYPE_BYTE_GRAY)
        return img
    if (img.getType() == BufferedImage.TYPE_BYTE_INDEXED) {
        ColorSpace cs = ColorSpace.getInstance(ColorSpace.CS_GRAY)
        def colorModel = new ComponentColorModel(cs, 8 as int[], false, true,
                Transparency.OPAQUE,
                DataBuffer.TYPE_BYTE)
        return new BufferedImage(colorModel, img.getRaster(), false, null)
    }
    BufferedImage imgGray = new BufferedImage(img.getWidth(), img.getHeight(), BufferedImage.TYPE_BYTE_GRAY)
    Graphics2D g2d = imgGray.createGraphics()
    g2d.drawImage(img, 0, 0, null)
    g2d.dispose()
    return imgGray
}