#!/usr/bin/env python
# coding=utf-8
# Copyright 2023 Statistics and Machine Learning Research Group at HKUST. All rights reserved.
"""A one-line summary of the module or program, terminated by a period.

Leave one blank line.  The rest of this docstring should contain an
overall description of the module or program.  Optionally, it may also
contain a brief description of exported classes and functions and/or usage
examples.

Typical usage example:

  foo = ClassFoo()
  bar = foo.FunctionBar()
"""

import sys
import os
# Apply VILA path patches before any VILA imports
try:
    import vila_path_patcher
except ImportError:
    print("Warning: vila_path_patcher not found, VILA paths may not be remapped", flush=True)

# sys.path.remove(os.path.abspath(os.path.dirname(sys.argv[0])))
from transformers import HfArgumentParser, AutoProcessor

from lmflow.args import (
    ModelArguments,
    MultiModalDatasetArguments,
    AutoArguments,
)

from lmflow.datasets.dataset import Dataset
from lmflow.models.auto_model import AutoModel
from lmflow.pipeline.auto_pipeline import AutoPipeline
from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

def main():
	# Parses arguments
    pipeline_name = "finetuner"
    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)

    parser = HfArgumentParser((ModelArguments, MultiModalDatasetArguments, PipelineArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, pipeline_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()

    # Robust multiprocessing defaults on HPC/Lustre to avoid SemLock/FileNotFound errors
    try:
        import torch
        import torch.multiprocessing as mp
        # Prefer fork on Linux when available; fall back silently if already set
        if mp.get_start_method(allow_none=True) != "fork":
            mp.set_start_method("fork", force=True)
        # Use file_system sharing to reduce reliance on OS semaphores
        torch.multiprocessing.set_sharing_strategy("file_system")
    except Exception:
        pass

    # Cap worker counts defensively to reduce spawn pressure on constrained clusters
    if hasattr(pipeline_args, "dataloader_num_workers") and pipeline_args.dataloader_num_workers:
        pipeline_args.dataloader_num_workers = max(0, min(4, int(pipeline_args.dataloader_num_workers)))
    if hasattr(data_args, "preprocessing_num_workers") and data_args.preprocessing_num_workers:
        data_args.preprocessing_num_workers = max(0, min(4, int(data_args.preprocessing_num_workers)))

    # Initialization
    finetuner = AutoPipeline.get_pipeline(
        pipeline_name=pipeline_name,
        model_args=model_args,
        data_args=data_args,
        pipeline_args=pipeline_args,
    )
    
    # Initialize model first to get tokenizer (needed for VILA mixtures)
    model = AutoModel.get_model(model_args)
    
    # Attach tokenizer to data_args for VILA mixture support
    data_args.tokenizer = model.tokenizer
    
    # Support backend selection via environment variable or default to custom_multi_modal
    dataset_backend = os.environ.get("DATASET_BACKEND", "webdataset_multi_modal")
    dataset = Dataset(data_args, backend=dataset_backend)
    
    # Set bd_size and mask_id on the backend model
    if model_args.mdm:
        model.backend_model.bd_size = model_args.bd_size
        model.backend_model.config.bd_size = model_args.bd_size
        if hasattr(model.backend_model, 'model') and hasattr(model.backend_model.model, 'language_model'):
            model.backend_model.model.language_model.mask_id = model.tokenizer.encode("|<MASK>|")[0]
    
        model.backend_model.complementary_mask = model_args.complementary_mask
        model.backend_model.config.complementary_mask = model_args.complementary_mask
        
        model.backend_model.always_mask_im_end = getattr(model_args, 'always_mask_im_end', False)
        model.backend_model.config.always_mask_im_end = getattr(model_args, 'always_mask_im_end', False)
        
        model.backend_model.flexible_bd_size = getattr(model_args, 'flexible_bd_size', False)
        model.backend_model.config.flexible_bd_size = getattr(model_args, 'flexible_bd_size', False)

        model.backend_model.use_block_causal_mask = getattr(model_args, 'use_block_causal_mask', False)
        model.backend_model.config.use_block_causal_mask = getattr(model_args, 'use_block_causal_mask', False)

        model.backend_model.anneal_block_size = getattr(model_args, 'anneal_block_size', False)
        model.backend_model.config.anneal_block_size = getattr(model_args, 'anneal_block_size', False)

        model.backend_model.minimum_noise_level = getattr(model_args, 'minimum_noise_level', 1e-3)
        model.backend_model.config.minimum_noise_level = getattr(model_args, 'minimum_noise_level', 1e-3)

        model.backend_model.entropy_loss = getattr(model_args, 'entropy_loss', False)
        model.backend_model.config.entropy_loss = getattr(model_args, 'entropy_loss', False)

        model.backend_model.entropy_loss_weight = getattr(model_args, 'entropy_loss_weight', 1.0)
        model.backend_model.config.entropy_loss_weight = getattr(model_args, 'entropy_loss_weight', 1.0)

        model.backend_model.enable_efficient_vision_embed = getattr(model_args, 'enable_efficient_vision_embed', False)
        model.backend_model.config.enable_efficient_vision_embed = getattr(model_args, 'enable_efficient_vision_embed', False)

        model.backend_model.block_causal_no_dynamic = getattr(model_args, 'block_causal_no_dynamic', False)
        model.backend_model.config.block_causal_no_dynamic = getattr(model_args, 'block_causal_no_dynamic', False)

        data_args.bd_size = model_args.bd_size
    else:
        data_args.bd_size = None

    # Optionally freeze the vision encoder to avoid DeepSpeed ZeRO gradient
    # allreduce deadlock when some ranks get text-only batches.
    if getattr(model_args, 'freeze_vision_encoder', False):
        frozen_count = 0
        for param in model.backend_model.model.visual.parameters():
            param.requires_grad = False
            frozen_count += 1
        print(f"🧊 Froze vision encoder: {frozen_count} parameters", flush=True)

    # Build Qwen2.5-VL processor for vision/text preprocessing in collate
    processor_id = getattr(model_args, "tokenizer_name", None) or getattr(model_args, "model_name_or_path", None)
    processor_kwargs = {}
    # Optional control via env vars (defaults to 448x448 to limit vision tokens)
    _min_px = os.environ.get("MIN_PIXELS")
    _max_px = os.environ.get("MAX_PIXELS", 705600)
    _rh = os.environ.get("RESIZED_HEIGHT")
    _rw = os.environ.get("RESIZED_WIDTH")
    if _min_px:
        processor_kwargs["min_pixels"] = int(_min_px)
    if _max_px:
        processor_kwargs["max_pixels"] = int(_max_px)
    if _rh:
        processor_kwargs["resized_height"] = int(_rh)
    if _rw:
        processor_kwargs["resized_width"] = int(_rw)

    # Log processor configuration for debugging
    if processor_kwargs:
        print(f"🔧 Processor kwargs from env vars: {processor_kwargs}", flush=True)
    else:
        print("ℹ️ No processor kwargs from env vars (using defaults)", flush=True)

    processor = AutoProcessor.from_pretrained(
        processor_id,
        trust_remote_code=True,
        **processor_kwargs,
    )
    
    # Verify processor configuration
    if hasattr(processor, 'image_processor'):
        img_proc = processor.image_processor
        if hasattr(img_proc, 'min_pixels'):
            print(f"✅ Processor min_pixels: {img_proc.min_pixels}", flush=True)
        if hasattr(img_proc, 'max_pixels'):
            print(f"✅ Processor max_pixels: {img_proc.max_pixels}", flush=True)
        if hasattr(img_proc, 'resized_height'):
            print(f"✅ Processor resized_height: {img_proc.resized_height}", flush=True)
        if hasattr(img_proc, 'resized_width'):
            print(f"✅ Processor resized_width: {img_proc.resized_width}", flush=True)

    # Finetuning
    # Always use LMFlow's Qwen collator (supports MDM)
    # VILA datasets are wrapped by VILADatasetAdapter to return messages format
    data_collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=model.tokenizer,
        bd_size=getattr(model_args, "bd_size", None) if model_args.mdm else None,
        max_length=getattr(data_args, "max_length", None),
        block_causal_no_dynamic=getattr(model_args, "block_causal_no_dynamic", False),
    )

    tuned_model = finetuner.tune(model=model, dataset=dataset, data_collator=data_collator)


if __name__ == '__main__':
    main()
