import os
import glob
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image
from umap import UMAP

sns.set_theme(style="whitegrid")

def load_ood_renders(method_dir):
    ood_dir = os.path.join(method_dir, "renders_ood")
    if not os.path.exists(ood_dir):
        return None
    images = sorted(glob.glob(os.path.join(ood_dir, "*.png")))
    if not images:
        return None
    
    feature_vector = []
    for img_path in images:
        img = Image.open(img_path).convert("RGB").resize((64, 64))
        feature_vector.append(np.array(img).flatten())
    
    return np.concatenate(feature_vector)

def main():
    outputs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "outputs"))
    plots_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "plots"))
    os.makedirs(plots_dir, exist_ok=True)
    
    features = []
    labels = []
    
    for folder in os.listdir(outputs_dir):
        method_dir = os.path.join(outputs_dir, folder)
        if not os.path.isdir(method_dir) or folder == "pathology_archive":
            continue
        
        vec = load_ood_renders(method_dir)
        if vec is not None:
            features.append(vec)
            labels.append(folder.split('_')[0])
            
    if not features:
        return
        
    X = np.stack(features)
    reducer = UMAP(n_components=2, random_state=42)
    embedding = reducer.fit_transform(X)
    
    df = pd.DataFrame({
        'UMAP_1': embedding[:, 0],
        'UMAP_2': embedding[:, 1],
        'Method': labels
    })
    
    plt.figure(figsize=(10, 8))
    sns.scatterplot(data=df, x='UMAP_1', y='UMAP_2', hue='Method', s=200, palette='tab10', edgecolor='k')
    plt.title("Functional Space Convergence (OOD Renders)")
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, "fig3_functional_manifold.png"))
    plt.close()

if __name__ == "__main__":
    main()