# -*- coding: utf-8 -*-
"""
该脚本用于对数据集中的文本提示（prompt）进行预编码。

它会执行以下步骤：
1. 使用 CustomDataset 加载定义在 JSON 文件中的数据集。
2. 在分布式环境中（支持多GPU）并行处理数据。
3. 使用预训练的 FLUX 模型对每个文本提示进行编码，生成 prompt embeddings。
4. 将编码后的 embeddings 保存为单独的 .pt 文件。
5. 生成一个最终的 JSON 元数据文件，将原始提示映射到其对应的 embedding 文件路径。

如何运行:
torchrun --nproc_per_node=NUM_GPUS your_script_name.py \
    --json_file /path/to/your/dataset_v0.1.json \
    --output_dir /path/to/save/embeddings \
    --model_path black-forest-labs/FLUX.1-Kontext-dev
"""
import os
import json
import random
import argparse
from tqdm import tqdm

import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

from PIL import Image
import torchvision.transforms as T

from diffusers import FluxKontextPipeline
from customization_kontext.dataset import CustomDataset

def encode_images(pipeline: FluxKontextPipeline, images: torch.Tensor):
    """
    Encodes the images into tokens and ids for FLUX pipeline.
    """
    images = pipeline.image_processor.preprocess(images)
    images = images.to(pipeline.device).to(pipeline.dtype)
    images = pipeline.vae.encode(images).latent_dist.sample()
    images = (
        images - pipeline.vae.config.shift_factor
    ) * pipeline.vae.config.scaling_factor
    images_tokens = pipeline._pack_latents(images, *images.shape)
    images_ids = pipeline._prepare_latent_image_ids(
        images.shape[0],
        images.shape[2],
        images.shape[3],
        pipeline.device,
        pipeline.dtype,
    )
    if images_tokens.shape[1] != images_ids.shape[0]:
        images_ids = pipeline._prepare_latent_image_ids(
            images.shape[0],
            images.shape[2] // 2,
            images.shape[3] // 2,
            pipeline.device,
            pipeline.dtype,
        )

    images_ids[..., 0] = 1 

    return images_tokens, images_ids

