# 标准库
import argparse
import math
import os
import json
import time
from collections import deque
import re, shutil
import pdb

# 第三方库
import torch
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from PIL import Image

# Diffusers 相关
from accelerate import Accelerator
from accelerate.utils import set_seed, ProjectConfiguration
from diffusers import FluxTransformer2DModel, AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
from diffusers.optimization import get_scheduler
from diffusers.utils import check_min_version

# 本地模块
from customization_kontext.dataset import LatentDataset
from customization_kontext.imagereward_scorer import ImageRewardScorer
from customization_kontext.vlm import VlmPairwiseScorer

# 版本检查（保持原位）
check_min_version("0.31.0")

def repeat_tensor(tensor):
    if tensor is None:
        return None
    return torch.repeat_interleave(tensor, args.num_generations, dim=0)

# No changes needed for these functions
def sd3_time_shift(shift, t):
    return (shift * t) / (1 + (shift - 1) * t)

def flux_step(
    model_output: torch.Tensor,
    latents: torch.Tensor,
    eta: float,
    sigmas: torch.Tensor,
    index: int,
    prev_sample: torch.Tensor,
    grpo: bool,
    sde_solver: bool,
):
    sigma = sigmas[index]
    dsigma = sigmas[index + 1] - sigma
    prev_sample_mean = latents + dsigma * model_output

    pred_original_sample = latents - sigma * model_output

    delta_t = sigma - sigmas[index + 1]
    std_dev_t = eta * math.sqrt(delta_t)

    if sde_solver:
        score_estimate = -(latents-pred_original_sample*(1 - sigma))/sigma**2
        log_term = -0.5 * eta**2 * score_estimate
        prev_sample_mean = prev_sample_mean + log_term * dsigma

    if grpo and prev_sample is None:
        prev_sample = prev_sample_mean + torch.randn_like(prev_sample_mean) * std_dev_t 
        

    if grpo:
        # log prob of prev_sample given prev_sample_mean and std_dev_t
        log_prob = (
            -((prev_sample.detach().to(torch.float32) - prev_sample_mean.to(torch.float32)) ** 2) / (2 * (std_dev_t**2))
        )
        - math.log(std_dev_t)- torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
        # mean along all but batch dimension
        log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
        return prev_sample, pred_original_sample, log_prob
    else:
        return prev_sample_mean,pred_original_sample

def encode_images(vae, images, device, dtype):
    """
    Encodes the images into tokens and ids for FLUX pipeline.
    """
    image_processor = VaeImageProcessor(16)
    images = image_processor.preprocess(images)
    images = images.to(device).to(dtype)
    images = vae.encode(images).latent_dist.sample()
    images = (
        images - vae.config.shift_factor
    ) * vae.config.scaling_factor

    images_tokens = pack_latents(images, *images.shape)
    images_ids = prepare_latent_image_ids(
        images.shape[0],
        images.shape[2],
        images.shape[3],
        device,
        dtype,
    )
    if images_tokens.shape[1] != images_ids.shape[0]:
        images_ids = prepare_latent_image_ids(
            images.shape[0],
            images.shape[2] // 2,
            images.shape[3] // 2,
            device,
            dtype,
        )

    images_ids[..., 0] = 1 

    return images_tokens, images_ids

def prepare_latent_image_ids(batch_size, height, width, device, dtype):
    latent_image_ids = torch.zeros(height, width, 3)
    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]

    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape

    latent_image_ids = latent_image_ids.reshape(
        latent_image_id_height * latent_image_id_width, latent_image_id_channels
    )

    return latent_image_ids.to(device=device, dtype=dtype)

def pack_latents(latents, batch_size, num_channels_latents, height, width):
    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
    latents = latents.permute(0, 2, 4, 1, 3, 5)
    latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)

    return latents

