import nibabel as nib
import numpy as np
import plotly.graph_objects as go
from scipy.ndimage import gaussian_filter
from plotly.subplots import make_subplots

def load_and_preprocess_nifti(nifti_path):
    print(f"Loading NIFTI file from {nifti_path}")
    nifti_img = nib.load(nifti_path)
    nifti_data = nifti_img.get_fdata()
    print("Normalizing data")
    normalized_data = (nifti_data - np.min(nifti_data)) / (np.max(nifti_data) - np.min(nifti_data))
    print("Data loaded and preprocessed")
    return normalized_data

def save_as_nrrd(data, file_path):
    """Save the numpy array as NRRD file if needed."""
    nrrd.write(file_path, data)

def downsample_3d(data, factor):
    """Downsample 3D data by a factor"""
    return data[::factor, ::factor, ::factor]

def compute_hog3d(ct_data, block_size=8, cell_size=4, stride=4, num_bins=12):
    print("Starting HOG3D computation")

    # Apply Gaussian smoothing to the data
    # Sigma value is set to 1, increasing it will result in more smoothing
    smoothed_data = gaussian_filter(ct_data, sigma=1)

    # Compute the gradient of the smoothed data
    gx, gy, gz = np.gradient(smoothed_data)

    # Compute the magnitude, phi and theta values
    magnitude = np.sqrt(gx**2 + gy**2 + gz**2)

    # Phi is the angle between the x-axis and the projection of the gradient vector onto the xy-plane
    phi = np.arctan2(np.sqrt(gx**2 + gy**2), gz)

    # Theta is the angle between the z-axis and the gradient vector
    theta = np.arctan2(gy, gx)

    # Initialize the lists to store the HOG3D features, positions and orientations
    hog3d_features = []
    hog3d_positions = []
    hog3d_orientations = []

    # Calculate total blocks and initialize the current block counter
    total_blocks = ((ct_data.shape[0] - block_size) // stride + 1) * \
                   ((ct_data.shape[1] - block_size) // stride + 1) * \
                   ((ct_data.shape[2] - block_size) // stride + 1)
    current_block = 0



    for x in range(0, ct_data.shape[0] - block_size + 1, stride): # Loop over the x-axis
        for y in range(0, ct_data.shape[1] - block_size + 1, stride): # Loop over the y-axis
            for z in range(0, ct_data.shape[2] - block_size + 1, stride): # Loop over the z-axis
                block_hist = np.zeros((2 * num_bins, (block_size // cell_size) ** 3)) # Initialize the block histogram

                cell_index = 0
                for i in range(0, block_size, cell_size): # Loop over the x-axis of the block
                    for j in range(0, block_size, cell_size): # Loop over the y-axis of the block
                        for k in range(0, block_size, cell_size): # Loop over the z-axis of the block

                            # Extract the cell from the phi, theta and magnitude arrays
                            cell_phi = phi[x+i:x+i+cell_size, y+j:y+j+cell_size, z+k:z+k+cell_size]
                            cell_theta = theta[x+i:x+i+cell_size, y+j:y+j+cell_size, z+k:z+k+cell_size]
                            cell_magnitude = magnitude[x+i:x+i+cell_size, y+j:y+j+cell_size, z+k:z+k+cell_size]

                            # 1D histogram for phi angles
                            hist_phi, _ = np.histogram(
                                cell_phi.ravel(),
                                bins=num_bins,
                                range=(-np.pi, np.pi),
                                weights=cell_magnitude.ravel()
                            )

                            # 1D histogram for theta angles
                            hist_theta, _ = np.histogram(
                                cell_theta.ravel(),
                                bins=num_bins,
                                range=(-np.pi, np.pi),
                                weights=cell_magnitude.ravel()
                            )

                            # Concatenate the histograms for phi and theta
                            block_hist[:, cell_index] = np.concatenate([hist_phi, hist_theta])
                            cell_index += 1

                block_hist = block_hist.ravel()  # Reshape the block_hist array
                block_hist /= np.linalg.norm(block_hist) + 1e-5  # Normalize the histogram

                hog3d_features.append(block_hist) # Append the block histogram to the HOG3D features list
                hog3d_positions.append([x + block_size // 2, y + block_size // 2, z + block_size // 2]) # Append the block position to the HOG3D positions list
                hog3d_orientations.append([np.mean(cell_phi), np.mean(cell_theta)]) # Append the block orientation to the HOG3D orientations list

                # statements to keep track of the progress
                current_block += 1
                if current_block % 10000 == 0 or current_block == total_blocks:
                    print(f"Processed {current_block}/{total_blocks} blocks")

    print("HOG3D computation completed")
    return np.array(hog3d_features), np.array(hog3d_positions), np.array(hog3d_orientations)

def visualize_hog3d(ct_data, hog3d_features, hog3d_positions, hog3d_orientations, threshold=0.75, save_path=None):
    print("Starting visualization")
    feature_magnitudes = np.linalg.norm(hog3d_features, axis=1)

    # Use a more discriminative threshold
    magnitude_threshold = np.percentile(feature_magnitudes, threshold * 100)
    mask = feature_magnitudes > magnitude_threshold

    filtered_positions = hog3d_positions[mask]
    filtered_magnitudes = feature_magnitudes[mask]
    filtered_orientations = hog3d_orientations[mask]

    # Apply non-linear transformation to magnitudes
    filtered_magnitudes = np.power(filtered_magnitudes, 0.3)  # Adjust the power for desired spread

    # Find the bounding box of important features
    min_x, min_y, min_z = np.min(filtered_positions, axis=0)
    max_x, max_y, max_z = np.max(filtered_positions, axis=0)

    # Add some padding to the bounding box
    padding = 10
    min_x, min_y, min_z = max(0, min_x - padding), max(0, min_y - padding), max(0, min_z - padding)
    max_x, max_y, max_z = min(ct_data.shape[0], max_x + padding), min(ct_data.shape[1], max_y + padding), min(ct_data.shape[2], max_z + padding)

    # Create subplots
    fig = make_subplots(
        rows=1, cols=3,
        specs=[[{'type': 'scene'}, {'type': 'scene'}, {'type': 'scene'}]],
        subplot_titles=("Feature Positions", "Orientations", "Combined")
    )

    # 1. Feature position points
    fig.add_trace(
        go.Scatter3d(
            x=filtered_positions[:, 0],
            y=filtered_positions[:, 1],
            z=filtered_positions[:, 2],
            mode='markers',
            marker=dict(
                size=2,
                color=filtered_magnitudes,
                colorscale='Viridis',
                opacity=1,
                colorbar=dict(title="Feature Magnitude", x=0.3)
            ),
            name='HOG3D Features'
        ),
        row=1, col=1
    )

    # 2. Orientation cones
    u = np.sin(filtered_orientations[:, 0]) * np.cos(filtered_orientations[:, 1])
    v = np.sin(filtered_orientations[:, 0]) * np.sin(filtered_orientations[:, 1])
    w = np.cos(filtered_orientations[:, 0])

    scale_factor = 5
    u *= scale_factor
    v *= scale_factor
    w *= scale_factor

    fig.add_trace(
        go.Cone(
            x=filtered_positions[:, 0],
            y=filtered_positions[:, 1],
            z=filtered_positions[:, 2],
            u=u,
            v=v,
            w=w,
            colorscale='Plotly3',
            sizemode="absolute",
            opacity=0.2,
            sizeref=3,

        ),
        row=1, col=2
    )

    # 3. Combined visualization
    fig.add_trace(
        go.Scatter3d(
            x=filtered_positions[:, 0],
            y=filtered_positions[:, 1],
            z=filtered_positions[:, 2],
            mode='markers',
            marker=dict(
                size=1,
                color=filtered_magnitudes,
                colorscale='Viridis',
                opacity=1,

            ),
            name='HOG3D Features'
        ),
        row=1, col=3
    )

    fig.add_trace(
        go.Cone(
            x=filtered_positions[:, 0],
            y=filtered_positions[:, 1],
            z=filtered_positions[:, 2],
            u=u,
            v=v,
            w=w,
            colorscale='Plotly3',
            sizemode="absolute",
            opacity=0.4,
            sizeref=5,

        ),
        row=1, col=3
    )

    # Update layout for all subplots
    for i in range(1, 4):
        fig.update_scenes(
            xaxis=dict(title="X", range=[min_x, max_x]),
            yaxis=dict(title="Y", range=[min_y, max_y]),
            zaxis=dict(title="Z", range=[min_z, max_z]),
            aspectmode='data',
            row=1, col=i
        )

    fig.update_layout(
        width=5000,
        height=1000,
        title="HOG3D Feature and Orientation Visualization (Cropped)"
    )

    fig.show()

    if save_path:
        fig.write_html(save_path)
        print(f"Visualization saved to {save_path}")

    print("Visualization completed")

!wget https://raw.githubusercontent.com/Pranav-Karra-3301/CVD_HOG3D/f6a7011b2edcf5a0bbdd2094a75a0b793945dd68/Img_001.nii.gz

nifti_file = "Img_001.nii.gz" # Replace with the correct file path
ct_data = load_and_preprocess_nifti(nifti_file)
print("Processed data shape:", ct_data.shape)

downsample_factor = 4
downsampled_ct_data = downsample_3d(ct_data, downsample_factor)
print("Downsampled data shape:", downsampled_ct_data.shape)

# Calculate expected shapes
"""
Block Size: Increase to capture more information in each block
Cell Size: Increase to capture more information in each cell
Stride: Increase to reduce the number of blocks
Num Bins: Increase to capture more information in each histogram.
"""
block_size = 4
cell_size = 2
stride = 2
num_bins = 18

num_blocks_x = (downsampled_ct_data.shape[0] - block_size) // stride + 1
num_blocks_y = (downsampled_ct_data.shape[1] - block_size) // stride + 1
num_blocks_z = (downsampled_ct_data.shape[2] - block_size) // stride + 1
total_blocks = num_blocks_x * num_blocks_y * num_blocks_z

features_per_block = (2 * num_bins) * ((block_size // cell_size) ** 3)


# Compute HOG3D
hog3d_features, hog3d_positions, hog3d_orientations = compute_hog3d(downsampled_ct_data, block_size=block_size, cell_size=cell_size, stride=stride, num_bins=num_bins)

# Print expected and computed shapes alternately
print("")
print("Downsample Factor:", downsample_factor)
print("Total Blocks: ",(total_blocks))
print(f"Block Size: {block_size}, Cell Size: {cell_size}, Stride: {stride}, Num Bins: {num_bins}")
print("")
print("Expected HOG3D features shape:", (total_blocks, features_per_block))
print("Computed HOG3D features shape:", hog3d_features.shape)
print("")
print("Expected HOG3D positions shape:", (total_blocks, 3))
print("Computed HOG3D positions shape:", hog3d_positions.shape)
print("")
print("Expected HOG3D orientations shape:", (total_blocks, 2))
print("Computed HOG3D orientations shape:", hog3d_orientations.shape)

# Visualize HOG3D features
save_path = "hog3d_visualization.html"  # Set to None if you don't want to save
visualize_hog3d(ct_data, hog3d_features, hog3d_positions, hog3d_orientations, threshold=0.2, save_path=save_path)