#!/usr/bin/env python3
"""Diagnostic for ZeRO-2 with unequal real ChartQA image-token lengths.

This is a bounded regression diagnostic for the OPD-only training failure where
one rank reached a Qwen MLP gradient all-reduce while a longer-image rank had
not yet reached the same autograd hook.  It loads the 0.5B student only, picks
the highest-resolution hard-correct ChartQA rows, and performs a few real VLM
backward/optimizer steps with the bounded production ZeRO-2 JSON.  Formal
OPD-only training uses native DDP; this script is retained only to inspect
DeepSpeed behavior when explicitly requested.
"""

from __future__ import annotations

import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import deepspeed
import torch
import torch.distributed as dist
from PIL import Image
from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
from transformers.models.llava_onevision.modeling_llava_onevision import image_size_to_num_patches

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from config.loader import load_config
from data_utils.chart.data_collector import prepare_chart_rl_data
from data_utils.chart.teacher_gate import filter_chartqa_rows_with_teacher_gate
from data_utils.commom_util import collate_fn
from opsd_utils.teacher_batching import student_batch_num_images_tensor


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="config/config_opd_only_7b_chartqa.yaml")
    parser.add_argument("--iterations", type=int, default=3)
    parser.add_argument("--samples-per-rank", type=int, default=2)
    return parser.parse_args()


def _deepspeed_config(*, world_size: int, micro_batch_size: int) -> dict:
    root = Path(__file__).resolve().parents[1]
    payload = json.loads((root / "configs/deepspeed/zero2_bf16.json").read_text(encoding="utf-8"))
    payload = copy.deepcopy(payload)
    payload.pop("scheduler", None)
    payload["optimizer"]["params"] = {
        "lr": 1.0e-6,
        "betas": [0.9, 0.999],
        "eps": 1.0e-8,
        "weight_decay": 0.0,
        "torch_adam": True,
        "adam_w_mode": True,
    }
    payload["gradient_accumulation_steps"] = 1
    payload["gradient_clipping"] = 0.0
    payload["train_micro_batch_size_per_gpu"] = micro_batch_size
    payload["train_batch_size"] = world_size * micro_batch_size
    return payload


def _highest_resolution_rows(config: dict, count: int) -> tuple[list[dict], list[int]]:
    rows = prepare_chart_rl_data(config["dataset"]["train_dataset"])
    for index, row in enumerate(rows):
        row["_dyme_index"] = index
    rows, _ = filter_chartqa_rows_with_teacher_gate(rows, config["dataset"]["teacher_gate"])

    processor = AutoProcessor.from_pretrained(config["model"]["pretrained_model_path"], local_files_only=True)
    grid = processor.image_processor.image_grid_pinpoints
    patch_size = processor.image_processor.size["height"]
    ranked: list[tuple[int, int, dict]] = []
    for index, row in enumerate(rows):
        with Image.open(row["image"]) as image:
            num_patches = image_size_to_num_patches(image.size, grid, patch_size)
        ranked.append((int(num_patches), index, row))
    ranked.sort(key=lambda item: (-item[0], item[1]))
    chosen = ranked[:count]
    return [row for _patches, _index, row in chosen], [patches for patches, _index, _row in chosen]


def main() -> None:
    args = parse_args()
    if args.iterations <= 0 or args.samples_per_rank <= 0:
        raise ValueError("--iterations and --samples-per-rank must be positive")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    deepspeed.init_distributed(dist_backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    device = torch.device("cuda", local_rank)
    try:
        config = load_config(args.config)
        candidate_count = world_size * args.samples_per_rank
        rows, patch_counts = _highest_resolution_rows(config, candidate_count)
        offset = rank * args.samples_per_rank
        local_rows = rows[offset : offset + args.samples_per_rank]
        local_patch_counts = patch_counts[offset : offset + args.samples_per_rank]
        if len(local_rows) != args.samples_per_rank:
            raise RuntimeError("not enough ChartQA rows for distributed high-resolution diagnostic")

        model_path = config["model"]["pretrained_model_path"]
        processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
        processor.tokenizer.padding_side = "left"
        model = LlavaOnevisionForConditionalGeneration.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            local_files_only=True,
            low_cpu_mem_usage=True,
        )
        model.base_model.vision_tower.requires_grad_(False)
        model.config.use_cache = False
        model.to(device)
        trainable_parameters = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)

        engine, _, _, _ = deepspeed.initialize(
            model=model,
            model_parameters=[parameter for parameter in model.parameters() if parameter.requires_grad],
            config=_deepspeed_config(world_size=world_size, micro_batch_size=args.samples_per_rank),
        )
        bucket_size = int(engine.optimizer.reduce_bucket_size)

        batch = collate_fn(local_rows, processor, label_id=None)
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        pixel_values = batch["pixel_values"].to(device=device, dtype=torch.bfloat16)
        image_sizes = batch["image_sizes"].to(device)
        batch_num_images = student_batch_num_images_tensor(pixel_values, input_ids.size(0))

        dist.barrier()
        started = time.perf_counter()
        for _ in range(args.iterations):
            outputs = engine(
                input_ids=input_ids,
                attention_mask=attention_mask,
                pixel_values=pixel_values,
                image_sizes=image_sizes,
                batch_num_images=batch_num_images,
                logits_to_keep=2,
            )
            loss = outputs.logits.float().mean()
            engine.backward(loss)
            engine.step()
        torch.cuda.synchronize(device)
        elapsed = time.perf_counter() - started
        print(
            "ZeRO-2 VLM bounded-bucket diagnostic passed: "
            f"rank={rank}/{world_size} patch_counts={local_patch_counts} "
            f"input_shape={tuple(input_ids.shape)} bucket={bucket_size} "
            f"trainable={trainable_parameters} iterations={args.iterations} "
            f"elapsed_s={elapsed:.3f} loss={float(loss.detach().item()):.6f}",
            flush=True,
        )
        dist.barrier()
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()


if __name__ == "__main__":
    main()