AvtomatskaPoravnavaMarkerjev/SeekTransformModule/SeekTransformModule.py

import os
import numpy as np
import scipy
from scipy.spatial.distance import cdist
from scipy.spatial.transform import Rotation as R
import slicer
import itertools
from DICOMLib import DICOMUtils
from collections import deque
import vtk
from slicer.ScriptedLoadableModule import *
import qt
import matplotlib.pyplot as plt
import csv
import time

#exec(open("C:/Users/lkomar/Documents/Prostata/FirstTryRegister.py").read())

class SeekTransformModule(ScriptedLoadableModule):
    """
    Module description shown in the module panel.
    """
    def __init__(self, parent):
        ScriptedLoadableModule.__init__(self, parent)
        self.parent.title = "Seek Transform module"
        self.parent.categories = ["Image Processing"]
        self.parent.contributors = ["Luka Komar (Onkološki Inštitut Ljubljana, Fakulteta za Matematiko in Fiziko Ljubljana)"]
        self.parent.helpText = "This module applies rigid transformations to CBCT volumes based on reference CT volumes."
        self.parent.acknowledgementText = "Supported by doc. Primož Peterlin & prof. Andrej Studen"

class SeekTransformModuleWidget(ScriptedLoadableModuleWidget):
    """
    GUI of the module.
    """
    def setup(self):
        ScriptedLoadableModuleWidget.setup(self)

        # Dropdown menu za izbiro metode
        self.rotationMethodComboBox = qt.QComboBox()
        self.rotationMethodComboBox.addItems(["Kabsch", "Horn", "Iterative Closest Point (Kabsch)"])
        self.layout.addWidget(self.rotationMethodComboBox)

        # Checkboxi za transformacije
        self.rotationCheckBox = qt.QCheckBox("Rotation")
        self.rotationCheckBox.setChecked(True)
        self.layout.addWidget(self.rotationCheckBox)

        self.translationCheckBox = qt.QCheckBox("Translation")
        self.translationCheckBox.setChecked(True)
        self.layout.addWidget(self.translationCheckBox)

        self.scalingCheckBox = qt.QCheckBox("Scaling")
        self.scalingCheckBox.setChecked(True)
        self.layout.addWidget(self.scalingCheckBox)
        
        self.writefileCheckBox = qt.QCheckBox("Write distances to csv file")
        self.writefileCheckBox.setChecked(True)
        self.layout.addWidget(self.writefileCheckBox)

        # Load button
        self.applyButton = qt.QPushButton("Find markers and transform")
        self.applyButton.toolTip = "Finds markers, computes optimal rigid transform and applies it to CBCT volumes."
        self.applyButton.enabled = True
        self.layout.addWidget(self.applyButton)

        # Connect button to logic
        self.applyButton.connect('clicked(bool)', self.onApplyButton)

        self.layout.addStretch(1)

    def onApplyButton(self):
        logic = MyTransformModuleLogic()
        selectedMethod = self.rotationMethodComboBox.currentText  # izberi metodo izračuna rotacije

        # Preberi stanje checkboxov
        applyRotation = self.rotationCheckBox.isChecked()
        applyTranslation = self.translationCheckBox.isChecked()
        applyScaling = self.scalingCheckBox.isChecked()
        writefilecheck = self.writefileCheckBox.isChecked()

        # Pokliči logiko z izbranimi nastavitvami
        logic.run(selectedMethod, applyRotation, applyTranslation, applyScaling, writefilecheck)


