import os
import json
from transformers import HfArgumentParser, set_seed
import torch.distributed as dist
from src.datasets import build_data_module
from src.custom_args import CustomTrainingArguments, ModelArguments
from src.train.trainers import CustomTrainer
from src.custom_utils import get_logger
from src.models import Ret2Config, Ret2Model

logger = get_logger(__name__)


if __name__ == '__main__':
    if not dist.is_initialized():
        dist.init_process_group('nccl')

    parser = HfArgumentParser((CustomTrainingArguments, ModelArguments))
    training_args, model_args, = parser.parse_args_into_dataclasses()
    set_seed(training_args.seed)

    logger.info(f"Start training code.")

    config = Ret2Config(
        text_name_or_path=model_args.text_name_or_path,
        text_frozen=model_args.text_frozen,
        text_layer_strategy=model_args.text_layer_strategy,
        text_global_feats_proj_size=model_args.text_global_feats_proj_size,
        vision_name_or_path=model_args.vision_name_or_path,
        vision_frozen=model_args.vision_frozen,
        vision_layer_strategy=model_args.vision_layer_strategy,
        vision_global_feats_proj_size=model_args.vision_global_feats_proj_size,
        tokenizer_max_length=model_args.tokenizer_max_length,
        w_final_size=model_args.w_final_size,
        logit_scale_init_value=model_args.logit_scale_init_value,
        logit_scale_clamp=model_args.logit_scale_clamp,
        hidden_size=model_args.hidden_size,
        dropout_p=model_args.dropout_p,
        attention_dropout=model_args.attention_dropout,
        input_gate_bias=model_args.input_gate_bias,
        forget_gate_bias=model_args.forget_gate_bias          
    )
    model = Ret2Model(config)

    data_module = build_data_module(
        training_args,
        tokenizer=model.tokenizer,
        image_processor=model.image_processor
    )

    trainer = CustomTrainer(
        model=model,
        args=training_args,
        model_args=model_args,
        data_collator=data_module['data_collator'],
        train_dataset=data_module['train_dataset'],
        eval_dataset=data_module['eval_dataset'],
        tokenizer=model.tokenizer,
        compute_metrics=None
    )

    if training_args.do_only_eval:
        eval_output = trainer.evaluate()
        for k, v in eval_output.items():
            logger.info(f"{k}: {v}", logger.info)
        with open(os.path.join(training_args.output_dir, 'eval_output.json'), 'w') as f:
            json.dump(eval_output, f)
        logger.info(f"Evaluation finished.")    
    else:
        resume_from_checkpoint = None
        if training_args.resume_from_checkpoint:
            resume_from_checkpoint = training_args.resume_from_checkpoint
        elif training_args.resume_from_last_checkpoint:
            resume_from_checkpoint = True

        train_output = trainer.train(
            resume_from_checkpoint=resume_from_checkpoint)
        
        trainer.save_model(training_args.output_dir)

        logger.info(f"Training finished. Checkpoint saved.")
