import pandas as pd
import os
import matplotlib.pyplot as plt
from PIL import Image

# 1. CONFIGURATION (Relative paths for portability)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# Strictly use the root results file generated by graficar_tsne_msi.py
CSV_PATH = os.path.join(BASE_DIR, "..", "msi_clusters_results.csv")

IMG_DIR = os.path.join(BASE_DIR, "..", "datos_para_ai")
OUTPUT_DIR = os.path.join(BASE_DIR, "..", "figures", "galleries")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 2. PROCESSING
if os.path.exists(CSV_PATH):
    print(f"Loading results from: {CSV_PATH}")
    df = pd.read_csv(CSV_PATH)

    # Detect the correct column for ion image names
    possible_cols = ['ion_image', 'ion', 'filename']
    img_col = next((c for c in possible_cols if c in df.columns), None)

    if not img_col:
        print(f"Error: Could not find image column in {CSV_PATH}. Available: {df.columns.tolist()}")
    else:
        for cluster_id in sorted(df['cluster'].unique()):
            print(f"Generating gallery for Cluster {cluster_id}...")
            # Take a 3x3 sample of the cluster
            cluster_ions = df[df['cluster'] == cluster_id][img_col].head(9).tolist()
            
            fig, axes = plt.subplots(3, 3, figsize=(10, 10))
            fig.suptitle(f"Cluster {cluster_id} - Ion Morphologies", fontsize=16)
            
            for i, ax in enumerate(axes.flat):
                if i < len(cluster_ions):
                    img_path = os.path.join(IMG_DIR, cluster_ions[i])
                    if os.path.exists(img_path):
                        img = Image.open(img_path)
                        # Using 'viridis' improves human visibility of low-concentration transport traces in the veins
                        ax.imshow(img, cmap='viridis')
                        ax.set_title(cluster_ions[i], fontsize=8)
                    else:
                        ax.text(0.5, 0.5, 'Image Not Found', ha='center', va='center', fontsize=6)
                ax.axis('off')
            
            plt.tight_layout()
            # Standardized filename for consistency
            plt.savefig(f"{OUTPUT_DIR}/Gallery_cluster_{cluster_id}.png")
            plt.close()

        print(f"\nGalleries successfully saved in: {OUTPUT_DIR}")
else:
    print(f"Error: Missing required results file: {CSV_PATH}")
    print("Please run 'python scripts_python/graficar_tsne_msi.py' first to generate results.")