import os
import numpy as np
from PIL import Image
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import logging
import time
import argparse

# Configure logging
logging.basicConfig(
    filename='process.log',
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def log_and_print(message, level='info'):
    """Log and print messages."""
    if level == 'info':
        logging.info(message)
    elif level == 'debug':
        logging.debug(message)
    elif level == 'error':
        logging.error(message)
    print(message)

def resize_image(image_path, new_size=(256, 256)):
    """Resize the image to a smaller size."""
    try:
        img = Image.open(image_path)
        # Handle PIL versions for resampling
        try:
            resample_method = Image.Resampling.LANCZOS
        except AttributeError:
            resample_method = Image.LANCZOS  # For older PIL versions
        img = img.resize(new_size, resample=resample_method)
        resized_dir = 'resized_images'
        os.makedirs(resized_dir, exist_ok=True)
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        resized_path = os.path.join(resized_dir, f"{base_name}{new_size[0]}x{new_size[1]}.png")
        img.save(resized_path)
        log_and_print(f"Image resized and saved to {resized_path}")
        return resized_path
    except Exception as e:
        log_and_print(f"Error resizing image: {e}", 'error')
        return None

def process_image(image_path, text_message):
    """Process image with t-SNE and save transformed image."""
    resized_image_path = resize_image(image_path)
    if not resized_image_path:
        return False

    try:
        log_and_print(f"Processing image: {resized_image_path}")
        
        # Read image
        img = Image.open(resized_image_path).convert('RGB')
        img_array = np.array(img)
        
        # Flatten image
        img_flattened = img_array.reshape(-1, 3)
        
        # Optional: Dimensionality reduction with PCA before t-SNE for efficiency
        # from sklearn.decomposition import PCA
        # pca = PCA(n_components=50, random_state=0)
        # img_pca = pca.fit_transform(img_flattened)
        # log_and_print("PCA reduced dimensionality before t-SNE.")
        
        # Apply t-SNE
        log_and_print("Applying t-SNE...")
        start_time = time.time()
        try:
            tsne = TSNE(n_components=2, random_state=0, verbose=1, n_jobs=-1, perplexity=30)
            img_tsne = tsne.fit_transform(img_flattened)
            end_time = time.time()
            log_and_print(f"t-SNE computation time: {end_time - start_time:.2f} seconds")
        except MemoryError:
            log_and_print("MemoryError: t-SNE failed due to insufficient memory.", 'error')
            log_and_print("Attempting to process with a smaller image.", 'error')
            return handle_t_sne_failure(image_path, text_message)
        except Exception as e:
            log_and_print(f"t-SNE failed: {e}", 'error')
            return False

        # Normalize t-SNE output for visualization
        tsne_min = np.min(img_tsne, axis=0)
        tsne_max = np.max(img_tsne, axis=0)
        img_tsne_normalized = (img_tsne - tsne_min) / (tsne_max - tsne_min)
        img_tsne_normalized *= 255
        img_tsne_normalized = img_tsne_normalized.astype(np.uint8)

        # Since t-SNE has 2 components, we can map them to RGB by adding a third component
        # For simplicity, we can set the third component to 128 or use another feature
        third_channel = np.full((img_tsne_normalized.shape[0], 1), 128, dtype=np.uint8)
        img_tsne_rgb = np.hstack((img_tsne_normalized, third_channel))

        # Reshape to image dimensions
        height, width = img_array.shape[0], img_array.shape[1]
        img_tsne_reshaped = img_tsne_rgb.reshape(height, width, 3)

        # Create a plot
        plt.figure(figsize=(8, 8))
        plt.imshow(img_tsne_reshaped)
        plt.axis('off')
        plt.title(text_message if text_message else 't-SNE Transformed Image')
        transformed_dir = 'transformed_images'
        os.makedirs(transformed_dir, exist_ok=True)
        transformed_path = os.path.join(transformed_dir, f"transformed{os.path.basename(resized_image_path)}")
        plt.savefig(transformed_path, bbox_inches='tight', pad_inches=0)
        plt.close()
        
        log_and_print(f"Image transformation completed successfully. Saved to {transformed_path}")
    except Exception as e:
        log_and_print(f"Error processing image: {e}", 'error')
        return False
    return True

def handle_t_sne_failure(image_path, text_message):
    """Handle t-SNE failure by resizing the image or skipping."""
    log_and_print("Handling t-SNE failure.")
    # For automation, we can attempt to resize again or skip
    # Here, we'll attempt to resize to a smaller size
    smaller_size = (128, 128)
    log_and_print(f"Attempting to resize image to smaller size: {smaller_size}")
    resized_path = resize_image(image_path, new_size=smaller_size)
    if resized_path:
        return process_image(resized_path, text_message)
    else:
        log_and_print("Failed to resize image for t-SNE.", 'error')
        return False

def read_text_from_file(file_path):
    """Read text message from a file."""
    try:
        with open(file_path, 'r') as file:
            text_message = file.read().strip()
        return text_message
    except Exception as e:
        log_and_print(f"Error reading text file: {e}", 'error')
        return ''

def main():
    # Argument parsing for flexibility
    parser = argparse.ArgumentParser(description='Process an image with t-SNE transformation.')
    parser.add_argument('image_path', type=str, help='Path to the input image.')
    args = parser.parse_args()

    image_path = args.image_path
    text_file_path = '/root/it/text_message.txt'
    text_message = read_text_from_file(text_file_path)

    if not os.path.isfile(image_path):
        log_and_print(f"Image file does not exist: {image_path}", 'error')
        return

    log_and_print(f"Starting process with image: {image_path}")
    log_and_print(f"Text message to display on the plot: {text_message}")

    success = process_image(image_path, text_message)
    if success:
        log_and_print("Process completed successfully.")
    else:
        log_and_print("Process failed.", 'error')

if __name__ == "__main__":
    main()
[t-SNE] Mean sigma: 0.000000