def main(args):
    # 1. 初始化分布式环境
    local_rank = int(os.getenv("RANK", 0))
    world_size = int(os.getenv("WORLD_SIZE", 1))
    print(f"World size: {world_size}, Local rank: {local_rank}")

    device = torch.device("cuda", local_rank)
    torch.cuda.set_device(local_rank)
    if not dist.is_initialized():
        dist.init_process_group(
            backend="nccl", init_method="env://", world_size=world_size, rank=local_rank
        )

    # 2. 创建输出目录 (仅在主进程中执行)
    if local_rank == 0:
        os.makedirs(args.output_dir, exist_ok=True)
        os.makedirs(os.path.join(args.output_dir, "prompt_embeds"), exist_ok=True)
        os.makedirs(os.path.join(args.output_dir, "text_ids"), exist_ok=True)
        os.makedirs(os.path.join(args.output_dir, "pooled_prompt_embeds"), exist_ok=True)
        os.makedirs(os.path.join(args.output_dir, "ref_image_tokens"), exist_ok=True)
        os.makedirs(os.path.join(args.output_dir, "ref_image_ids"), exist_ok=True)
    dist.barrier() # 等待主进程创建好目录

    # 3. 准备数据集和 DataLoader
    train_dataset = CustomDataset(json_file=args.json_file)
    sampler = DistributedSampler(
        train_dataset, rank=local_rank, num_replicas=world_size, shuffle=False # 编码时无需打乱
    )
    train_dataloader = DataLoader(
        train_dataset,
        sampler=sampler,
        batch_size=args.train_batch_size,
        num_workers=args.dataloader_num_workers,
    )
    
    # 4. 加载模型
    pipe = FluxKontextPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16).to(device)

    # 5. 循环处理数据并保存
    json_data_local = []
    # 为 tqdm 添加进程判断，只在主进程显示进度条
    pbar = tqdm(train_dataloader, disable=(local_rank != 0))
    
    for data_batch in pbar:
        try:
            with torch.inference_mode():
                # 从批次中获取 prompts 和 image_paths
                prompts = data_batch["prompt"]
                image_paths = data_batch["image_path"]
                ref_image_paths = data_batch["ref_image_path"]
                subjects = data_batch["subject"]

                # 使用 FLUX pipeline 编码文本
                # 注意：FLUX 的 encode_prompt 通常需要 prompt 和 prompt_2
                prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt(
                    prompt=prompts, prompt_2=prompts
                )

                # 逐个保存批次中的每个样本
                for i in range(len(prompts)):
                    # 从图像路径生成唯一的文件名 (例如: 'my_image.png' -> 'my_image')
                    base_name = os.path.basename(image_paths[i])
                    file_stem = os.path.splitext(base_name)[0]

                    # 定义各个 embedding 的保存路径
                    prompt_embed_path = os.path.join("prompt_embeds", f"{file_stem}.pt")
                    pooled_prompt_embeds_path = os.path.join("pooled_prompt_embeds", f"{file_stem}.pt")
                    text_ids_path = os.path.join("text_ids", f"{file_stem}.pt")

                    # 保存张量
                    torch.save(prompt_embeds[i].cpu(), os.path.join(args.output_dir, prompt_embed_path))
                    torch.save(pooled_prompt_embeds[i].cpu(), os.path.join(args.output_dir, pooled_prompt_embeds_path))
                    torch.save(text_ids[i].cpu(), os.path.join(args.output_dir, text_ids_path))
                    
                    # 准备写入 JSON 的元数据
                    item_metadata = {
                        "prompt_embed_path": prompt_embed_path,
                        "pooled_prompt_embeds_path": pooled_prompt_embeds_path,
                        "text_ids_path": text_ids_path,
                        "prompt": prompts[i],
                        # 保留原始图像路径作为参考
                        "original_image_path": image_paths[i],
                        "ref_image_path": ref_image_paths[i],
                        "subject": subjects[i]
                    }
                    json_data_local.append(item_metadata)

        except Exception as e:
            print(f"Rank {local_rank} encountered an error: {repr(e)}")
            dist.barrier() # 发生错误时，同步所有进程并终止
            raise e
            
    # 6. 聚合所有进程的数据并写入最终的 JSON 文件
    dist.barrier() # 确保所有进程都已完成编码
    
    gathered_data = [None] * world_size
    dist.all_gather_object(gathered_data, json_data_local)

    if local_rank == 0:
        print("Gathering results from all processes...")
        # 将收集到的数据列表扁平化
        all_json_data = [item for sublist in gathered_data for item in sublist]
        
        output_json_path = os.path.join(args.output_dir, "preprocessd_data.json")
        with open(output_json_path, "w", encoding='utf-8') as f:
            json.dump(all_json_data, f, indent=4, ensure_ascii=False)
        print(f"Successfully saved metadata for {len(all_json_data)} items to {output_json_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Pre-encode prompts using FLUX model.")
    parser.add_argument(
        "--json_file", 
        type=str, 
        required=True,
        help="Path to the input dataset JSON file."
    )
    parser.add_argument(
        "--output_dir", 
        type=str, 
        required=True,
        help="Directory to save the output embeddings and metadata JSON."
    )
    parser.add_argument(
        "--model_path", 
        type=str, 
        default="path/to/FLUX.1-Kontext-dev",
        help="Path or Hugging Face ID of the pre-trained FLUX model."
    )
    parser.add_argument(
        "--train_batch_size", 
        type=int, 
        default=8,
        help="Batch size per GPU for processing."
    )
    parser.add_argument(
        "--dataloader_num_workers", 
        type=int, 
        default=4,
        help="Number of worker processes for data loading."
    )

    args = parser.parse_args()
    main(args)