import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
import gzip

def split_fft(image):
    fshift_channels = []
    magnitude_spectrums = []
    for i in range(3):  # Process B, G, R channels separately
        channel = image[:, :, i]
        f = np.fft.fft2(channel)
        fshift = np.fft.fftshift(f)
        magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)
        fshift_channels.append(fshift)
        magnitude_spectrums.append(magnitude_spectrum)
    return fshift_channels, magnitude_spectrums

def save_fft_parts_compressed(path_prefix, fshift_channels):
    for i, fshift in enumerate(fshift_channels):
        file_path = f"{path_prefix}_channel_{i}.npz.gz"
        with gzip.GzipFile(file_path, 'w') as f:
            np.save(f, fshift)
        print(f"Saved compressed FFT part: {file_path}")

def save_magnitude_spectrums_as_images(path_prefix, magnitude_spectrums):
    for i, mag in enumerate(magnitude_spectrums):
        normalized = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
        img = normalized.astype(np.uint8)
        file_path = f"{path_prefix}_magnitude_{i}.png"
        cv2.imwrite(file_path, img)
        print(f"Saved magnitude spectrum image: {file_path}")

def load_fft_parts_compressed(path_prefix):
    fshift_channels = []
    for i in range(3):
        file_path = f"{path_prefix}_channel_{i}.npz.gz"
        with gzip.GzipFile(file_path, 'r') as f:
            part = np.load(f)
        fshift_channels.append(part)
        print(f"Loaded compressed FFT part: {file_path}")
    return fshift_channels

def merge_fft(fshift_channels):
    reconstructed_channels = []
    for fshift in fshift_channels:
        f_ishift = np.fft.ifftshift(fshift)
        img_back = np.fft.ifft2(f_ishift)
        img_back = np.abs(img_back)
        img_back = np.clip(img_back, 0, 255).astype(np.uint8)
        reconstructed_channels.append(img_back)
    reconstructed_image = cv2.merge(reconstructed_channels)
    return reconstructed_image

def load_image(path):
    if not os.path.exists(path):
        raise FileNotFoundError(f"Image file '{path}' not found")
    image = cv2.imread(path)
    if image is None:
        raise ValueError(f"Failed to load image from '{path}'")
    return image

def plot_images(original=None, magnitude_spectrums=None, reconstructed=None):
    if original.any():
        plt.figure(figsize=(12, 6))
        plt.subplot(1, 3, 1)
        plt.imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
        plt.title('Original Image')
        plt.axis('off')

    if magnitude_spectrums:
        plt.subplot(1, 3, 2)
        combined_mag = np.stack(magnitude_spectrums, axis=2)
        combined_mag = np.clip(combined_mag / np.max(combined_mag), 0, 1)
        plt.imshow(combined_mag)
        plt.title('Magnitude Spectrum (BGR)')
        plt.axis('off')

    if reconstructed.any():
        plt.subplot(1, 3, 3)
        plt.imshow(cv2.cvtColor(reconstructed, cv2.COLOR_BGR2RGB))
        plt.title('Reconstructed Image')
        plt.axis('off')

    plt.show()

def main(image_path, fft_prefix):
    image = load_image(image_path)
    fshift_channels, magnitude_spectrums = split_fft(image)
    save_fft_parts_compressed(fft_prefix, fshift_channels)
    save_magnitude_spectrums_as_images(fft_prefix, magnitude_spectrums)

    # To test loading and reconstructing:
    loaded_fshift_channels = load_fft_parts_compressed(fft_prefix)
    reconstructed_image = merge_fft(loaded_fshift_channels)

    plot_images(image, magnitude_spectrums, reconstructed_image)
    #plot_images(None, None, reconstructed_image)

if __name__ == "__main__":
    input_image_path = 'your_image.jpg'        # Replace with your image path
    fft_parts_prefix = 'fft_parts/your_image'  # Directory + prefix for saving FFT parts
    os.makedirs(os.path.dirname(fft_parts_prefix), exist_ok=True)
    main(input_image_path, fft_parts_prefix)