class MyTransformModuleLogic(ScriptedLoadableModuleLogic):
    """
    Core logic of the module.
    """
    
    
    def run(self, selectedMethod, applyRotation, applyTranslation, applyScaling, writefilecheck):
        start_time = time.time()
        print("Calculating...")
        
        def group_points(points, threshold):
            # Function to group points that are close to each other
            grouped_points = []
            while points:
                point = points.pop()  # Take one point from the list
                group = [point]  # Start a new group
                
                # Find all points close to this one
                distances = cdist([point], points)  # Calculate distances from this point to others
                close_points = [i for i, dist in enumerate(distances[0]) if dist < threshold]
                
                # Add the close points to the group
                group.extend([points[i] for i in close_points])
                
                # Remove the grouped points from the list
                points = [point for i, point in enumerate(points) if i not in close_points]
                
                # Add the group to the result
                grouped_points.append(group)
            
            return grouped_points

        def region_growing(image_data, seed, intensity_threshold, max_distance):
            dimensions = image_data.GetDimensions()
            visited = set()
            region = []
            queue = deque([seed])

            while queue:
                x, y, z = queue.popleft()
                if (x, y, z) in visited:
                    continue

                visited.add((x, y, z))
                voxel_value = image_data.GetScalarComponentAsDouble(x, y, z, 0)
                
                if voxel_value >= intensity_threshold:
                    region.append((x, y, z))
                    # Add neighbors within bounds
                    for dx, dy, dz in [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)]:
                        nx, ny, nz = x + dx, y + dy, z + dz
                        if 0 <= nx < dimensions[0] and 0 <= ny < dimensions[1] and 0 <= nz < dimensions[2]:
                            if (nx, ny, nz) not in visited:
                                queue.append((nx, ny, nz))

            return region

        
        def compute_optimal_scaling_per_axis(moving_points, fixed_points):
            """Computes optimal scaling factors for each axis (X, Y, Z) to align moving points (CBCT) to fixed points (CT).

            Args:
                moving_points (list of lists): List of (x, y, z) moving points (CBCT).
                fixed_points (list of lists): List of (x, y, z) fixed points (CT).

            Returns:
                tuple: Scaling factors (sx, sy, sz).
            """
            moving_points_np = np.array(moving_points)
            fixed_points_np = np.array(fixed_points)

            # Compute centroids
            centroid_moving = np.mean(moving_points_np, axis=0)
            centroid_fixed = np.mean(fixed_points_np, axis=0)

            # Compute absolute distances of each point from its centroid along each axis
            distances_moving = np.abs(moving_points_np - centroid_moving)
            distances_fixed = np.abs(fixed_points_np - centroid_fixed)

            # Compute scaling factors as the ratio of mean absolute distances per axis
            scale_factors = np.mean(distances_fixed, axis=0) / np.mean(distances_moving, axis=0)

            return tuple(scale_factors)

        def compute_scaling(cbct_points, scaling_factors):
            """Applies non-uniform scaling to CBCT points.

            Args:
                cbct_points (list of lists): List of (x, y, z) points.
                scaling_factors (tuple): Scaling factors (sx, sy, sz) for each axis.

            Returns:
                np.ndarray: Scaled CBCT points.
            """
            sx, sy, sz = scaling_factors  # Extract scaling factors
            scaling_matrix = np.diag([sx, sy, sz])  # Create diagonal scaling matrix

            cbct_points_np = np.array(cbct_points)  # Convert to numpy array
            scaled_points = cbct_points_np @ scaling_matrix.T  # Apply scaling

            return scaled_points.tolist()  # Convert back to list

        def compute_Kabsch_rotation(moving_points, fixed_points):
            """
            Computes the optimal rotation matrix to align moving_points to fixed_points.
            
            Parameters:
            moving_points (list or ndarray): List of points to be rotated CBCT
            fixed_points (list or ndarray): List of reference points CT

            Returns:
            ndarray: Optimal rotation matrix.
            """
            assert len(moving_points) == len(fixed_points), "Point lists must be the same length."

            # Convert to numpy arrays
            moving = np.array(moving_points)
            fixed = np.array(fixed_points)

            # Compute centroids
            centroid_moving = np.mean(moving, axis=0)
            centroid_fixed = np.mean(fixed, axis=0)

            # Center the points
            moving_centered = moving - centroid_moving
            fixed_centered = fixed - centroid_fixed

            # Compute covariance matrix
            H = np.dot(moving_centered.T, fixed_centered)

            # SVD decomposition
            U, _, Vt = np.linalg.svd(H)
            Rotate_optimal = np.dot(Vt.T, U.T)

            # Correct improper rotation (reflection)
            if np.linalg.det(Rotate_optimal) < 0:
                Vt[-1, :] *= -1
                Rotate_optimal = np.dot(Vt.T, U.T)

            return Rotate_optimal


        def compute_Horn_rotation(moving_points, fixed_points):
            """
            Computes the optimal rotation matrix using quaternions.

            Parameters:
            moving_points (list or ndarray): List of points to be rotated.
            fixed_points (list or ndarray): List of reference points.

            Returns:
            ndarray: Optimal rotation matrix.
            """
            assert len(moving_points) == len(fixed_points), "Point lists must be the same length."
            
            moving = np.array(moving_points)
            fixed = np.array(fixed_points)
            
            # Compute centroids
            centroid_moving = np.mean(moving, axis=0)
            centroid_fixed = np.mean(fixed, axis=0)
            
            # Center the points
            moving_centered = moving - centroid_moving
            fixed_centered = fixed - centroid_fixed
            
            # Construct the cross-dispersion matrix
            M = np.dot(moving_centered.T, fixed_centered)
            
            # Construct the N matrix for quaternion solution
            A = M - M.T
            delta = np.array([A[1, 2], A[2, 0], A[0, 1]])
            trace = np.trace(M)
            
            N = np.zeros((4, 4))
            N[0, 0] = trace
            N[1:, 0] = delta
            N[0, 1:] = delta
            N[1:, 1:] = M + M.T - np.eye(3) * trace
            
            # Compute the eigenvector corresponding to the maximum eigenvalue
            eigvals, eigvecs = np.linalg.eigh(N)
            q_optimal = eigvecs[:, np.argmax(eigvals)]  # Optimal quaternion
            
            # Convert quaternion to rotation matrix
            w, x, y, z = q_optimal
            R = np.array([
                [1 - 2*(y**2 + z**2), 2*(x*y - z*w), 2*(x*z + y*w)],
                [2*(x*y + z*w), 1 - 2*(x**2 + z**2), 2*(y*z - x*w)],
                [2*(x*z - y*w), 2*(y*z + x*w), 1 - 2*(x**2 + y**2)]
            ])
            
            return R

        def icp_algorithm(moving_points, fixed_points, max_iterations=100, tolerance=1e-5):
            """
            Iterative Closest Point (ICP) algorithm to align moving_points to fixed_points.
            
            Parameters:
            moving_points (list or ndarray): List of points to be aligned.
            fixed_points (list or ndarray): List of reference points.
            max_iterations (int): Maximum number of iterations.
            tolerance (float): Convergence tolerance.

            Returns:
            ndarray: Transformed moving points.
            ndarray: Optimal rotation matrix.
            ndarray: Optimal translation vector.
            """
            # Convert to numpy arrays
            moving = np.array(moving_points)
            fixed = np.array(fixed_points)

            # Initialize transformation
            R = np.eye(3)  # Identity matrix for rotation
            t = np.zeros(3)  # Zero vector for translation

            prev_error = np.inf  # Initialize previous error to a large value

            for iteration in range(max_iterations):
                # Step 1: Find the nearest neighbors (correspondences)
                distances = np.linalg.norm(moving[:, np.newaxis] - fixed, axis=2)
                nearest_indices = np.argmin(distances, axis=1)
                nearest_points = fixed[nearest_indices]

                # Step 2: Compute the optimal rotation and translation
                R_new = compute_Kabsch_rotation(moving, nearest_points)
                centroid_moving = np.mean(moving, axis=0)
                centroid_fixed = np.mean(nearest_points, axis=0)
                t_new = centroid_fixed - np.dot(R_new, centroid_moving)

                # Step 3: Apply the transformation
                moving = np.dot(moving, R_new.T) + t_new

                # Update the cumulative transformation
                R = np.dot(R_new, R)
                t = np.dot(R_new, t) + t_new

                # Step 4: Check for convergence
                mean_error = np.mean(np.linalg.norm(moving - nearest_points, axis=1))
                if np.abs(prev_error - mean_error) < tolerance:
                    print(f"ICP converged after {iteration + 1} iterations.")
                    break
                prev_error = mean_error

            else:
                print(f"ICP reached maximum iterations ({max_iterations}).")

            return moving, R, t

        def match_points(cbct_points, ct_points, auto_weights=True, fallback_if_worse=True, normalize_lengths=True):
            def side_lengths(points):
                lengths = [
                    np.linalg.norm(points[0] - points[1]),
                    np.linalg.norm(points[1] - points[2]),
                    np.linalg.norm(points[2] - points[0])
                ]
                if normalize_lengths:
                    total = sum(lengths) + 1e-6  # da ne delimo z 0
                    lengths = [l / total for l in lengths]
                return lengths

            def triangle_angles(points):
                a = np.linalg.norm(points[1] - points[2])
                b = np.linalg.norm(points[0] - points[2])
                c = np.linalg.norm(points[0] - points[1])
                angle_A = np.arccos(np.clip((b**2 + c**2 - a**2) / (2 * b * c), -1.0, 1.0))
                angle_B = np.arccos(np.clip((a**2 + c**2 - b**2) / (2 * a * c), -1.0, 1.0))
                angle_C = np.pi - angle_A - angle_B
                return [angle_A, angle_B, angle_C]

            def permutation_score(perm, ct_lengths, ct_angles, w_len, w_ang):
                perm_lengths = side_lengths(perm)
                perm_angles = triangle_angles(perm)

                # normaliziraj
                def normalize(vec):
                    norm = np.linalg.norm(vec)
                    return vec / norm if norm > 0 else vec

                perm_lengths = normalize(perm_lengths)
                perm_angles = normalize(perm_angles)
                ct_lengths_n = normalize(ct_lengths)
                ct_angles_n = normalize(ct_angles)

                score_len = sum(abs(a - b) for a, b in zip(perm_lengths, ct_lengths_n))
                score_ang = sum(abs(a - b) for a, b in zip(perm_angles, ct_angles_n))
                return w_len * score_len + w_ang * score_ang

            cbct_points = list(cbct_points)
            ct_lengths = side_lengths(np.array(ct_points))
            ct_angles = triangle_angles(np.array(ct_points))

            if auto_weights:
                var_len = np.var(ct_lengths)
                var_ang = np.var(ct_angles)
                total_var = var_len + var_ang + 1e-6
                weight_length = (1 - var_len / total_var)
                weight_angle = (1 - var_ang / total_var)
            else:
                weight_length = 0.5
                weight_angle = 0.5

            best_perm = None
            best_score = float('inf')

            for perm in itertools.permutations(cbct_points):
                perm = np.array(perm)
                score = permutation_score(perm, ct_lengths, ct_angles, weight_length, weight_angle)
                if score < best_score:
                    best_score = score
                    best_perm = perm

            if fallback_if_worse:
                original_score = permutation_score(np.array(cbct_points), ct_lengths, ct_angles, weight_length, weight_angle)
                if original_score <= best_score or np.allclose(cbct_points, best_perm):
                    print("Fallback to original points due to worse score of the permutation.")
                    return list(cbct_points)

            return list(best_perm)

        def compute_translation(moving_points, fixed_points, rotation_matrix):
            """
            Computes the translation vector to align moving_points to fixed_points given a rotation matrix.
            
            Parameters:
            moving_points (list or ndarray): List of points to be translated.
            fixed_points (list or ndarray): List of reference points.
            rotation_matrix (ndarray): Rotation matrix.

            Returns:
            ndarray: Translation vector.
            """
            # Convert to numpy arrays
            moving = np.array(moving_points)
            fixed = np.array(fixed_points)

            # Compute centroids
            centroid_moving = np.mean(moving, axis=0)
            centroid_fixed = np.mean(fixed, axis=0)

            # Compute translation
            translation = centroid_fixed - np.dot(centroid_moving, rotation_matrix)

            return translation

        def create_vtk_transform(rotation_matrix, translation_vector): 
            """
            Creates a vtkTransform from a rotation matrix and a translation vector.
            """
            # Create a 4x4 transformation matrix
            transform_matrix = np.eye(4)  # Start with an identity matrix
            transform_matrix[:3, :3] = rotation_matrix  # Set rotation part
            transform_matrix[:3, 3] = translation_vector  # Set translation part

            # Convert to vtkMatrix4x4
            vtk_matrix = vtk.vtkMatrix4x4()
            for i in range(4):
                for j in range(4):
                    vtk_matrix.SetElement(i, j, transform_matrix[i, j])
            print("Transform matrix:")
            for i in range(4):
               print(" ".join(f"{vtk_matrix.GetElement(i, j):.6f}" for j in range(4)))
            # Create vtkTransform and set the matrix
            transform = vtk.vtkTransform()
            transform.SetMatrix(vtk_matrix)
            return transform
        
        
        def detect_points_region_growing(volume_name, yesCbct, create_marker, intensity_threshold=3000, x_min=90, x_max=380, y_min=190, y_max=380, z_min=80, z_max=140, max_distance=9, centroid_merge_threshold=5):
            volume_node = slicer.util.getNode(volume_name)
            if not volume_node:
                raise RuntimeError(f"Volume {volume_name} not found.")
            
            image_data = volume_node.GetImageData()
            matrix = vtk.vtkMatrix4x4()
            volume_node.GetIJKToRASMatrix(matrix)

            dimensions = image_data.GetDimensions()
            #detected_regions = []

            if yesCbct: #je cbct ali ct?
                valid_x_min, valid_x_max = 0, dimensions[0] - 1
                valid_y_min, valid_y_max = 0, dimensions[1] - 1
                valid_z_min, valid_z_max = 0, dimensions[2] - 1
            else:
                valid_x_min, valid_x_max = max(x_min, 0), min(x_max, dimensions[0] - 1)
                valid_y_min, valid_y_max = max(y_min, 0), min(y_max, dimensions[1] - 1)
                valid_z_min, valid_z_max = max(z_min, 0), min(z_max, dimensions[2] - 1)

            visited = set()

            def grow_region(x, y, z):
                if (x, y, z) in visited:
                    return None

                voxel_value = image_data.GetScalarComponentAsDouble(x, y, z, 0)
                if voxel_value < intensity_threshold:
                    return None

                region = region_growing(image_data, (x, y, z), intensity_threshold, max_distance=max_distance)
                if region:
                    for point in region:
                        visited.add(tuple(point))
                    return region
                return None

            regions = []
            for z in range(valid_z_min, valid_z_max + 1):
                for y in range(valid_y_min, valid_y_max + 1):
                    for x in range(valid_x_min, valid_x_max + 1):
                        region = grow_region(x, y, z)
                        if region:
                            regions.append(region)

            # Collect centroids using intensity-weighted average
            centroids = []
            for region in regions:
                points = np.array([matrix.MultiplyPoint([*point, 1])[:3] for point in region])
                intensities = np.array([image_data.GetScalarComponentAsDouble(*point, 0) for point in region])
                
                if intensities.sum() > 0:
                    weighted_centroid = np.average(points, axis=0, weights=intensities)
                    max_intensity = intensities.max()
                    centroids.append((np.round(weighted_centroid, 2), max_intensity))

            unique_centroids = []
            for centroid, intensity in centroids:
                if not any(np.linalg.norm(centroid - existing_centroid) < centroid_merge_threshold for existing_centroid, _ in unique_centroids):
                    unique_centroids.append((centroid, intensity))
            
            if create_marker:        
                markups_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"Markers_{volume_name}")
                for centroid, intensity in unique_centroids:
                    markups_node.AddControlPoint(*centroid)
                    markups_node.SetDisplayVisibility(False)
                    #print(f"Detected Centroid (RAS): {centroid}, Max Intensity: {intensity}")

            return unique_centroids

        def find_table_top_z(ct_volume_name, writefilecheck, yesCbct):
            """
            Najde višino zgornjega roba mize v CT/CBCT volumnu in doda markerje.
            
            :param ct_volume_name: Ime volumna v slicerju
            :param writefilecheck: Če je True, zapiše rezultat v CSV
            :param yesCbct: Če je True, uporabi CBCT thresholde
            :return: Višina zgornjega roba mize v mm
            """
            # Pridobi volumen
            ct_volume_node = slicer.util.getNode(ct_volume_name)
            image_data = ct_volume_node.GetImageData()
            spacing = ct_volume_node.GetSpacing()
            dims = image_data.GetDimensions()
            #origin = ct_volume_node.GetOrigin()

            # Pretvori volumen v numpy array
            np_array = slicer.util.arrayFromVolume(ct_volume_node)

            # Izračunaj sredinske IJK koordinate
            mid_ijk = [dims[0] // 2, dims[1] // 2, dims[2] // 2]

            # Preveri, da so IJK indeksi v mejah volumna
            mid_ijk = [max(0, min(dims[i] - 1, mid_ijk[i])) for i in range(3)]

            # Pretvorba IJK → RAS
            ijkToRasMatrix = vtk.vtkMatrix4x4()
            ct_volume_node.GetIJKToRASMatrix(ijkToRasMatrix)

            

            # Sredinski Z slice
            mid_z_voxel = mid_ijk[2]
            slice_data = np_array[mid_z_voxel, :, :]  # (Y, X)

            # Sredinski stolpec
            mid_x_voxel = mid_ijk[0] - 15  # 15 pikslov levo od sredine da ne merimo pri hrbtenici
            column_values = slice_data[:, mid_x_voxel]  # Y smer

            # Doda marker v RAS koordinatah
            #mid_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"Sredina_{ct_volume_name}")
            #mid_ras = np.array(ijkToRasMatrix.MultiplyPoint([*mid_ijk, 1]))[:3]
            #mid_node.AddControlPoint(mid_ras)

            # Določi threshold glede na CBCT ali CT
            threshold = -300 if yesCbct else -100 # če je cbct iščemo vrednost -300, sicer pri CT 0

            # Poišči rob mize (drugi dvig)
            previous_value = -1000
            edge_count = 0
            table_top_y = None
            min_jump = 100 if yesCbct else 50  # Minimalni skok za CBCT in CT

            for y in range(len(column_values) - 1, -1, -1):  # Od spodaj navzgor
                intensity = column_values[y]
                #if column_values[y] > threshold and previous_value <= threshold:
                if (intensity - previous_value) > min_jump and intensity > threshold:  # Namesto primerjave s threshold
                    if yesCbct:
                        table_top_y = y + 1 #Da nismo že v trupu
                        #print(f"Zgornji rob mize najden pri Y = {table_top_y}") # CBCT
                        break
                    if edge_count == 0 or (edge_count == 1 and previous_value < -200):  # Check if intensity is back lower than -400
                        edge_count += 1
                        #print(f"Zaznan rob mize pri X, Y, Z = {mid_x_voxel},{y}, {mid_z_voxel}")
                        if edge_count == 2:  # Drugi dvig = zgornji rob mize
                            table_top_y = y + 1
                            #print(f"Zgornji rob mize najden pri Y = {table_top_y}")
                            break
                previous_value = column_values[y]

            if table_top_y is None:
                print("❌ Zgornji rob mize ni bil najden!")
                return None

            # Pretvorba Y IJK → RAS
            table_ijk = [mid_x_voxel, table_top_y, mid_z_voxel]
            table_ras = np.array(ijkToRasMatrix.MultiplyPoint([*table_ijk, 1]))[:3]

            # Doda marker za višino mize
            table_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"VišinaMize_{ct_volume_name}")
            table_node.AddControlPoint(table_ras)
            table_node.SetDisplayVisibility(False)

            # Izračun višine v mm
            #image_center_y = dims[1] // 2
            pixel_offset = table_top_y
            #mm_offset = pixel_offset * spacing[1]

            #print(f"📏 Miza je {abs(mm_offset):.2f} mm {'nižja' if mm_offset > 0 else 'višja'} od središča.")
            #print(f"📏 Miza je {abs(pixel_offset)} pixlov {'nižja' if pixel_offset > 0 else 'višja'} od središča.")
            #print(f"📏 CBCT table_top_y: {table_top_y}, image_center_y: {image_center_y}, pixel_offset: {pixel_offset}, mm_offset: {mm_offset}")
            #print(f"🛏 Absolutna višina mize (RAS Y): {table_ras[1]} mm")
            # Shrani v CSV
            if writefilecheck:
                file_path = os.path.join(os.path.dirname(__file__), "heightdata.csv")
                with open(file_path, mode='a', newline='') as file:
                    writer = csv.writer(file)
                    modality = "CBCT " if yesCbct else "CT "
                    writer.writerow([modality, ct_volume_name, f" Upper part of table detected at Z = {table_ras[1]:.2f} mm"])

            return table_ras[1], pixel_offset
        
        def align_cbct_to_ct(volumeNode, scan_type, offset, CT_offset=None, CT_spacing=None):
            """
            Aligns CBCT volume to CT volume based on height offset.

            Args:
                volumeNode (vtkMRMLScalarVolumeNode): The volume node to be aligned.
                scan_type (str): The type of scan ("CT" or "CBCT").
                offset (float): The height offset of the current volume from the center in mm.
                CT_offset (float, optional): The height offset of the CT volume from the center. Required for CBCT alignment.
                CT_spacing (float, optional): The voxel spacing of the CT volume in mm (for scaling the offset).

            Returns:
                float: The alignment offset applied to the CBCT volume (if applicable).
            """
            if scan_type == "CT":
                CT_offset = offset
                CT_spacing = volumeNode.GetSpacing()[1]
                #print(f"CT offset set to: {CT_offset}, CT spacing: {CT_spacing} mm/voxel")
                return CT_offset, CT_spacing
            else:
                if CT_offset is None or CT_spacing is None:
                    raise ValueError("CT_offset and CT_spacing must be provided to align CBCT to CT.")

                CBCT_offset = offset

                # Razlika v mm brez skaliranja na CBCT_spacing
                alignment_offset_mm = CT_offset - CBCT_offset

                #print(f"CT offset: {CT_offset}, CBCT offset: {CBCT_offset}")
                #print(f"CT spacing: {CT_spacing} mm/voxel, CBCT spacing: {volumeNode.GetSpacing()[1]} mm/voxel")
                #print(f"Aligning CBCT with CT. Offset in mm: {alignment_offset_mm}")

                # Uporabi transformacijo
                transform = vtk.vtkTransform()
                transform.Translate(0, alignment_offset_mm, 0)

                transformNode = slicer.vtkMRMLTransformNode()
                slicer.mrmlScene.AddNode(transformNode)
                transformNode.SetAndObserveTransformToParent(transform)
                volumeNode.SetAndObserveTransformNodeID(transformNode.GetID())
                slicer.vtkSlicerTransformLogic().hardenTransform(volumeNode)
                slicer.mrmlScene.RemoveNode(transformNode)

                return alignment_offset_mm
    
        def print_orientation(volume_name):
            node = slicer.util.getNode(volume_name)
            matrix = vtk.vtkMatrix4x4()
            node.GetIJKToRASMatrix(matrix)
            print(f"{volume_name} IJK→RAS:")
            for i in range(3):
                print([matrix.GetElement(i, j) for j in range(3)])

        def prealign_by_centroid(cbct_points, ct_points):
            """
            Predporavna CBCT markerje na CT markerje glede na centrične točke.

            Args:
                cbct_points: List ali ndarray točk iz CBCT.
                ct_points: List ali ndarray točk iz CT.

            Returns:
                List: CBCT točke premaknjene tako, da so centrične točke usklajene.
            """
            cbct_points = np.array(cbct_points)
            ct_points = np.array(ct_points)
            cbct_centroid = np.mean(cbct_points, axis=0)
            ct_centroid = np.mean(ct_points, axis=0)
            translation_vector = ct_centroid - cbct_centroid
            aligned_cbct = cbct_points + translation_vector
            return aligned_cbct
        
        def choose_best_translation(cbct_points, ct_points, rotation_matrix):
            """
            Izbere boljšo translacijo: centroidno ali povprečno po rotaciji (retranslation).
            
            Args:
                cbct_points (array-like): Točke iz CBCT (še ne rotirane).
                ct_points (array-like): Ciljne CT točke.
                rotation_matrix (ndarray): Rotacijska matrika.

            Returns:
                tuple: (best_translation_vector, transformed_cbct_points, used_method)
            """
            cbct_points = np.array(cbct_points)
            ct_points = np.array(ct_points)
            
            # 1. Rotiraj CBCT točke
            rotated_cbct = np.dot(cbct_points, rotation_matrix.T)
            
            # 2. Centroid translacija
            centroid_moving = np.mean(cbct_points, axis=0)
            centroid_fixed = np.mean(ct_points, axis=0)
            translation_centroid = centroid_fixed - np.dot(centroid_moving, rotation_matrix)
            transformed_centroid = rotated_cbct + translation_centroid
            error_centroid = np.mean(np.linalg.norm(transformed_centroid - ct_points, axis=1))

            # 3. Retranslacija (srednja razlika)
            translation_recomputed = np.mean(ct_points - rotated_cbct, axis=0)
            transformed_recomputed = rotated_cbct + translation_recomputed
            error_recomputed = np.mean(np.linalg.norm(transformed_recomputed - ct_points, axis=1))

            # 4. Izberi boljšo
            if error_recomputed < error_centroid:
                print(f"✅ Using retranslation (error: {error_recomputed:.2f} mm)")
                return translation_recomputed, transformed_recomputed, "retranslation"
            else:
                print(f"✅ Using centroid-based translation (error: {error_centroid:.2f} mm)")
                return translation_centroid, transformed_centroid, "centroid"

        def visualize_point_matches_in_slicer(cbct_points, ct_points, study_name="MatchVisualization"):
            assert len(cbct_points) == len(ct_points), "Mora biti enako število točk!"

            # Ustvari markups za CBCT
            cbct_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"{study_name}_CBCT")
            cbct_node.GetDisplayNode().SetSelectedColor(0, 0, 1)  # modra

            # Ustvari markups za CT
            ct_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsFiducialNode", f"{study_name}_CT")
            ct_node.GetDisplayNode().SetSelectedColor(1, 0, 0)  # rdeča

            # Dodaj točke
            for i, (cbct, ct) in enumerate(zip(cbct_points, ct_points)):
                cbct_node.AddControlPoint(*cbct, f"CBCT_{i}")
                ct_node.AddControlPoint(*ct, f"CT_{i}")

            # Ustvari model z linijami med pari
            points = vtk.vtkPoints()
            lines = vtk.vtkCellArray()

            for i, (p1, p2) in enumerate(zip(cbct_points, ct_points)):
                id1 = points.InsertNextPoint(p1)
                id2 = points.InsertNextPoint(p2)

                line = vtk.vtkLine()
                line.GetPointIds().SetId(0, id1)
                line.GetPointIds().SetId(1, id2)
                lines.InsertNextCell(line)

            polyData = vtk.vtkPolyData()
            polyData.SetPoints(points)
            polyData.SetLines(lines)

            # Model node
            modelNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLModelNode", f"{study_name}_Connections")
            modelNode.SetAndObservePolyData(polyData)

            modelDisplay = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLModelDisplayNode")
            modelDisplay.SetColor(0, 0, 0)  # črna
            modelDisplay.SetLineWidth(2)
            modelDisplay.SetVisibility(True)

            modelNode.SetAndObserveDisplayNodeID(modelDisplay.GetID())
            modelNode.SetAndObservePolyData(polyData)

            print(f"✅ Vizualizacija dodana za {study_name} (točke + povezave)")
        
        # Globalni seznami za končno statistiko
        prostate_size_est = []
        ctcbct_distance = []

        # Pridobimo SubjectHierarchyNode
        shNode = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
        
        studyItems = vtk.vtkIdList()
        shNode.GetItemChildren(shNode.GetSceneItemID(), studyItems)

        for i in range(studyItems.GetNumberOfIds()):
            studyItem = studyItems.GetId(i)
            studyName = shNode.GetItemName(studyItem)
            print(f"\nProcessing study: {studyName}")
            # **LOKALNI** seznami, resetirajo se pri vsakem study-ju
            cbct_list = []
            ct_list = []
            volume_points_dict = {}
            CT_offset = 0

            # Get child items of the study item
            volumeItems = vtk.vtkIdList()
            shNode.GetItemChildren(studyItem, volumeItems)
            
            # Iteracija čez vse volumne v posameznem studyju
            for j in range(volumeItems.GetNumberOfIds()):
                intermediateItem = volumeItems.GetId(j)

                finalVolumeItems = vtk.vtkIdList()
                shNode.GetItemChildren(intermediateItem, finalVolumeItems)  # Išči globlje!

                for k in range(finalVolumeItems.GetNumberOfIds()):
                    volumeItem = finalVolumeItems.GetId(k)
                    volumeNode = shNode.GetItemDataNode(volumeItem)
                    
                    dicomUIDs = volumeNode.GetAttribute("DICOM.instanceUIDs")
                    if not dicomUIDs:
                        print("❌ This is an NRRD volume!")
                        continue  # Preskoči, če ni DICOM volume             
                    
                                            
                    volumeName = volumeNode.GetName()
                    imageItem = shNode.GetItemByDataNode(volumeNode)
                    modality = shNode.GetItemAttribute(imageItem, "DICOM.Modality")                                 #deluje!

                    #dimensions = volumeNode.GetImageData().GetDimensions()
                    #spacing = volumeNode.GetSpacing()
                    #print(f"Volume {volumeNode.GetName()} - Dimenzije: {dimensions}, Spacing: {spacing}")

                    if modality != "CT":
                        print("Not a CT")
                        continue  # Preskoči, če ni CT

                    # Preveri, če volume obstaja v sceni
                    if not slicer.mrmlScene.IsNodePresent(volumeNode):
                        print(f"Volume {volumeName} not present in the scene.")
                        continue

                    # Preverimo proizvajalca (DICOM metapodatki)
                    manufacturer = shNode.GetItemAttribute(imageItem, 'DICOM.Manufacturer')
                    #manufacturer = volumeNode.GetAttribute("DICOM.Manufacturer")
                    #manufacturer = slicer.dicomDatabase.fileValue(uid, "0008,0070")
                    #print(manufacturer)
                    
                    # Določimo, ali gre za CBCT ali CT
                    if "varian" in manufacturer.lower() or "elekta" in manufacturer.lower():
                        cbct_list.append(volumeName)
                        scan_type = "CBCT"
                        yesCbct = True
                    else:  # Siemens ali Philips
                        ct_list.append(volumeName)
                        scan_type = "CT"
                        yesCbct = False

                    if volumeNode and volumeNode.IsA("vtkMRMLScalarVolumeNode"):
                        print(f"✔️ {scan_type} {volumeNode.GetName()} (ID: {volumeItem})")
                
                    if not volumeNode or not volumeNode.IsA("vtkMRMLScalarVolumeNode"):
                        print("Can't find volumeNode")
                        #continue  # Preskoči, če ni veljaven volume
                        
                        
                    # Detekcija točk v volumnu
                    ustvari_marker = not yesCbct  # Ustvari markerje pred poravnavo na mizo
                    grouped_points = detect_points_region_growing(volumeName, yesCbct, ustvari_marker, intensity_threshold=3000)
                    #print(f"Populating volume_points_dict with key ('{scan_type}', '{volumeName}')")
                    volume_points_dict[(scan_type, volumeName)] = grouped_points
                    #print(volume_points_dict)  # Check if the key is correctly added

            # Če imamo oba tipa volumna (CBCT in CT) **znotraj istega studyja**
            if cbct_list and ct_list:
                ct_volume_name = ct_list[0]  # Uporabi prvi CT kot referenco
                print(f"\nProcessing CT: {ct_volume_name}")
                yesCbct = False
                                               
                mm_offset, pixel_offset = find_table_top_z(ct_volume_name, writefilecheck, yesCbct)
                CT_offset, CT_spacing = align_cbct_to_ct(slicer.util.getNode(ct_volume_name), "CT", mm_offset)                                    
                     
                
                ct_points = [centroid for centroid, _ in volume_points_dict[("CT", ct_volume_name)]]
                
                print(f"CT points: {ct_points}")
                if len(ct_points) < 3:
                    print(f"CT volume {ct_volume_name} doesn't have enough points for registration. Points: {len(ct_points)}")
                    continue
                else:
                    for cbct_volume_name in cbct_list:
                        print(f"\nProcessing CBCT Volume: {cbct_volume_name}")                        
                        yesCbct = True
                        scan_type = "CBCT" #redundant but here we are
                        cbct_volume_node = slicer.util.getNode(cbct_volume_name)
                        
                        mm_offset, pixel_offset = find_table_top_z(cbct_volume_name, writefilecheck, yesCbct)                   
                                            
                        cbct_points = [centroid for centroid, _ in volume_points_dict[("CBCT", cbct_volume_name)]] #zastareli podatki
                        cbct_points_array = np.array(cbct_points)  # Pretvorba v numpy array
                        
                        #print_orientation(ct_volume_name)
                        #print_orientation(cbct_volume_name)
                        
                        #for i, (cb, ct) in enumerate(zip(cbct_points, ct_points)):
                        #    print(f"Pair {i}: CBCT {cb}, CT {ct}, diff: {np.linalg.norm(cb - ct):.2f}")
                        
                        align_cbct_to_ct(cbct_volume_node, scan_type, mm_offset, CT_offset, CT_spacing)
                        #cbct_points = [centroid for centroid, _ in volume_points_dict[("CBCT", cbct_volume_name)]] #zastareli podatki
                        # for i in range(len(cbct_points)):
                        #     x, y, z = cbct_points[i]
                        #     y_new = y + mm_offset  # Upoštevaj premik zaradi poravnave mize
                        #     cbct_points[i] = (x, y_new, z)  # Posodobi marker s premaknjeno višino
                        ustvari_marker = False  # Ustvari markerje
                        cbct_points = [centroid for centroid, _ in detect_points_region_growing(cbct_volume_name, yesCbct, ustvari_marker, intensity_threshold=3000)]                     
                        #cbct_points = detect_points_region_growing(cbct_volume_name, yesCbct, intensity_threshold=3000)   
                        #cbct_points = [centroid for centroid, _ in volume_points_dict[("CBCT", cbct_volume_name)]] #zastareli podatki
                                                
                                                
                        if len(cbct_points) < 3:
                            print(f"CBCT Volume '{cbct_volume_name}' doesn't have enough points for registration. Points: {len(cbct_points)}")
                            continue
                        
                        
                        
                        
                        #Sortiramo točke po X/Y/Z da se izognemo težavam pri poravnavi
                        cbct_points = match_points(cbct_points, ct_points)
                        
                        #visualize_point_matches_in_slicer(cbct_points, ct_points, studyName) #poveže pare markerjev
        
                        
                        # Shranjevanje razdalj
                        distances_ct_cbct = []
                        distances_internal = {"A-B": [], "B-C": [], "C-A": []}

                        cbct_volume_node = slicer.util.getNode(cbct_volume_name)
                        
                        # Sortiramo točke po Z-koordinati (ali X/Y, če raje uporabljaš drugo os)
                        cbct_points_sorted = cbct_points_array[np.argsort(cbct_points_array[:, 2])]

                        # Razdalje med CT in CBCT (SORTIRANE točke!)
                        d_ct_cbct = np.linalg.norm(cbct_points_sorted - ct_points, axis=1)
                        distances_ct_cbct.append(d_ct_cbct)

                        # Razdalje med točkami znotraj SORTIRANIH cbct_points
                        d_ab = np.linalg.norm(cbct_points_sorted[0] - cbct_points_sorted[1])
                        d_bc = np.linalg.norm(cbct_points_sorted[1] - cbct_points_sorted[2])
                        d_ca = np.linalg.norm(cbct_points_sorted[2] - cbct_points_sorted[0])

                        # Sortiramo razdalje po velikosti, da so vedno v enakem vrstnem redu
                        sorted_distances = sorted([d_ab, d_bc, d_ca])

                        distances_internal["A-B"].append(sorted_distances[0])
                        distances_internal["B-C"].append(sorted_distances[1])
                        distances_internal["C-A"].append(sorted_distances[2])
                        
                        # Dodamo ime študije za v statistiko
                        studyName = shNode.GetItemName(studyItem)
                        
                        # **Shrani razdalje v globalne sezname**
                        prostate_size_est.append({"Study": studyName, "Distances": sorted_distances})
                        ctcbct_distance.append({"Study": studyName, "Distances": list(distances_ct_cbct[-1])})  # Pretvorimo v seznam

                        # Izberi metodo glede na uporabnikov izbor
                        chosen_rotation_matrix = np.eye(3)
                        chosen_translation_vector = np.zeros(3)
                        print("Markerji pred transformacijo:", cbct_points, ct_points)
                        if applyScaling:
                            scaling_factors = compute_optimal_scaling_per_axis(cbct_points, ct_points)
                            #print("Scaling factors: ", scaling_factors)
                            cbct_points = compute_scaling(cbct_points, scaling_factors)
                        
                        initial_error = np.mean(np.linalg.norm(np.array(cbct_points) - np.array(ct_points), axis=1))
                        if initial_error > 30:
                            print("⚠️ Initial distance too large, applying centroid prealignment.")
                            cbct_points = prealign_by_centroid(cbct_points, ct_points)
                        
                        if applyRotation:
                            if selectedMethod == "Kabsch":
                                chosen_rotation_matrix = compute_Kabsch_rotation(cbct_points, ct_points)
                            elif selectedMethod == "Horn":
                                chosen_rotation_matrix = compute_Horn_rotation(cbct_points, ct_points)
                            elif selectedMethod == "Iterative Closest Point (Kabsch)":
                                _, chosen_rotation_matrix, _ = icp_algorithm(cbct_points, ct_points)
                            #print("Rotation Matrix:\n", chosen_rotation_matrix)

                        if applyTranslation:
                            chosen_translation_vector, cbct_points_transformed, method_used = choose_best_translation(cbct_points, ct_points, chosen_rotation_matrix) #Izbere optimalno translacijo
                            
                            per_axis_error = np.abs(np.array(cbct_points_transformed) - np.array(ct_points))
                            dz_mean = np.mean(per_axis_error[:, 2])
                            print(f"per-axis error: {per_axis_error}, dz_mean err: {dz_mean:.2f} mm")
                            
                            #chosen_translation_vector = compute_translation(cbct_points, ct_points, chosen_rotation_matrix)
                            #print("Translation Vector:\n", chosen_translation_vector)

                        
                        
                        
                        # Ustvari vtkTransformNode in ga poveži z CBCT volumenom
                        imeTransformNoda = cbct_volume_name + " Transform"
                        transform_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLTransformNode", imeTransformNoda)

                        # Kreiraj transformacijo in jo uporabi
                        vtk_transform = create_vtk_transform(chosen_rotation_matrix, chosen_translation_vector)
                        transform_node.SetAndObserveTransformToParent(vtk_transform)

                        # Pridobi CBCT volumen in aplikacijo transformacije
                        cbct_volume_node = slicer.util.getNode(cbct_volume_name)
                        cbct_volume_node.SetAndObserveTransformNodeID(transform_node.GetID())

                        # Uporabi transformacijo na volumnu (fizična aplikacija)
                        slicer.vtkSlicerTransformLogic().hardenTransform(cbct_volume_node) #aplicira transformacijo na volumnu
                        slicer.mrmlScene.RemoveNode(transform_node) # Odstrani transformacijo iz scene
                        #print("Transform successful on ", cbct_volume_name)
                        ustvari_marker = True  # Ustvari markerje
                        cbct_points = [centroid for centroid, _ in detect_points_region_growing(cbct_volume_name, yesCbct, ustvari_marker, intensity_threshold=3000)]
                        print("Markerji po transformaciji:\n", cbct_points, ct_points)
                        
                        # Izračun razdalj
                        errors = [np.linalg.norm(cbct - ct) for cbct, ct in zip(cbct_points, ct_points)]
                        mean_error = np.mean(errors)

                        print("Individualne napake:", errors)
                        print("📏 Povprečna napaka poravnave: {:.2f} mm".format(mean_error))
                        
                        

            else:
                print(f"Study {studyItem} doesn't have any appropriate CT or CBCT volumes.")

        # Izpis globalne statistike
        
        
        if writefilecheck:
            print("Distances between CT & CBCT markers: ", ctcbct_distance)
            print("Distances between pairs of markers for each volume: ", prostate_size_est)
            
            # Define file paths
            prostate_size_file = os.path.join(os.path.dirname(__file__), "prostate_size.csv")
            ctcbct_distance_file = os.path.join(os.path.dirname(__file__), "ct_cbct_distance.csv")

            # Write prostate size data
            with open(prostate_size_file, mode='w', newline='') as file:
                writer = csv.writer(file)
                writer.writerow(["Prostate Size"])
                for size in prostate_size_est:
                    writer.writerow([size])
            print("Prostate size file written at ", prostate_size_file)

            # Write CT-CBCT distance data
            with open(ctcbct_distance_file, mode='w', newline='') as file:
                writer = csv.writer(file)
                writer.writerow(["CT-CBCT Distance"])
                for distance in ctcbct_distance:
                    writer.writerow([distance])
            print("CT-CBCT distance file written at ", ctcbct_distance_file)

        end_time = time.time()
        # Calculate and print elapsed time
        elapsed_time = end_time - start_time
        # Convert to minutes and seconds
        minutes = int(elapsed_time // 60)
        seconds = elapsed_time % 60

        print(f"Execution time: {minutes} minutes and {seconds:.6f} seconds")