import os
import cv2
import pandas as pd
import numpy as np
from skimage.metrics import structural_similarity as ssim
from haarPsi import haar_psi_numpy


def adjust_contrast(image, amount):
    """
    Adjust the contrast of the input image.

    :param image: Input image (numpy array)
    :param amount: Amount of contrast adjustment
    :return: Adjusted image (numpy array)
    """
    LUT = np.interp(
        np.arange(256),
        np.array([0, 64, 192, 256]),
        np.array([0, (0.25 - amount / 4) * 255, (0.75 + amount / 4) * 255, 255]),
    ).astype("uint8")
    adjusted_image = cv2.LUT(image, LUT)
    return adjusted_image


def gradient_descent(x_start, gradient, step_size, direction_changes):
    x = x_start
    n = 0  
    direction = np.sign(gradient(x))  
    
    while n < direction_changes:
        x -= step_size * direction * gradient(x)  
        new_direction = np.sign(gradient(x))   
        
        if new_direction != direction:
            step_size *= 0.5  
            n += 1   
            direction = new_direction  
            
    return x


def main(input_folder, output_dir, save_images=False):
    # Constants and Hyperparameters
    TARGET_HAAR_PSI = 0.85
    LEARNING_RATE = 0.05
    NUM_ITERATIONS = 200
    ERROR_WEIGHT = 0.1
    CONVERGENCE_THRESHOLD = 0.005
    DIRECTION_CHANGE_LIMIT = 10

    results = []

    for file in os.listdir(input_folder):
        file_path = os.path.join(input_folder, file)

        # Read the image
        image = cv2.imread(file_path, cv2.IMREAD_COLOR)

        # Check if the image is valid (not None)
        if image is not None:
            contrast_amount = 0.44
            prev_contrast_amount = 0.0  # Store the previous contrast_amount

        # Gradient Descent Loop
        for i in range(NUM_ITERATIONS):
            # Apply contrast modification
            contrast_image = adjust_contrast(image, contrast_amount)

            # Calculate haar_psi_numpy() for the current contrast
            haar_psi_c, _, _ = haar_psi_numpy(image, contrast_image)

            # Calculate the error and derivative
            error = haar_psi_c - TARGET_HAAR_PSI
            d_contrast = np.sign(error) * ERROR_WEIGHT

            # Update contrast_amount 
            contrast_amount = gradient_descent(contrast_amount, d_contrast, LEARNING_RATE, DIRECTION_CHANGE_LIMIT)

            # Clip contrast_amount between 0.3 and 0.8
            contrast_amount = np.clip(contrast_amount, 0.3, 0.8)

                # Print the status after each iteration
            print(
                    f"Image: {file}  Iteration: {i + 1}/{NUM_ITERATIONS}  Contrast Amount: {contrast_amount:.4f}  haar_psi_numpy: {haar_psi_c:.4f}"
                )

                # Check convergence criterion
            if abs(contrast_amount - prev_contrast_amount) < CONVERGENCE_THRESHOLD:
                    print("Convergence criterion met. Stopping the gradient descent loop.")
                    break

            prev_contrast_amount = contrast_amount

            # Modify the image using the updated contrast_amount value
            contrast_image = adjust_contrast(image, contrast_amount)

            if save_images:
                # Save distorted images
                cv2.imwrite(os.path.join(output_dir, file[:-4] + "_contrast.jpg"), contrast_image)

            # Append results to the list
            results.append([file, "contrast", contrast_amount, round(haar_psi_c, 3)])

    # Create a pandas DataFrame from the results list
    df = pd.DataFrame(results, columns=["Image Name", "Distortion Type", "Parameter", "haar_psi_numpy"])

    return df


input_folder = "C:/Users/bob/Downloads/100sampled/Sub/"
output_dir = "C:/Users/bob/Downloads/100sampled/Sub/distorted_images"
save_images = False

# Create output directory if it doesn't exist
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# Run the main function and get the DataFrame
results_df = main(input_folder, output_dir, save_images)
print(results_df)

# Group the data by 'Distortion Type' column
grouped_df = results_df.groupby("Distortion Type")

Returns
  Cell In[4], line 35 in gradient_descent
    direction = np.sign(gradient(x))

TypeError: 'numpy.float64' object is not callable