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

import numpy as np
import torch
import torch.nn as nn
from accelerate import PartialState
from datasets import builder, load_dataset
from tqdm import tqdm
from transformers import (
    AutoConfig,
    AutoModel,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    PretrainedConfig,
    PreTrainedModel,
    Trainer,
)
from transformers.modeling_outputs import SequenceClassifierOutputWithPast

from trl import ModelConfig, RewardConfig


builder.has_sufficient_disk_space = lambda needed_bytes, directory=".": True

tqdm.pandas()


@dataclass
class ScriptArguments:
    """
    The name of the Casual LM model we wish to fine with RewardTrainer
    """

    dataset_name: Optional[str] = field(
        default="mnoukhov/dpo_20konly_1b_fp16.yml_1a838_generations", metadata={"help": "the dataset name"}
    )
    tokenizer_name: Optional[str] = field(default=None, metadata={"help": "the tokenizer name"})


def get_kbit_device_map() -> Optional[Dict[str, int]]:
    # if is_xpu_available():
    #     return {"": f"xpu:{PartialState().local_process_index}"}
    if torch.cuda.is_available():
        return {"": PartialState().local_process_index}
    else:
        return None


def get_quantization_config(model_config: ModelConfig) -> Optional[BitsAndBytesConfig]:
    if model_config.load_in_4bit:
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=model_config.torch_dtype,  # For consistency with model weights, we use the same value as `torch_dtype`
            bnb_4bit_quant_type=model_config.bnb_4bit_quant_type,
            bnb_4bit_use_double_quant=model_config.use_bnb_nested_quant,
        )
    elif model_config.load_in_8bit:
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=True,
        )
    else:
        quantization_config = None

    return quantization_config


def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.normal_(layer.weight, std=std)
    torch.nn.init.constant_(layer.bias, val=bias_const)
    return layer


class ScalarModelConfig(PretrainedConfig):
    def __init__(
        self,
        base_model: str = "EleutherAI/pythia-160m",
        base_config: PretrainedConfig = AutoConfig.from_pretrained("EleutherAI/pythia-160m"),
        hidden_size: int = 768,
        bias: float = 0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.base_model = base_model
        self.base_config = base_config
        self.hidden_size = hidden_size
        self.bias = bias


class ScalarModel(PreTrainedModel):
    config_class = ScalarModelConfig
    _supports_flash_attn_2 = True

    def __init__(self, config: ScalarModelConfig, **base_model_kwargs):
        super().__init__(config)
        # if config.base_model == "models/EleutherAI/pythia-6.9b-deduped/sft_model_55513":
        #     config.base_model = "vwxyzjn/EleutherAI_pythia-6.9b-deduped__sft__tldr"
        #     config.base_model_revision = "sft__55513__1706646024"
        #     config.base_config["_name_or_path"] = "vwxyzjn/EleutherAI_pythia-6.9b-deduped__sft__tldr"
        #     config.base_config["revision"] = "sft__55513__1706646024"
        self.config = config
        self.lm_backbone = AutoModel.from_pretrained(
            config.base_model,
            revision=getattr(config, "base_model_revision", None),
            config=self.config.base_config,
            trust_remote_code=True,
            **base_model_kwargs,
        )
        self.scalar_head = layer_init(
            nn.Linear(self.config.hidden_size, 1),
            std=1 / np.sqrt(self.config.hidden_size + 1),
        )

    def forward(self, input_ids, attention_mask, output_hidden_states=True, return_dict=True, **kwargs):
        output = self.lm_backbone(input_ids, attention_mask, output_hidden_states=True, return_dict=True, **kwargs)
        reward_logits = self.scalar_head(output.hidden_states[-1]) - self.config.bias

        sequence_lengths = first_true_indices(input_ids[:, :] == self.config.pad_token_id) - 1
        # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454
        reward = reward_logits[
            torch.arange(reward_logits.size(0), device=reward_logits.device), sequence_lengths
        ].squeeze(-1)

        return_dict = kwargs.pop("return_dict", None)
        if not return_dict:
            return_values = (reward,) + output[1:]

        return SequenceClassifierOutputWithPast(loss=None, logits=reward)


def first_true_indices(bools, dtype=torch.long):
    """
    Takes an N-dimensional bool tensor and returns an (N-1)-dimensional tensor of integers giving
    the position of the first True in each "row".

    Returns the length of the rows (bools.size(-1)) if no element is True in a given row.
    """
    row_len = bools.size(-1)
    zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device)
    return torch.min(zero_or_index, dim=-1).values


if __name__ == "__main__":
    parser = HfArgumentParser((RewardConfig, ModelConfig, ScriptArguments))
    reward_config, model_config, script_args = parser.parse_args_into_dataclasses()
    reward_config.gradient_checkpointing_kwargs = dict(use_reentrant=False)

    ################
    # Model & Tokenizer
    ################
    torch_dtype = (
        model_config.torch_dtype
        if model_config.torch_dtype in ["auto", None]
        else getattr(torch, model_config.torch_dtype)
    )
    tokenizer = AutoTokenizer.from_pretrained(script_args.tokenizer_name)
    tokenizer.add_special_tokens({"pad_token": "[PAD]"})

    scalar_model_config = ScalarModelConfig.from_pretrained(
        model_config.model_name_or_path,
        revision=model_config.model_revision,
    )
    # hack to remove the path
    # models/EleutherAI/pythia-6.9b-deduped/sft_model_55513 -> EleutherAI/pythia-6.9b-deduped
    if scalar_model_config.base_model.startswith("models/"):
        original_model = scalar_model_config.base_config["_name_or_path"].split("/")[2]
        sft_model = f"vwxyzjn/EleutherAI_{original_model}__sft__tldr"
        scalar_model_config.base_config["_name_or_path"] = sft_model
        scalar_model_config.base_model = sft_model
        scalar_model_config.base_model_revision = "sft__55513__1706646024"

    quantization_config = get_quantization_config(model_config)
    model = ScalarModel.from_pretrained(
        model_config.model_name_or_path,
        revision=model_config.model_revision,
        device_map=get_kbit_device_map() if quantization_config is not None else None,
        quantization_config=quantization_config,
        config=scalar_model_config,
    )

    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        args=reward_config,
    )

    datasets = []
    epochs = [1] + list(range(1000, 10000, 1000))
    for epoch in epochs:
        dataset = load_dataset(script_args.dataset_name, revision=str(epoch), split="train[:100]")

        dataset = dataset.map(
            lambda example: tokenizer(
                example["generations"] + "<|endoftext|>",
                padding="max_length",
                max_length=reward_config.max_length,
                truncation=True,
            ),
            batched=True,
        )

        results = trainer.predict(dataset)
        # raw_datasets = raw_datasets.filter(
        #     lambda x: len(x["input_ids_chosen"]) <= reward_config.max_length
        #     and len(x["input_ids_rejected"]) <= reward_config.max_length
        # )
        # train_dataset = raw_datasets["train"]
        # eval_dataset = raw_datasets["test"]