import logging
from dataclasses import dataclass, field
from typing import Optional

import transformers
from peft import LoraConfig, get_peft_model
from transformers import (
    Trainer,
    TrainingArguments,
    AutoTokenizer,
    PreTrainedTokenizer
)

from src.dataset.collator import InstructionTuningCollator
from src.dataset.dataset import InstructionTuningDataset
from src.model.init_llemr import init_llemr
from src.model.modeling_llemr import LlemrForConditionalGeneration
from src.model.utils import find_all_linear_names

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    name_or_path: Optional[str] = field(default=None)
    llm_pretrained_model_name_or_path: Optional[str] = field(default="Qwen/Qwen2-0.5B-Instruct")
    train_type: Optional[str] = field(
        default="train_both",
        metadata={
            "help": """
            1. train_multi_modal_projector
            2. train_both
            """
        },
    )
    use_lora: Optional[bool] = field(default=True)
    lora_r: int = 32
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_bias: str = "none"
    vision_hidden_size: int = 768


@dataclass
class DataArguments:
    source: Optional[str] = field(default="note")


def load_model(model_args: ModelArguments):
    if model_args.name_or_path is not None:
        logging.warning(f"Load model {model_args.name_or_path} from pretrained")
        model = LlemrForConditionalGeneration.from_pretrained(
            model_args.name_or_path
        )
        tokenizer = AutoTokenizer.from_pretrained(
            model_args.name_or_path,
            padding_side="left"
        )
    else:
        logging.warning(f"Init model {model_args.llm_pretrained_model_name_or_path}")
        model, tokenizer = init_llemr(
            model_args.llm_pretrained_model_name_or_path, model_args.vision_hidden_size
        )

    assert model_args.train_type in ["train_multi_modal_projector", "train_both"]
    if model_args.train_type == "train_multi_modal_projector":
        logging.warning("Train multi_modal_projector")
        for param in model.language_model.parameters():
            param.requires_grad = False
    else:
        logging.warning("Train both")

    if model_args.use_lora:
        assert model_args.train_type == "train_both"
        logging.warning("Use Lora")

        config = LoraConfig(
            r=model_args.lora_r,
            lora_alpha=model_args.lora_alpha,
            target_modules=find_all_linear_names(model),
            lora_dropout=model_args.lora_dropout,
            bias=model_args.lora_bias,
            task_type="CAUSAL_LM",
            modules_to_save=["multi_modal_projector"],
        )
        model = get_peft_model(model, config)

    else:
        logging.warning("Not use Lora")

    return model, tokenizer


def load_data(data_args: DataArguments, tokenizer: PreTrainedTokenizer):
    train_dataset = InstructionTuningDataset(
        split="train",
        source=data_args.source,
    )
    val_dataset = InstructionTuningDataset(
        split="val",
        source=data_args.source,
    )
    collator = InstructionTuningCollator(
        tokenizer=tokenizer,
    )
    return train_dataset, val_dataset, collator


def train():
    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    model, tokenizer = load_model(model_args)
    train_dataset, val_dataset, collator = load_data(data_args, tokenizer)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=collator,
    )

    tokenizer.save_pretrained(training_args.output_dir)
    trainer.train()


if __name__ == "__main__":
    train()