guiscript=true

import javafx.collections.FXCollections
import javafx.geometry.Insets
import javafx.geometry.Pos
import javafx.scene.Scene
import javafx.scene.control.Button
import javafx.scene.control.ChoiceBox
import javafx.scene.layout.GridPane
import javafx.scene.layout.StackPane
import javafx.scene.text.Font
import javafx.scene.text.FontWeight
import javafx.scene.text.Text
import javafx.stage.Stage
import qupath.lib.gui.prefs.PathPrefs
import qupath.lib.gui.QuPathGUI.Modes
import qupath.lib.gui.QuPathGUI
import qupath.lib.images.ImageData
import qupath.lib.images.servers.ImageServer
import qupath.lib.images.servers.ImageServerProvider
import qupath.lib.io.PathIO
import qupath.lib.objects.PathObject
import qupath.lib.scripting.QP
import qupath.lib.common.GeneralTools

import groovy.io.FileType
import qupath.lib.scripting.QPEx

import javax.script.ScriptContext
import javax.script.ScriptEngine
import javax.script.ScriptEngineManager
import javax.script.SimpleBindings
import javax.script.SimpleScriptContext
import java.awt.image.BufferedImage

/** Pre-analysis script for selecting annotations to run a script on
 *
 *
 * @author Caleb Grenko
 *
 **/

Preanalysis.run()

class Preanalysis {

    static void run() {
        def selections = getSelections()
        print("When this appears, you should have just hit 'Done'")
        def script = getScript()
        print(script)
        runScriptOnSelections(script, selections)
    }
    /**
     * Checks for selections. If no annotations are selected, gives option of either launching selection viewer
     * or using the existing annotations. However, it disables the annotation button if there are no annotations.
     * After the choice is made, it selects the objects which you chose.
     *
     */
    static void getSelections() {
        print("Getting selections")

        def selections = QP.getSelectedObjects().findAll() { it.getDisplayedName() == "PathAnnotationObject" }.toSet()
        def hasSelections = (!selections.isEmpty())

        GridPane grid = new GridPane()
        grid.setAlignment(Pos.CENTER)
        grid.setHgap(50)
        grid.setVgap(15)
        grid.setPadding(new Insets(15, 10, 15, 10))

        //Initial prompt
        Text text = new Text("Select which \nobjects to use")
        text.setFont(Font.font("Arial", FontWeight.LIGHT, 14))
        StackPane stack = new StackPane()
        stack.getChildren().add(text)
        grid.add(stack, 0, 1, 2, 1)

        //Open region selector option
        Button openSelectorButton = new Button("Launch region selector")
        openSelectorButton.setOnAction {
            runRegionSelector()
        }
        openSelectorButton.setPrefWidth(150)
        grid.add(openSelectorButton, 2, 1)

        //Use currently selected objects
        Button useSelectionsButton = new Button("Use current selections")

        if (!hasSelections) {
            useSelectionsButton.setDisable(true)
        }
        useSelectionsButton.setOnAction {
            useSelectionsButton.getScene().getWindow().hide()
            QP.selectAnnotations()
        }
        useSelectionsButton.setPrefWidth(150)
        grid.add(useSelectionsButton, 2, 2)

        //Use all annotations
        Button useAnnotationsButton = new Button("Use all annotations")

        if (QP.getAnnotationObjects().isEmpty()) {
            useAnnotationsButton.setDisable(true)
        }
        useAnnotationsButton.setOnAction {
            useAnnotationsButton.getScene().getWindow().hide()
            QP.selectAnnotations()
        }
        useAnnotationsButton.setPrefWidth(150)
        grid.add(useAnnotationsButton, 2, 3)

        //Cancel the process
        Button cancelButton = new Button("Close")
        cancelButton.setCancelButton(true)
        cancelButton.setOnAction {
            cancelButton.getScene().getWindow().hide()
        }
        cancelButton.setPrefWidth(150)
        grid.add(cancelButton, 2, 4)

        def stage = new Stage()
        stage.setTitle('Select Objects To Use')
        stage.setResizable(false)
        stage.setScene(new Scene(grid, 300, 200))
        stage.initOwner(QuPathGUI.getInstance().getStage())
        stage.showAndWait()
    }

