"""Simple layer profile plots"""

import numpy as np
import nibabel as nb
import matplotlib.pyplot as plt

# Output nifti from 01_simulate_layers
FILE1 = "/home/faruk/gdrive/LAYNII/demo_big3/M_brain_rim_metric_equidist_simulated_layers_noised.nii.gz"
# FILE1 = "/home/faruk/gdrive/LAYNII/demo_big3/M_brain_rim_metric_equidist_simulated_layers_brain_draining_veins_noised.nii.gz"

# Metric file generated by LN2_LAYERS
FILE2 = "/home/faruk/gdrive/LAYNII/demo_big3/M_brain_rim_metric_equidist.nii.gz"
FILE3 = "/home/faruk/gdrive/LAYNII/demo_big3/M_brain_rim_layers_equidist.nii.gz"

# Column or area definitions
FILE4 = "/home/faruk/gdrive/LAYNII/demo_big3/M_brain_rim_columns33.nii.gz"

# =============================================================================
# Load image data, you can think of this as an activation map.
nii1 = nb.load(FILE1)
data = nii1.get_fdata()
# Load normalized depths
nii2 = nb.load(FILE2)
norm_depth = nii2.get_fdata()
# Load layers (quantized normalized depths)
nii3 = nb.load(FILE3)
layers = nii3.get_fdata()
idx_layers = np.unique(layers.astype("int"))
idx_layers = idx_layers[1:]  # Remove layers with index 0
nr_layers = idx_layers.size
layer_bins = np.zeros(nr_layers + 1)  # +1 for zeros indices

# Load columns
nii4 = nb.load(FILE4)
colums = nii4.get_fdata()
idx_columns = np.unique(colums)
idx_columns = idx_columns[1:]  # Remove columns with index 0

# Prepare plot
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle('Layer profiles in two alternative ways')

for j in idx_columns:
    # Take voxels of a single column
    idx = colums == j
    data_roi = data[idx]
    norm_depth_roi = norm_depth[idx]
    layers_roi = layers[idx]

    # Before layer quantization
    ax1.scatter(norm_depth_roi, data_roi, alpha=0.1, marker=".")
    ax1.set_xlim((0, 1))
    ax1.set_ylim((90, 120))
    ax1.set_xlabel("Normalized cortical depth")
    ax1.set_ylabel("Voxel value")

    # After layer quantization
    for i in idx_layers:  # Compute bin averages
        layer_bins[i] = np.mean(data_roi[layers_roi == i])
    ax2.plot(idx_layers, layer_bins[1:], linewidth=1)
    ax2.set_ylim((90, 120))
    ax2.set_xlabel("Layers (0=white matter)")
    ax2.set_ylabel("Voxel value")

plt.show()

print("Finished.")
