import numpy as np
import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os
import csv

# 1. Function to apply Fast Fourier Transform to a colored image (separate for each channel)
def apply_fft(image):
    """Apply Fast Fourier Transform to a colored image (separate for each channel)"""
    fft_channels = []
    for channel in cv2.split(image):  # Split the image into its color channels (B, G, R)
        fft = np.fft.fft2(channel)
        fft_shifted = np.fft.fftshift(fft)  # Shift the zero frequency to the center
        fft_channels.append(fft_shifted)
    return fft_channels

# 2. Function to display frequency components of each color channel in interactive 3D plots using Matplotlib
def show_frequency_components_3d(fft_channels, title='Frequency Components'):
    """Display the magnitude and phase of the FFT for each channel (R, G, B) in 3D with Matplotlib"""
    
    channel_names = ['Blue Channel', 'Green Channel', 'Red Channel']
    
    # Create a figure for 3D plotting
    fig = plt.figure(figsize=(18, 6))

    # Loop through each channel's FFT data
    for i, fft_data in enumerate(fft_channels):
        magnitude = np.abs(fft_data)  # Magnitude spectrum
        phase = np.angle(fft_data)    # Phase spectrum

        # Generate grid for the 3D plot (x, y grid of frequencies)
        rows, cols = magnitude.shape
        x = np.linspace(-cols // 2, cols // 2, cols)
        y = np.linspace(-rows // 2, rows // 2, rows)
        X, Y = np.meshgrid(x, y)

        # Create a subplot for each channel's magnitude and phase
        ax = fig.add_subplot(1, 6, 2 * i + 1, projection='3d')
        ax.set_title(f'{channel_names[i]} - Magnitude')
        ax.plot_surface(X, Y, magnitude, cmap='viridis', edgecolor='none')

        ax = fig.add_subplot(1, 6, 2 * i + 2, projection='3d')
        ax.set_title(f'{channel_names[i]} - Phase')
        ax.plot_surface(X, Y, phase, cmap='inferno', edgecolor='none')

    plt.suptitle(title, fontsize=16)
    plt.tight_layout()
    plt.show()

# 3. Function to apply percentage-based filtering to the FFT (for each channel)
def filter_fft_percentage(fft_channels, percentage):
    """Apply percentage-based filtering, keeping the highest-magnitude frequency components"""
    filtered_fft = []
    for fft_data in fft_channels:
        magnitude = np.abs(fft_data)
        flat_mag = magnitude.flatten()
        sorted_mag = np.sort(flat_mag)[::-1]  # Sort magnitudes in descending order
        
        # Determine the threshold for the given percentage
        num_elements_to_keep = int(len(sorted_mag) * percentage / 100)
        threshold = sorted_mag[num_elements_to_keep - 1] if num_elements_to_keep > 0 else 0
        
        # Create a mask to keep only the top frequencies
        mask = magnitude >= threshold
        filtered_fft.append(fft_data * mask)  # Apply mask to the FFT data
    return filtered_fft

# 4. Function to apply inverse Fourier transform to reconstruct the color image
def inverse_fft(filtered_fft):
    """Apply inverse Fourier transform to reconstruct the image from filtered FFT"""
    reconstructed_channels = []
    for fft_data in filtered_fft:
        fft_ishift = np.fft.ifftshift(fft_data)  # Reverse FFT shift
        img_reconstructed = np.fft.ifft2(fft_ishift)  # Apply inverse FFT
        img_reconstructed = np.abs(img_reconstructed)  # Get the magnitude of the result
        
        # Normalize and convert to uint8 for image format
        img_normalized = cv2.normalize(img_reconstructed, None, 0, 255, cv2.NORM_MINMAX)
        reconstructed_channels.append(img_normalized.astype(np.uint8))

    # Merge the channels back into a color image (BGR format)
    return cv2.merge(reconstructed_channels)

# 5. Main function to process images in batches
def process_images():
    """Main function to process images, apply FFT, and filter results"""
    # Create directories if they don't exist
    os.makedirs('Modified', exist_ok=True)
    
    # Ask user for the percentage of top frequencies to keep
    percentage = float(input("Enter the percentage of highest-magnitude frequencies to keep (0-100): "))
    print(f"Applying percentage-based filtering: Keeping top {percentage}% of frequencies.")

    # Create CSV logging file
    with open('fft_features.csv', 'w', newline='') as csvfile:
        csv_writer = csv.writer(csvfile)
        csv_writer.writerow(['Image', 'Max Magnitude', 'Mean Magnitude', 'Non-zero Count'])
        
        # Process each image in the 'original' folder
        image_filenames = [filename for filename in os.listdir('original') 
                           if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff'))]
        
        # Process all images
        for i, filename in enumerate(image_filenames):
            # Read image (colored)
            img_path = os.path.join('original', filename)
            img = cv2.imread(img_path)
            
            # Apply FFT to each channel (RGB)
            fft_channels = apply_fft(img)
            
            # Apply percentage-based filtering
            filtered_fft = filter_fft_percentage(fft_channels, percentage)
            
            # Reconstruct the image using inverse FFT
            reconstructed = inverse_fft(filtered_fft)

            # Show frequency components (filtered FFT) as interactive 3D plots using Matplotlib
            show_frequency_components_3d(filtered_fft, f'Filtered FFT - {filename}')
            
            # Log FFT features for the first channel (R)
            magnitude = np.abs(filtered_fft[0])  # Just checking the first channel's magnitude
            non_zero = magnitude > 0
            csv_writer.writerow([ 
                filename,
                np.max(magnitude),
                np.mean(magnitude[non_zero]) if np.any(non_zero) else 0,
                np.count_nonzero(non_zero)
            ])

            # Save reconstructed image
            cv2.imwrite(os.path.join('Modified', filename), reconstructed)

        print("Processing completed successfully!")

# Run the main function
if __name__ == "__main__":
    process_images()