    /**
     * Allows user to repeatedly draw regions. When the "Done" button is pressed, it selects the regions just drawn.
     *
     * @return List of regions that the user selects
     *
     * This works by toggling the "returnToMoveMode" preference and then selecting the Rectangle mode. When the "Done"
     * button is pressed, it toggles back to the previous setting and changes back to the move marker. When the window is
     * opened, a list is created of all pre-existing annotations. After the user is done, another list is created and the
     * contents of the old list are removed from the new list. The remainder of the objects in the new list are selected.
     *
     * TODO: On drag release, add to empty list for more efficiency
     *
     */
    static LinkedList<PathObject> runRegionSelector() {
        def oldObjects = new LinkedList<PathObject>()
        def qupath = QuPathGUI.getInstance()
        def previousPref = PathPrefs.getReturnToMoveMode()
        def hierarchy = qupath.getImageData().getHierarchy()

        //Region selection scene
        GridPane selectionGrid = new GridPane()
        selectionGrid.setAlignment(Pos.CENTER)
        selectionGrid.setHgap(30)
        selectionGrid.setVgap(10)
        selectionGrid.setPadding(new Insets(25, 10, 25, 10))
        Text selectText = new Text("Hit 'Done!' when finished creating regions")
        selectText.setFont(Font.font("Arial", FontWeight.LIGHT, 14))
        StackPane selectionStack = new StackPane()
        selectionStack.getChildren().add(selectText)
        selectionGrid.add(selectionStack, 0, 1, 2, 1)

        //Done Button: resets user preferences for ReturnToMoveMode
        Button doneButton = new Button()
        doneButton.setText("Done!")

        doneButton.setOnAction() {
            PathPrefs.setReturnToMoveMode(previousPref)
            qupath.setMode(Modes.MOVE)
            def newObjectsList = qupath.getImageData().getHierarchy().getFlattenedObjectList(new LinkedList<PathObject>())
            newObjectsList.removeAll(oldObjects)
            doneButton.getScene().getWindow().hide()
            hierarchy.getSelectionModel().selectObjects(newObjectsList)
        }
        selectionGrid.add(doneButton, 3, 3)

        def stage = new Stage()
        stage.setTitle('Select regions')

        stage.setOnShown {

            oldObjects = hierarchy.getFlattenedObjectList(new LinkedList<PathObject>())
            PathPrefs.setReturnToMoveMode(false)
            qupath.setMode((Modes.RECTANGLE))
        }

        stage.setOnCloseRequest {
            PathPrefs.setReturnToMoveMode(previousPref)
            qupath.setMode(Modes.MOVE)
        }

        stage.setScene(new Scene(selectionGrid, 300, 100))
        stage.setResizable(false)
        stage.setAlwaysOnTop(true)

        stage.initOwner(qupath.getStage())
        stage.showAndWait()
        return QP.getSelectedObjects()
    }

    static String getScript() {
        //Fetches image data, sets to Brightfield (other) and estimates stains if not set already
        def imageData = QP.getCurrentImageData()
        print("Current image type: " + imageData.getImageType())
        //Gets directory of projects (assumes parent folder for current project)
        def generalProjectDirectory = getProjectBaseDirectory().substring(0, getProjectBaseDirectory().lastIndexOf("\\"))

        GridPane grid = new GridPane()
        grid.setAlignment(Pos.CENTER)
        grid.setHgap(0)
        grid.setVgap(15)
        grid.setPadding(new Insets(15, 10, 15, 10))

        Text text = new Text("Select script to run on selected regions")
        text.setFont(Font.font("Arial", FontWeight.LIGHT, 14))
        StackPane stack = new StackPane()
        stack.getChildren().add(text)
        grid.add(stack, 0, 0,)




        // Create a File object
        def folder = new File("${generalProjectDirectory}\\Analysis\\Scripts")

        // If it doesn't exist
        if( !folder.exists() ) {
            // Create all folders up-to and including B
            folder.mkdirs()
        }

        def list = []

        folder.eachFileRecurse(FileType.FILES) { file ->
            list << file
        }

        def listNames = new ArrayList<String>()

        def name
        for (item in list) {
            name = item.toString()
            listNames.add(name.substring(name.lastIndexOf('\\') + 1, name.lastIndexOf('.')))
        }

        ChoiceBox scripts = new ChoiceBox(FXCollections.observableArrayList(listNames))
        grid.add(scripts, 0, 1)

        Button ok = new Button("Ok")
        ok.setOnAction {
            ok.getScene().getWindow().hide()
        }

        grid.add(ok, 1, 1)
        def stage = new Stage()
        stage.setTitle('Select script')
        stage.setResizable(false)
        stage.setScene(new Scene(grid, 300, 100))
        stage.initOwner(QuPathGUI.getInstance().getStage())
        stage.showAndWait()
        return list.get(scripts.getSelectionModel().getSelectedIndex())
    }

    /** Executes a given script on a specified collection of regions
     *
     * @param scriptPath the analysis script to be used on selected regions
     * @param selections a collection of PathAnnotationObject where the script will be executed
     *
     * TODO: Get it to actually work
     *
     */
    static void runScriptOnSelections(String scriptPath, Collection<PathObject> selections) {
        ScriptEngineManager manager = new ScriptEngineManager()
        ScriptEngine engine = manager.getEngineByExtension("groovy")
        String script = GeneralTools.readFileAsString(scriptPath)

        engine.eval(script)

    }
}