import cv2
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.colors import LinearSegmentedColormap, TwoSlopeNorm
import os

colors = [
    (67, 65, 143),
    (255, 255, 255),  
    (230, 160, 2)
]

colors_normalized = [(r/255, g/255, b/255) for r, g, b in colors]

custom_cmap = LinearSegmentedColormap.from_list("custom_cmap", colors_normalized)

color_scores = {
    'red': 6,
    'orange': 5,
    'yellow': 4,
    'green': 3,
    'blue': 2,
    'purple': 1,
    'other': 0
}


def get_color_score(bgr):
    blue, green, red = bgr
    if blue < 100 and green < 100 and red > 190:
        return color_scores['red']
    elif blue < 100 and 100 < green < 150 and red > 150:
        return color_scores['orange']
    elif blue < 100 and green > 150 and red > 150:
        return color_scores['yellow']
    elif blue < 160 and green > 90 and red < 160:
        return color_scores['green']
    elif blue > 160 and green < 90 and red < 90:
        return color_scores['blue']
    elif blue > 100 and green < 100 and red > 100:
        return color_scores['purple']
    else:
        return color_scores['other']


ctrl_image_path = '******.jpg'
den_image_path = '******.jpg'

ctrl_image = cv2.imread(ctrl_image_path)
den_image = cv2.imread(den_image_path)

height, width, _ = map(int, ctrl_image.shape)

ctrl_scores = np.zeros((height, width), dtype=np.float32)
den_scores = np.zeros((height, width), dtype=np.float32)

for y in range(height):
    for x in range(width):
        ctrl_scores[y, x] = get_color_score(ctrl_image[y, x])
        den_scores[y, x] = get_color_score(den_image[y, x])

diff = ctrl_scores - den_scores

min_val = -6
max_val = 6
normalized_diff = (diff - min_val) / (max_val - min_val) * 2 - 1

normalized_diff = np.clip(normalized_diff, -1, 1)


norm = TwoSlopeNorm(vmin=-1, vcenter=0, vmax=1)

plt.figure(figsize=(10, 8))
plt.imshow(normalized_diff, cmap=custom_cmap, norm=norm, interpolation='nearest')
plt.colorbar(label='Difference Magnitude')
plt.axis('off')

output_folder = 'Heatmaps'
os.makedirs(output_folder, exist_ok=True)

heatmap_path = os.path.join(output_folder, 'EP.png')
plt.savefig(heatmap_path)
print(f"The corrected heatmap image has been saved to ‘{heatmap_path}’.")

plt.show()

red_side_total = np.sum(normalized_diff[normalized_diff > 0])
blue_side_total = np.sum(normalized_diff[normalized_diff < 0])

plt.figure(figsize=(6, 5))
plt.title("Sum of Normalized Differences by Color Tone")
bars = plt.bar(['Blue (Positive)', 'Red (Negative)'], [red_side_total, abs(blue_side_total)], color=['blue', 'red'])
plt.ylabel('Sum of Normalized Differences (Absolute)')

bargraph_path = os.path.join(output_folder, '******')
plt.savefig(bargraph_path)
print(f"The difference sum graph has been saved to ‘{bargraph_path}’.")

plt.show()

data_with_coords = []

for y in range(height):
    for x in range(width):
        data_with_coords.append({
            'x': x,
            'y': y,
            'normalized_diff': normalized_diff[y, x]
        })

df_with_coords = pd.DataFrame(data_with_coords)

df_positive = df_with_coords[df_with_coords['normalized_diff'] > 0]

df_negative = df_with_coords[df_with_coords['normalized_diff'] < 0]

positive_path = os.path.join(output_folder, 'normalized_diff_positive_only.xlsx')
negative_path = os.path.join(output_folder, 'normalized_diff_negative_only.xlsx')

df_positive.to_excel(positive_path, index=False)
df_negative.to_excel(negative_path, index=False)

print(f"Only positive differences were saved to ‘{positive_path}’.")
print(f"Only negative differences have been saved to ‘{negative_path}’.")