"""
Created on Thu Dec 10 10:37:56 2020
@author: Anja Hajnal

Adjusted on 31.03.2026
@author: Jan Hünting

1. Place the binary black-and-white images in the "Images" folder.
2. Adjust the parameters in the "User settings" section if required.
3. Run the script.
4. The generated plots and Excel results are automatically saved in the "Plots" and "Results" folders, respectively. 
"""

import cv2
import numpy as np
import glob
import matplotlib.pyplot as plt
import xlsxwriter
import os


# ============================================================
# User settings
# ============================================================

# Sample information
sample_name = "XX"

# Input/output directories
directory = r"Images"
plot_dir = r"Plots"
results_dir = r"Results" 

# Image acquisition
dt = 0.5  # time between consecutive images [min]

# Image segmentation
pixel_threshold = 50       # grayscale threshold for black/white segmentation
height_threshold = 0.3     # relative threshold for height determination
gaussian_kernel_size = 51  # kernel size for smoothing row widths

# Calibration
known_distance = 3.9  # known reference distance [cm]

# Moving-average smoothing window [number of data points]
smoothing_window = 21



print("Script started")

os.makedirs(plot_dir, exist_ok=True)
os.makedirs(results_dir, exist_ok=True)

path = os.path.join(directory, "*.jpg")

filenames = glob.glob(path)
filenames.sort()

if len(filenames) == 0:
    raise ValueError("No images found. Check path")

images = [cv2.imread(image) for image in filenames]

print(f"{len(images)} images loaded")   


volume_all = []
mid_d_all = []
h_all = []


# Calibration
gray_0 = cv2.cvtColor(images[0], cv2.COLOR_BGR2GRAY)
local_pixel_d0 = np.count_nonzero(gray_0 <= pixel_threshold, axis=1)
median_pixel_d = np.median(local_pixel_d0)
calib = median_pixel_d / known_distance
print("Calibration factor (pixel/cm):", calib)

for i in range(len(images)):  

        # Convert image to grayscale
        gray = cv2.cvtColor(images[i], cv2.COLOR_BGR2GRAY)
       
        # Local width in pixels and centimeters
        local_pixel_d = np.count_nonzero(gray <= pixel_threshold, axis=1) 
        # Convert to centimeters
        local_d = local_pixel_d / calib

        gel = local_d[local_d > 0]         

        # Spatial increment corresponding to one pixel [cm]    
        dx = 1 / calib

        # Local and total proxy volume       
        volume = np.sum(local_d**2) * dx
        volume_all.append(volume) 

        # Mean local width         
        if len(gel) > 0:
         mid_d = np.mean(gel)
        else:
         mid_d = 0

        mid_d_all.append(mid_d)
 #     
        # Gel Height        
        mask = (gray <= pixel_threshold)

        # number of black pixels per row
        row_width = np.sum(mask, axis=1)

        # Smoothing
        row_width_smooth = cv2.GaussianBlur(row_width.astype(np.float32), (gaussian_kernel_size, 1), 0)

        if i % 50 == 0:
            print("min/max width:", np.min(row_width), np.max(row_width)) 

        # Normalization
        row_width_norm = row_width_smooth / np.max(row_width_smooth)

        gel_rows = np.where(
            row_width_norm > height_threshold
        )[0]

        if len(gel_rows) > 0:
            y_min = gel_rows[0]
            y_max = gel_rows[-1]
            height_px = y_max - y_min
        else:
            height_px = 0 

        height_cm = height_px / calib
        h_all.append(height_cm)


time = [i * dt for i in range(len(mid_d_all))] 

# --- Normalization after complete calculation ---
V0 = volume_all[0] if volume_all[0] != 0 else 1
normalized_volume = [v / V0 for v in volume_all]

# --- Smoothing of normalized volume ---
normalized_volume_smooth = np.convolve(
    normalized_volume,
    np.ones(smoothing_window) / smoothing_window,
    mode='same' 
)

D0 = mid_d_all[0] if mid_d_all[0] != 0 else 1
normalized_mid_d_all = [d / D0 for d in mid_d_all]

H0 = h_all[0] if h_all[0] != 0 else 1
normalized_h = [h / H0 for h in h_all]  


# Plot the data       
# normalized width
plt.clf()
plt.figure(1)
plt.plot(time, normalized_mid_d_all, 'b') 
plt.title(sample_name) 
plt.xlabel('Time [min]')
plt.ylabel('Normalized mean local width D/D(0) [-]')  

plt.savefig(os.path.join(plot_dir, "Width.png"))

# normalized height
plt.clf() 
plt.figure(2)
plt.plot(time, normalized_h, 'g') 
plt.title(sample_name)   
plt.xlabel('Time [min]')
plt.ylabel('Normalized height H/H(0) [-]') 

plt.savefig(os.path.join(plot_dir, "Height.png"))

# Plot of height and width
plt.clf()
plt.figure(3)
plt.plot(time, normalized_mid_d_all, 'b') 
plt.plot(time, normalized_h, 'g')
plt.title(sample_name)  
plt.xlabel('Time [min]')
plt.ylabel('Dimensionless unit [-]')

plt.savefig(os.path.join(plot_dir, "Height_Width.png")) 

# normalized volume
plt.clf()
plt.figure(4)
plt.plot(time, normalized_volume, 'r') 
plt.title(sample_name)
plt.xlabel('Time [min]')
plt.ylabel('Normalized volume V/V(0) [-]')

plt.savefig(os.path.join(plot_dir, "Normalized_volume.png")) 

# smoothed normalized volume
plt.clf()
plt.figure(5)
plt.plot(time, normalized_volume, 'r', alpha=0.3, label='raw')
plt.plot(time, normalized_volume_smooth, 'k', label='smoothed')
plt.title(sample_name)
plt.xlabel('Time [min]')
plt.ylabel('Smoothed normalized volume V/V(0) [-]') 

plt.savefig(os.path.join(plot_dir, "Normalized_volume_smoothed.png"))  


# Save data to excel file
columns = [time, volume_all, normalized_volume, mid_d_all, normalized_mid_d_all, h_all, normalized_h]

# Specify the path and file name 
excel_path = os.path.join(results_dir, "Results_Kinetics.xlsx")
workbook = xlsxwriter.Workbook(excel_path) 

worksheet = workbook.add_worksheet() 
worksheet.write_string(0,0,'Time [min]') 
worksheet.write_string(0,1,'Geometric volume proxy [cm^3]')
worksheet.write_string(0,2,'Normalized volume V/V0 [-]')
worksheet.write_string(0,3,'Mean local width [cm]') 
worksheet.write_string(0,4,'Normalized mean local width D/D0 [-]')
worksheet.write_string(0,5,'Height [cm]')
worksheet.write_string(0,6,'Normalized height H/H0 [-]')
for col,data in enumerate(columns):
    worksheet.write_column(1,col,data) 
    
workbook.close()