####BiosignalGroup####Idiap###2022####
####FullAuto_VolumeStraightened_3dS###

# Import libraries
from __main__ import vtk, slicer
import ExtractCenterline
import CurvedPlanarReformat
import time
from PIL import Image
import numpy as np
import os
import nrrd

# Starting the clock and variable of cases, for the log file
start_clock = time.time()
var_num_cases = 0

# Setting the for structure / Verify here the origin of the data.
directory = os.fsencode("C:/data")
for file in os.listdir(directory):
    filename = os.fsdecode(file)

    if filename.endswith(".nrrd"):

        ## Opening data ##
        index = filename.rfind(".")
        file_names = filename[0:index]
        nodepoints = filename[0:index]

        markk = 'C:/data/'+file_names+'.json'
        markks = 'C:/data/'+file_names+'.nrrd'

        slicer.util.loadMarkups(markk)
        slicer.util.loadVolume(markks)

        ## First Process: Definition initial points ##
        # Take coordinates of points and creation of center point base on them
        pun1 = vtk.vtkVector3d(0, 0, 0)
        pun2 = vtk.vtkVector3d(0, 0, 0)
        pointListNode = slicer.util.getNode(nodepoints)
        pointListNode.GetNthControlPointPosition(0, pun1)
        pointListNode.GetNthControlPointPosition(1, pun2)
        punc = vtk.vtkVector3d((pun1[0]+pun2[0])/2, (pun1[1]+pun2[1])/2, (pun1[2]+pun2[2])/2)

        ## Process: Definition ROI ##
        # Create a new ROI
        roiNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLAnnotationROINode")
        roiNode.SetXYZ(punc)
        punc_d = vtk.vtkVector3d((abs(pun1[0]-pun2[0]))/2, (abs(pun1[1]-pun2[1]))/2, (abs(pun1[2]-pun2[2]))/2)
        roiNode.SetRadiusXYZ(punc_d)

        ## Process: Crop ##
        # Make the crop and create new volume
        inputVolume = slicer.util.getNode('vtkMRMLScalarVolumeNode1')
        cropVolumeLogic = slicer.modules.cropvolume.logic()
        cropVolumeParameterNode = slicer.vtkMRMLCropVolumeParametersNode()
        cropVolumeParameterNode.SetROINodeID(roiNode.GetID())
        cropVolumeParameterNode.SetInputVolumeNodeID(inputVolume.GetID())
        cropVolumeParameterNode.SetVoxelBased(True)
        cropVolumeParameterNode.SetSpacingScalingConst(0.0)
        cropVolumeLogic.Apply(cropVolumeParameterNode)
        croppedVolume = slicer.mrmlScene.GetNodeByID(cropVolumeParameterNode.GetOutputVolumeNodeID())

        ## Process: Segmentation ##
        # Create segmentation
        segmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode")
        segmentationNode.CreateDefaultDisplayNodes()
        segmentationNode.SetReferenceImageGeometryParameterFromVolumeNode(croppedVolume)
        addedSegmentID = segmentationNode.GetSegmentation().AddEmptySegment("skin")

        # Create segment editor to get access to effects
        segmentEditorWidget = slicer.qMRMLSegmentEditorWidget()
        segmentEditorWidget.setMRMLScene(slicer.mrmlScene)
        segmentEditorNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentEditorNode")
        segmentEditorWidget.setMRMLSegmentEditorNode(segmentEditorNode)
        segmentEditorWidget.setSegmentationNode(segmentationNode)
        segmentEditorWidget.setMasterVolumeNode(croppedVolume)

        # Thresholding
        segmentEditorWidget.setActiveEffectByName("Threshold")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("MinimumThreshold", "160")
        effect.setParameter("MaximumThreshold", "700")
        effect.self().onApply()

        # Clean up the segmentEditorWidget
        segmentEditorWidget = None
        slicer.mrmlScene.RemoveNode(segmentEditorNode)

        ## Process: Extract center line ##
        # Segment data needed
        segmentName = 'skin'
        segmentID = segmentationNode.GetSegmentation().GetSegmentIdBySegmentName(segmentName)
        extractLogic = ExtractCenterline.ExtractCenterlineLogic()

        # Preprocess the surface
        inputSurfacePolyData = extractLogic.polyDataFromNode(segmentationNode, segmentID)
        targetNumberOfPoints = 5000.0
        decimationAggressiveness = 4
        subdivideInputSurface = False
        preprocessedPolyData = extractLogic.preprocess(inputSurfacePolyData, targetNumberOfPoints, decimationAggressiveness, subdivideInputSurface)

        # Extract the centerline
        centerlineCurveNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsCurveNode", "Centerline curve")
        centerlinePolyData, voronoiDiagramPolyData = extractLogic.extractCenterline(preprocessedPolyData, pointListNode)
        centerlinePropertiesTableNode = None
        extractLogic.createCurveTreeFromCenterline(centerlinePolyData, centerlineCurveNode, centerlinePropertiesTableNode)

        ## Process: CurvedPlanarReformat ##
        # Settings of the process
        fieldOfView = [40.0, 40.0]
        outputSpacing = [0.5, 0.5, 1.0]

        # Calling logic section of the module CurvedPlanarReformat
        logic = CurvedPlanarReformat.CurvedPlanarReformatLogic()

        # Creating node and using the module
        straighteningTransformNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLTransformNode','Straightening transform')
        logic.computeStraighteningTransform(straighteningTransformNode, centerlineCurveNode, fieldOfView, outputSpacing[2])
        straightenedVolume = slicer.modules.volumes.logic().CloneVolume(inputVolume,inputVolume.GetName()+' straightened')
        logic.straightenVolume(straightenedVolume, inputVolume, outputSpacing, straighteningTransformNode)

        # Remove this node because its not necessary to show it
        slicer.mrmlScene.RemoveNode(straighteningTransformNode)

        ## Outputs ##
        # Saving objects
        V_olumecropped = slicer.util.getNode('vtkMRMLScalarVolumeNode2')
        S_egmen = slicer.util.getNode('vtkMRMLSegmentationNode1')
        C_enterLine = slicer.util.getNode('vtkMRMLMarkupsCurveNode1')
        St_raightenedVolume = slicer.util.getNode('vtkMRMLScalarVolumeNode3')

        # Creating the main folder
        index = filename.rfind(".")
        file_name = filename[0:index]
        newpath = r'C:/Results/'+file_name
        if not os.path.exists(newpath):
            os.makedirs(newpath)

        # Saving objets
        slicer.util.saveNode(V_olumecropped, newpath+'/croppvolume.nrrd')
        slicer.util.saveNode(S_egmen, newpath+'/segmentation.seg.nrrd')
        slicer.util.saveNode(C_enterLine, newpath+'/centerline.json')
        slicer.util.saveNode(St_raightenedVolume, newpath+'/straightenedVolume.nrrd')

        # Saving images
        def manipulating_nrrd_contrast(img, level):
            img_c = img.astype(int).copy()
            factor = (32 * (level + 255)) / (255 * (259 - level))  # Contrast set
            img_c = factor * (img_c - 128) + 128
            img_c = np.clip(img_c, 0, 255)
            return img_c.astype(np.uint8)

        readdata, header = nrrd.read(newpath+'/straightenedVolume.nrrd')
        ph_1 = np.asarray(readdata[:, 38, :]).astype(int)
        final = Image.fromarray(manipulating_nrrd_contrast(ph_1, 128))
        final.save(newpath+'/90°.png')
        ph_2 = np.asarray(readdata[38, :, :]).astype(int)
        final = Image.fromarray(manipulating_nrrd_contrast(ph_2, 128))
        final.save(newpath+'/00°.png')

        # Increase variable counter
        var_num_cases += 1

## Log file ##
# Closing clock
end_clock = time.time()
duration = "{:.2f}".format(end_clock - start_clock)

# Creating log file with results
with open('C:/Results/Log_file.txt', 'w') as f:
    f.writelines(['Biosignal Group - Idiap / FullAuto_VolumeStraightened_3dS\n', '--------------------------------------------------------- \n'])
    f.writelines(['Info about the run:\n', '- Time of execution: '+duration+' seconds.\n', '- Num of cases: '+str(var_num_cases)+' \n', '--- \n'])
    f.writelines(['Settings Crop:\n', '- Spacing Scaling Const: (0.0)\n', '--- \n'])
    f.writelines(['Settings Segmentation:\n', '- Active Effect: Threshold\n', '- MinimumThreshold: 160\n', "- MaximumThreshold: 700\n", '--- \n'])
    f.writelines(['Settings CenterLine:\n', '- Number Points = 5000.0\n', '- Decimation Aggressiveness = 4\n', '--- \n'])
    f.writelines(['Settings CurvedPlanarReformat:\n', '- Field of View: '+str(fieldOfView)+'\n', '- Out put Spacing: '+str(outputSpacing)+'\n'])
    f.write('-------------------End Log-------------------------------- \n')