from datasets import load_from_disk
import numpy as np
import torch
from PIL import Image
from collections import defaultdict
from diffusers import AutoencoderKLQwenImage
import gc

def analyze_dataset_by_size(dataset_path):
    """
    Группирует датасет по размерам изображений и выводит базовую информацию.
    """
    # Настройка устройства и типа данных
    dtype = torch.float32
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Загрузка VAE модели
    print("Загрузка VAE модели...")
    vae = AutoencoderKLQwenImage.from_pretrained("vae",torch_dtype=dtype).to(device).eval()
    shift_factor = getattr(vae.config, "shift_factor", 0.0)
    if shift_factor is None:
        shift_factor = 0.0
    
    scaling_factor = getattr(vae.config, "scaling_factor", 1.0)
    if scaling_factor is None:
        scaling_factor = 1.0
        
    mean = getattr(vae.config, "latents_mean", None)
    std = getattr(vae.config, "latents_std", None)
    if mean is not None and std is not None:
        latents_std = torch.tensor(std, device=device, dtype=dtype).view(1, len(std), 1, 1)
        latents_mean = torch.tensor(mean, device=device, dtype=dtype).view(1, len(mean), 1, 1)
    
    # Загружаем датасет
    print(f"Загрузка датасета из {dataset_path}...")
    dataset = load_from_disk(dataset_path)

    print(f"Осталось примеров после фильтрации: {len(dataset)}")
    
    # Группируем примеры по размерам
    print("\nГруппировка примеров по размерам...")
    size_to_indices = defaultdict(list)
    
    # Собираем примеры с одинаковыми размерами
    # Собираем примеры с одинаковыми размерами (оптимизированная версия)
    widths = dataset["width"]
    heights = dataset["height"]
    for i, (w, h) in enumerate(zip(widths, heights)):
        size_to_indices[(w, h)].append(i)
    
    # Сортируем размеры по количеству примеров
    print("\nСортируем...")
    size_stats = [(size, len(indices)) for size, indices in size_to_indices.items()]
    size_stats.sort(key=lambda x: x[1], reverse=True)
    
    # Выводим информацию о каждой группе и показываем первый пример
    for size, count in size_stats:
        width, height = size
        first_idx = size_to_indices[size][1]
        example = dataset[first_idx]
        
        print(f"\n--- Батч {width}x{height}: {count} примеров ---")
        
        # Декодируем латентное представление для первого примера
        latent = torch.tensor(example["vae"], dtype=dtype).unsqueeze(0).to(device)
        
        # 1. Снова обманываем VAE, превращая картинку в "видео из 1 кадра" [B, C, 1, H, W]
        if latent.ndim == 4:
            latent = latent.unsqueeze(2)
        
        with torch.no_grad():
            if latents_mean is not None and latents_std is not None:
                latent = latent * latents_std + latents_mean
            
            print(f"Min of latent_for_vae: {latent.min()}")
            print(f"Max of latent_for_vae: {latent.max()}")
            print(f"Mean of latent_for_vae: {latent.mean()}")
            print(f"Std: {latent.std().item():.4f}")
            if torch.isnan(latent).any() or torch.isinf(latent).any():
                print("WARNING: Raw latents contain NaN or Inf values!")
            
            reconstructed_image = vae.decode(latent).sample
        
        # 2. Вытаскиваем обычную 3D-картинку [C, H, W] из 5D-видеотензора
        if reconstructed_image.ndim == 5:
            # Берем нулевой батч, все каналы, нулевой кадр, всю высоту и ширину
            img_tensor = reconstructed_image[0, :, 0, :, :] 
        else:
            img_tensor = reconstructed_image.squeeze(0) # На всякий случай, если VAE вернул 4D
        
        img_array = img_tensor.cpu().numpy()
        img_array = np.transpose(img_array, (1, 2, 0))
        img_array = (img_array + 1) / 2  # Нормализация к [0, 1]
        img_array = np.clip(img_array * 255, 0, 255).astype(np.uint8)  # Преобразуем в uint8 для PIL
        
        # Создаем PIL изображение из массива
        pil_image = Image.fromarray(img_array)
        print(f"Текст: {example['text']}")
        print(f"Ключи: {', '.join(example.keys())}")
        print(f"latent: {latent.shape}")
        pil_image.save("1.jpg")
    
    # Очистка памяти
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()
    
    return size_to_indices  # Возвращаем словарь с индексами по группам

# Использование
if __name__ == "__main__":
    # Путь к датасету
    save_path = "datasets/ds234_640_vae_qwen"
    
    # Анализ датасета
    size_groups = analyze_dataset_by_size(save_path)