def unpack_latents(latents, height, width, vae_scale_factor):
    batch_size, num_patches, channels = latents.shape

    # VAE applies 8x compression on images but we must also account for packing which requires
    # latent height and width to be divisible by 2.
    height = 2 * (int(height) // (vae_scale_factor * 2))
    width = 2 * (int(width) // (vae_scale_factor * 2))

    latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
    latents = latents.permute(0, 3, 1, 4, 2, 5)

    latents = latents.reshape(batch_size, channels // (2 * 2), height, width)

    return latents

def run_sample_step(
        args,
        accelerator,
        z,
        progress_bar,
        sigma_schedule,
        transformer,
        encoder_hidden_states, 
        pooled_prompt_embeds, 
        text_ids,
        image_ids,
        ref_image_tokens,
        ref_image_ids,
        grpo_sample,
    ):
    if grpo_sample:
        all_latents = [z]
        all_log_probs = []
        transformer.eval()

        for i in progress_bar:
            B = encoder_hidden_states.shape[0]
            sigma = sigma_schedule[i]
            timestep_value = int(sigma * 1000)
            timesteps = torch.full([B], timestep_value, device=accelerator.device, dtype=torch.long)


            # 拼接参考图
            if ref_image_tokens is not None:
                latent_input = torch.cat([z, ref_image_tokens], dim=1)

            if image_ids is not None:
                latent_ids_input = torch.cat([image_ids, ref_image_ids], dim=0)  # dim 0 is sequence dimension

            with torch.autocast("cuda", torch.bfloat16):
                noise_pred = transformer(
                    hidden_states=latent_input, #修改输入
                    encoder_hidden_states=encoder_hidden_states,
                    timestep=timesteps/1000,
                    guidance=torch.tensor([3.5], device=accelerator.device),
                    txt_ids=text_ids.repeat(encoder_hidden_states.shape[1],1),
                    pooled_projections=pooled_prompt_embeds,
                    img_ids=latent_ids_input,
                    joint_attention_kwargs=None,
                    return_dict=False,
                )[0]
                pred = noise_pred[:, :z.size(1)]
            
            z, pred_original, log_prob = flux_step(pred, z.to(torch.float32), args.eta, sigmas=sigma_schedule, index=i, prev_sample=None, grpo=True, sde_solver=True)
            all_latents.append(z)
            all_log_probs.append(log_prob)
        
        latents = pred_original
        all_latents = torch.stack(all_latents, dim=1)
        all_log_probs = torch.stack(all_log_probs, dim=1)
        return z, latents, all_latents, all_log_probs

def grpo_one_step(
            args,
            accelerator,
            latents,
            pre_latents,
            encoder_hidden_states, 
            pooled_prompt_embeds, 
            text_ids,
            image_ids,
            ref_image_tokens,
            ref_image_ids,
            transformer,
            timesteps,
            i,
            sigma_schedule,
):
    device = accelerator.device
    B = encoder_hidden_states.shape[0]
    transformer.train() # Ensure model is in train mode
    if ref_image_tokens is not None:
        latent_input = torch.cat([latents, ref_image_tokens], dim=1)

    if image_ids is not None:
        latent_ids_input = torch.cat([image_ids.squeeze(0), ref_image_ids], dim=0)  # dim 0 is sequence dimension

    with torch.autocast("cuda", torch.bfloat16):
        noise_pred = transformer(
            hidden_states=latent_input.to(device),
            encoder_hidden_states=encoder_hidden_states.to(device),
            timestep=timesteps/1000,
            guidance=torch.tensor([3.5], device=device),
            txt_ids=text_ids.repeat(encoder_hidden_states.shape[1],1).to(device),
            pooled_projections=pooled_prompt_embeds.to(device),
            img_ids=latent_ids_input.to(device),
            joint_attention_kwargs=None,
            return_dict=False,
        )[0]
        pred = noise_pred[:, :latents.size(1)]
    
    _, _, log_prob = flux_step(pred, latents.to(torch.float32), args.eta, sigma_schedule, i, prev_sample=pre_latents.to(torch.float32), grpo=True, sde_solver=True)
    return log_prob

def sample_reference_model(
    args,
    accelerator,
    transformer,
    vae,
    encoder_hidden_states, 
    pooled_prompt_embeds, 
    text_ids,
    ref_image_tokens,
    ref_image_ids,
    reward_model,
    prompt,
    image_path,
    ref_image_path,
    subject
):
    device = accelerator.device # Use accelerator's device
    w, h, t = args.w, args.h, args.t
    sample_steps = args.sampling_steps
    sigma_schedule = torch.linspace(1, 0, args.sampling_steps + 1, device=device)
    sigma_schedule = sd3_time_shift(args.shift, sigma_schedule)
    B = encoder_hidden_states.shape[0]
    SPATIAL_DOWNSAMPLE = 8
    IN_CHANNELS = 16
    latent_w, latent_h = w // SPATIAL_DOWNSAMPLE, h // SPATIAL_DOWNSAMPLE

    batch_size = 1
    batch_indices = torch.chunk(torch.arange(B), B // batch_size)

    all_latents, all_log_probs, all_rewards, all_image_ids = [], [], [], []

    if args.init_same_noise:
        input_latents = torch.randn((batch_size, IN_CHANNELS, latent_h, latent_w), device=device, dtype=torch.bfloat16)

    images = []
    for index, batch_idx in enumerate(batch_indices):
        batch_encoder_hidden_states = encoder_hidden_states[batch_idx]
        batch_pooled_prompt_embeds = pooled_prompt_embeds[batch_idx]
        batch_text_ids = text_ids[batch_idx]
        batch_ref_image_tokens = ref_image_tokens[batch_idx]
        batch_prompt = [prompt[i] for i in batch_idx]
        if not args.init_same_noise:
            input_latents = torch.randn(
                    (len(batch_idx), IN_CHANNELS, latent_h, latent_w),  #（c,t,h,w)
                    device=device,
                    dtype=torch.bfloat16,
                )
        input_latents_new = pack_latents(input_latents, len(batch_idx), IN_CHANNELS, latent_h, latent_w)
        image_ids = prepare_latent_image_ids(len(batch_idx), latent_h // 2, latent_w // 2, device, torch.bfloat16)
        grpo_sample=True
        progress_bar = tqdm(range(0, sample_steps), desc="Sampling Progress")
        with torch.no_grad(): # Inference context
            z, latents, batch_latents, batch_log_probs = run_sample_step(
                args, accelerator, input_latents_new.to(device),
                progress_bar,
                sigma_schedule,
                transformer,
                batch_encoder_hidden_states.to(device),
                batch_pooled_prompt_embeds.to(device),
                batch_text_ids.to(device),
                image_ids.to(device),
                batch_ref_image_tokens.to(device),
                ref_image_ids.to(device),
                grpo_sample,
            )
        all_image_ids.append(image_ids)
        all_latents.append(batch_latents)
        all_log_probs.append(batch_log_probs)

        unwrapped_vae = accelerator.unwrap_model(vae)
        unwrapped_vae.enable_tiling()
        image_processor = VaeImageProcessor(16)

        with torch.inference_mode():
            with torch.autocast("cuda", dtype=torch.bfloat16):
                latents = unpack_latents(latents, h, w, 8)
                latents = (latents / 0.3611) + 0.1159
                image = unwrapped_vae.decode(latents, return_dict=False)[0]
                decoded_image = image_processor.postprocess(
                image)
        images.append(decoded_image[0])

    if isinstance(reward_model, VlmPairwiseScorer):
        bs = args.train_batch_size
        for i in range(bs):
            print(f"GPU {accelerator.process_index}: Computing VLM rewards for batch {i}")
            iref = Image.open(ref_image_path[i]).convert("RGB")
            iprompt = prompt[i * args.num_generations]
            isubject = subject[i]
            batch_images = images[i*args.num_generations:(i+1)*args.num_generations]
            reward = reward_model(
                ref=iref, 
                images=batch_images, 
                instruction=iprompt,
                subject=isubject
            )
            all_rewards.extend(reward)
    elif isinstance(reward_model, ImageRewardScorer):
        for i, image in enumerate(images):
            reward = reward_model(prompt[i], [image])
            all_rewards.append(reward)
    else:
        raise NotImplementedError("Only VLM and ImageRewardScorer are supported currently.")
    
    rank = accelerator.process_index
    logdir = os.path.join(args.output_dir, "images", f"rank_{rank}")
    os.makedirs(logdir, exist_ok=True)
    for i in range(args.train_batch_size):
        batch_images = images[i*args.num_generations:(i+1)*args.num_generations]
        batch_rewards = all_rewards[i*args.num_generations:(i+1)*args.num_generations]
        batch_logdir = os.path.join(logdir, f"batch_{i}")
        if os.path.exists(batch_logdir):
            shutil.rmtree(batch_logdir)
        os.makedirs(batch_logdir, exist_ok=True)
        for j, (img, reward) in enumerate(zip(batch_images, batch_rewards)):
            img_save_path = os.path.join(
                batch_logdir,
                f"img{j}_reward{reward:.2f}.png",
            )
            img.save(img_save_path)
        shutil.copy(ref_image_path[i], os.path.join(batch_logdir, "ref_image.png"))
        with open(os.path.join(batch_logdir, "prompt.txt"), "w") as f:
            f.write(prompt[i * args.num_generations])

    all_latents = torch.cat(all_latents, dim=0)
    all_log_probs = torch.cat(all_log_probs, dim=0)
    all_rewards = torch.tensor(all_rewards,device=device)
    all_image_ids = torch.stack(all_image_ids, dim=0)
    
    return all_rewards, all_latents, all_log_probs, sigma_schedule, all_image_ids

def train_one_step(
    args,
    accelerator: Accelerator,
    transformer,
    vae,
    reward_model,
    batch,
    optimizer,
    lr_scheduler,
):

    (
        encoder_hidden_states, 
        pooled_prompt_embeds, 
        text_ids,
        ref_image,
        prompt,
        image_path,
        ref_image_path,
        subject,
    ) = batch
    device = accelerator.device # Use accelerator's device
    ref_image_tokens, ref_image_ids = encode_images(vae,ref_image, device, torch.bfloat16)

    if args.use_group:
        encoder_hidden_states = repeat_tensor(encoder_hidden_states)
        pooled_prompt_embeds = repeat_tensor(pooled_prompt_embeds)
        text_ids = repeat_tensor(text_ids)
        ref_image_tokens = repeat_tensor(ref_image_tokens)
        if isinstance(prompt, str):
            prompt = [prompt] * args.num_generations
        elif isinstance(prompt, list):
            prompt = [item for item in prompt for _ in range(args.num_generations)]
        else:
            raise ValueError(f"Unsupported prompt type: {type(prompt)}")

    reward, all_latents, all_log_probs, sigma_schedule, all_image_ids = sample_reference_model(args, accelerator, transformer, vae, encoder_hidden_states, pooled_prompt_embeds, text_ids, ref_image_tokens, ref_image_ids, reward_model, prompt, image_path, ref_image_path, subject)
    
    batch_size = all_latents.shape[0]
    timestep_value = [int(sigma * 1000) for sigma in sigma_schedule][:args.sampling_steps]
    timestep_values = [timestep_value[:] for _ in range(batch_size)]
    device = all_latents.device
    timesteps =  torch.tensor(timestep_values, device=all_latents.device, dtype=torch.long)

    samples = {
        "timesteps": timesteps.detach().clone()[:, :-1],
        "latents": all_latents[:, :-1][:, :-1],
        "next_latents": all_latents[:, 1:][:, :-1],
        "log_probs": all_log_probs[:, :-1],
        "rewards": reward.to(torch.float32),
        "image_ids": all_image_ids,
        "text_ids": text_ids,
        "encoder_hidden_states": encoder_hidden_states,
        "pooled_prompt_embeds": pooled_prompt_embeds,
        "ref_image_tokens": ref_image_tokens,
    }

    gathered_reward = accelerator.gather_for_metrics(samples["rewards"])
    if accelerator.is_main_process:
        mean_reward = gathered_reward.mean().item()
        print("gathered_mean_reward", mean_reward)
        with open(f'{args.output_dir}/reward.txt', 'a') as f: 
            f.write(f"{mean_reward}\n")

    #计算advantage
    if args.use_group:
        n = len(samples["rewards"]) // (args.num_generations)
        advantages = torch.zeros_like(samples["rewards"])
        
        for i in range(n):
            start_idx = i * args.num_generations
            end_idx = (i + 1) * args.num_generations
            group_rewards = samples["rewards"][start_idx:end_idx]
            group_mean = group_rewards.mean()
            group_std = group_rewards.std() + 1e-8
            advantages[start_idx:end_idx] = (group_rewards - group_mean) / group_std
        
        samples["advantages"] = advantages
    else:
        advantages = (samples["rewards"] - gathered_reward.mean())/(gathered_reward.std()+1e-8)
        samples["advantages"] = advantages

    
    perms = torch.stack(
        [
            torch.randperm(len(samples["timesteps"][0]))
            for _ in range(batch_size)
        ]
    ).to(device) 
    for key in ["timesteps", "latents", "next_latents", "log_probs"]:
        samples[key] = samples[key][
            torch.arange(batch_size).to(device) [:, None],
            perms,
        ]
    samples_batched = {
        k: v.unsqueeze(1)
        for k, v in samples.items()
    }
    # dict of lists -> list of dicts for easier iteration
    samples_batched_list = [
        dict(zip(samples_batched, x)) for x in zip(*samples_batched.values())
    ]
    train_timesteps = int(len(samples["timesteps"][0]) * args.timestep_fraction)

    total_loss = 0.0
    
    for i, sample in enumerate(samples_batched_list):
        for _ in range(train_timesteps):
            clip_range = args.clip_range
            adv_clip_max = args.adv_clip_max
            new_log_probs = grpo_one_step(
                args,
                accelerator,
                sample["latents"][:,_],
                sample["next_latents"][:,_],
                sample["encoder_hidden_states"],
                sample["pooled_prompt_embeds"],
                sample["text_ids"],
                sample["image_ids"],
                sample["ref_image_tokens"],
                ref_image_ids,
                transformer,
                sample["timesteps"][:,_],
                perms[i][_],
                sigma_schedule,
            )
            advantages = torch.clamp(
                sample["advantages"],
                -adv_clip_max,
                adv_clip_max,
            )

            ratio = torch.exp(new_log_probs - sample["log_probs"][:,_])

            unclipped_loss = -advantages * ratio
            clipped_loss = -advantages * torch.clamp(
                ratio,
                1.0 - clip_range,
                1.0 + clip_range,
            )

            loss = torch.mean(torch.maximum(unclipped_loss, clipped_loss)) / train_timesteps
            accelerator.backward(loss)
            total_loss += accelerator.gather(loss.detach()).mean().item()

    return total_loss


def main(args):
    if args.reward_model == "vlm":
        reward_model = VlmPairwiseScorer()
    elif args.reward_model == "ir":
        reward_model = ImageRewardScorer(
            dtype=torch.float32,
            model_dir=args.pretrained_model_name_or_path,
        )
        accelerator.print("Using ImageReward Scorer for training")
    else:
        raise NotImplementedError("Only VLM and ImageRewardScorer are supported currently.")

    args.output_dir = os.path.join(args.output_dir, args.run_name)
    project_config = ProjectConfiguration(
        project_dir=args.output_dir, 
        logging_dir=os.path.join(args.output_dir, "logs"),
        automatic_checkpoint_naming=True,
        total_limit=args.max_checkpoints_to_keep,
    )
    
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        log_with="wandb",
        project_config=project_config,
        mixed_precision=args.mixed_precision,
    )
        
    accelerator.init_trackers(
        project_name=args.wandb_project, 
        config=vars(args), # 可选：记录所有超参数
        init_kwargs={"wandb": {"name": args.run_name}} # 传递 run name
    )

    if accelerator.is_main_process:
        os.makedirs(f"{args.output_dir}/images", exist_ok=True)
        # save args to output directory
        with open(os.path.join(args.output_dir, "training_args.json"), "w") as f:
            json.dump(vars(args), f, indent=4)

    accelerator.print(f"Starting training on {accelerator.num_processes} processes with strategy: {accelerator.distributed_type}")
    device = accelerator.device

    if args.seed is not None:
        set_seed(args.seed, device_specific=True)



    accelerator.print(f"--> loading model from {args.pretrained_model_name_or_path}")
    transformer = FluxTransformer2DModel.from_pretrained(
        f"{args.pretrained_model_name_or_path}/FLUX.1-Kontext-dev", subfolder="transformer", torch_dtype=torch.bfloat16
    )
    vae = AutoencoderKL.from_pretrained(
        f"{args.pretrained_model_name_or_path}/FLUX.1-Kontext-dev", subfolder="vae", torch_dtype=torch.bfloat16
    )
    vae.requires_grad_(False).eval()
    vae.to(device)
    
    if args.gradient_checkpointing:
        transformer.enable_gradient_checkpointing()
    transformer.train()
    
    # --- Optimizer and Scheduler ---
    params_to_optimize = list(filter(lambda p: p.requires_grad, transformer.parameters()))
    optimizer = torch.optim.AdamW(
        params_to_optimize, lr=args.learning_rate, betas=(0.9, 0.999), weight_decay=args.weight_decay, eps=1e-8
    )
    lr_scheduler = get_scheduler(
        args.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps
    )

    # --- Dataset and DataLoader ---
    train_dataset = LatentDataset(args.data_json_path)
    train_dataloader = DataLoader(
        train_dataset,
        shuffle=True, # Set shuffle here
        batch_size=args.train_batch_size,
        num_workers=args.dataloader_num_workers,
        pin_memory=True,
        drop_last=True,
    )

    accelerator.print("***** Running training *****")
    accelerator.print(f"  Num examples = {len(train_dataset)}")
    accelerator.print(f"  Dataloader size = {len(train_dataloader)}")
    accelerator.print(f"  Instantaneous batch size per device = {args.train_batch_size}")
    accelerator.print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    accelerator.print(f"  Total optimization steps per epoch = {args.max_train_steps}")
    accelerator.print(f"  Total training parameters  {sum(p.numel() for p in params_to_optimize) / 1e9} B")
    accelerator.print(f"  Master weight dtype: {transformer.parameters().__next__().dtype}")

    # Prepare the main training components TOGETHER. This creates the DeepSpeed engine.
    accelerator.print("Preparing main training components with Accelerator...")
    transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        transformer, optimizer, train_dataloader, lr_scheduler
    )

    # 检查是否正确使用了 DeepSpeed
    if hasattr(transformer, 'module'):
        accelerator.print("✓ DeepSpeed engine initialized successfully")
    else:
        accelerator.print("⚠ Warning: DeepSpeed may not be properly initialized")

    progress_bar = tqdm(range(args.max_train_steps))

    step_times = deque(maxlen=100)
    completed_steps = 0
    
    # --- 主训练循环 ---
    while completed_steps < args.max_train_steps:
        for batch in train_dataloader:
            start_time = time.time()
            with accelerator.accumulate(transformer):
                loss = train_one_step(
                    args, accelerator,transformer, vae, reward_model, batch, optimizer, lr_scheduler
                )

            # --- Logging and Checkpointing ---
            if accelerator.sync_gradients:
                grad_norm = accelerator.clip_grad_norm_(transformer.parameters(), args.max_grad_norm)
                completed_steps += 1

                step_time = time.time() - start_time
                step_times.append(step_time)
                avg_step_time = sum(step_times) / len(step_times)
                progress_bar.set_postfix(
                    {
                        "loss": f"{loss:.4f}",
                        "step_time": f"{step_time:.2f}s",
                    }
                )
                progress_bar.update(1)
                accelerator.log(
                    {
                        "train_loss": loss,
                        "learning_rate": lr_scheduler.get_last_lr()[0],
                        "step_time": step_time,
                        "avg_step_time": avg_step_time,
                        "grad_norm": grad_norm,
                    },
                    step=completed_steps,
                )

                if completed_steps > 0 and completed_steps % args.checkpointing_steps == 0:
                    accelerator.save_state()
                    accelerator.print(f"Saved checkpoint")

                # 5. 检查是否需要提前退出循环
                if completed_steps >= args.max_train_steps:
                    break

    accelerator.end_training()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # dataset & dataloader
    parser.add_argument("--data_json_path", type=str, required=True)
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=10,
        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
    )
    parser.add_argument(
        "--train_batch_size",
        type=int,
        default=16,
        help="Batch size (per device) for the training dataloader.",
    )
    # text encoder & vae & diffusion model
    parser.add_argument("--pretrained_model_name_or_path", type=str)

    # validation & logs
    parser.add_argument(
        "--seed", type=int, default=None, help="A seed for reproducible training."
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default=None,
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--run_name",
        type=str,
        default=None,
        help="The output sub directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument(
        "--wandb_project",
        type=str,
        default=None,
        help="The wandb project name",
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=int,
        default=500,
        help=(
            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
            " training using `--resume_from_checkpoint`."
        ),
    )
    parser.add_argument(
        "--max_checkpoints_to_keep",
        type=int,
        default=10,
        help=(
            "Maximum number of checkpoints to keep. If the number of checkpoints exceeds this value, the oldest checkpoints will be deleted."
        ),
    )
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help=(
            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
        ),
    )
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )

    # optimizer & scheduler & Training
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-4,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--lr_warmup_steps",
        type=int,
        default=10,
        help="Number of steps for the warmup in the lr scheduler.",
    )
    parser.add_argument(
        "--max_grad_norm", default=2.0, type=float, help="Max gradient norm."
    )
    parser.add_argument(
        "--gradient_checkpointing",
        action="store_true",
        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default=None,
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
        ),
    )

    # lr_scheduler
    parser.add_argument(
        "--lr_scheduler",
        type=str,
        default="constant_with_warmup",
        help=(
            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
            ' "constant", "constant_with_warmup"]'
        ),
    )
    parser.add_argument(
        "--weight_decay", type=float, default=0.01, help="Weight decay to apply."
    )

    #GRPO training
    parser.add_argument(
        "--h",
        type=int,
        default=None,   
        help="video height",
    )
    parser.add_argument(
        "--w",
        type=int,
        default=None,   
        help="video width",
    )
    parser.add_argument(
        "--t",
        type=int,
        default=None,   
        help="video length",
    )
    parser.add_argument(
        "--sampling_steps",
        type=int,
        default=None,   
        help="sampling steps",
    )
    parser.add_argument(
        "--eta",
        type=float,
        default=None,   
        help="noise eta",
    )
    parser.add_argument(
        "--sampler_seed",
        type=int,
        default=None,   
        help="seed of sampler",
    )
    parser.add_argument(
        "--use_group",
        action="store_true",
        default=False,
        help="whether compute advantages for each prompt",
    )
    parser.add_argument(
        "--num_generations",
        type=int,
        default=16,   
        help="num_generations per prompt",
    )
    parser.add_argument(
        "--init_same_noise",
        action="store_true",
        default=False,
        help="whether use the same noise within each prompt",
    )
    parser.add_argument(
        "--shift",
        type = float,
        default=1.0,
        help="shift for timestep scheduler",
    )
    parser.add_argument(
        "--timestep_fraction",
        type = float,
        default=1.0,
        help="timestep downsample ratio",
    )
    parser.add_argument(
        "--clip_range",
        type = float,
        default=1e-4,
        help="clip range for grpo",
    )
    parser.add_argument(
        "--adv_clip_max",
        type = float,
        default=5.0,
        help="clipping advantage",
    )
    parser.add_argument(
        "--reward_model",
        type = str,
        choices=["ir", "vlm"],
        required=True,
    )

    args = parser.parse_args()
    main(args)