import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from scipy.interpolate import interp1d
import os
import uuid


def calculate_intensity(image):
    """
    Calculate the intensity of the image by converting it to grayscale
    and taking the mean value.
    """
    gray_image = image.convert('L')  # Convert to grayscale
    intensity = np.mean(np.array(gray_image))
    return intensity


def process_image(image_path):
    """
    Process the image by dividing it into 32x32 grids and calculating normalized intensity.
    """
    # Open the image
    image = Image.open(image_path)
    width, height = image.size

    # Determine the ideal grid size
    grid_width = width // 32
    grid_height = height // 32

    intensities = []

    # Divide the image into grids and calculate intensity for each grid
    for y in range(32):
        row_intensities = []
        for x in range(32):
            left = x * grid_width
            upper = y * grid_height
            right = left + grid_width
            lower = upper + grid_height

            # Crop the image to the grid
            grid = image.crop((left, upper, right, lower))
            intensity = calculate_intensity(grid)
            row_intensities.append(intensity)
        intensities.append(row_intensities)

    intensities = np.array(intensities)

    # Normalize the intensities to the range 0-1
    min_intensity = intensities.min()
    max_intensity = intensities.max()
    normalized_intensities = (intensities - min_intensity) / (max_intensity - min_intensity)

    return normalized_intensities


def intensity_to_oxygen_level(normalized_intensities):
    """
    Map normalized intensity values to oxygen levels using interpolation.
    """
    # Provided mapping table
    ksv_values = [0, 0.009303080328, 0.19176, 0.5196879311, 0.6252571452, 0.7115134738,
                  0.8476158622, 0.9531850764, 1.039441405, 1.145010619, 1.175543793,
                  1.231266948, 1.281113007, 1.47293855, 1.609040939]
    oxygen_levels = [0, 0.68, 1, 2, 2.5, 3, 4, 5, 6, 7.5, 8, 9, 10, 15, 20]

    # Normalize intensities to match the range of Ksv values
    min_ksv = min(ksv_values)
    max_ksv = max(ksv_values)
    normalized_intensities = (normalized_intensities * (max_ksv - min_ksv)) + min_ksv

    # Interpolation function
    interp_func = interp1d(ksv_values, oxygen_levels, kind='linear', fill_value='extrapolate')
    oxygen_levels_mapped = interp_func(normalized_intensities)

    return oxygen_levels_mapped


def generate_heatmap(data, output_path):
    """
    Generate a high-resolution heatmap from data and save it as an image.
    """
    plt.figure(figsize=(10, 8))
    plt.imshow(data, cmap='RdBu_r', interpolation='nearest', norm=mcolors.Normalize(vmin=0, vmax=20))
    plt.colorbar(label='Oxygen Level (%)')

    # Save the figure with high resolution
    plt.savefig(output_path, dpi=300)
    plt.close()


def process_and_generate_heatmap(image):
    """
    Process the image and generate a heatmap, saving the result to a file.
    """
    # Generate a unique name for the input image
    unique_id = str(uuid.uuid4())
    input_dir = 'Gradio_Images'
    input_image_path = os.path.join(input_dir, f'uploaded_image_{unique_id}.jpg')
    image.save(input_image_path)

    # Process the image to get normalized intensities
    normalized_intensities = process_image(input_image_path)
    oxygen_levels = intensity_to_oxygen_level(normalized_intensities)

    # Generate the output path for the heatmap based on the unique name
    output_dir = 'Heatmap_Images'

    output_image_path = os.path.join(output_dir, f'heatmap_{unique_id}.jpg')

    generate_heatmap(oxygen_levels, output_image_path)

    return output_image_path