

import enum
import os
import sys
sys.path.append("/.../.../.../fastchat/fastchat/train")
sys.path.append("/.../.../.../fastchat")

from dataclasses import dataclass, field
import json
import math
import pathlib
from typing import Dict, Optional, Sequence

import numpy as np
import torch
from torch.utils.data import Dataset
import transformers
from transformers import Trainer
from transformers.trainer_pt_utils import LabelSmoother

from fastchat.conversation import SeparatorStyle

IGNORE_TOKEN_ID = LabelSmoother.ignore_index
from tqdm import tqdm


local_rank = None


def rank0_print(*args):
    if local_rank == 0:
        print(*args)

def rank0_init_wandb():
    if local_rank == 0:
        import wandb
        wandb.init(project="....fschat.numadd", name="fschat-numadd", notes="this is a hello world numadd test with fschat")

@dataclass
class ModelArguments:
    model_name_or_path: Optional[str] = field(default="facebook/...-125m")
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": "Whether or not to allow for custom models defined on the Hub in their own modeling files"
        },
    )
    padding_side: str = field(
        default="right", metadata={"help": "The padding side in tokenizer"}
    )


@dataclass
class DataArguments:
    data_path: str = field(
        default=None, metadata={"help": "Path to the training data."}
    )
    eval_data_path: str = field(
        default=None, metadata={"help": "Path to the evaluation data."}
    )
    lazy_preprocess: bool = False


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    ...im: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )


def trainer_save_model_safe(trainer: transformers.Trainer):
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    from torch.distributed.fsdp import StateDictType, FullStateDictConfig

    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(
        trainer.model, StateDictType.FULL_STATE_DICT, save_policy
    ):
        trainer.save_model()


def preprocess(
    sources,
    tokenizer: transformers.PreTrainedTokenizer,
    tokenizer_max_length = 16
) -> Dict:
    conversations = []
    for i, source in enumerate(sources):
        conversations.append(source['input'] + source['output'])

    input_ids = tokenizer(
        conversations,
        return_tensors="pt",
        padding="max_length",
        max_length=tokenizer_max_length,
        truncation=True,
    ).input_ids
    targets = input_ids.clone()

    for source, target in zip(sources, targets):
        o_target = target.clone()
        total_len = int(target.ne(tokenizer.pad_token_id).sum())

        cur_len = 1
        target[:cur_len] = IGNORE_TOKEN_ID
        instruction_len = len(tokenizer(source['input']).input_ids) - 1
        turn_len = len(tokenizer(source['input'] + source['output']).input_ids)

        target[cur_len : cur_len + instruction_len] = IGNORE_TOKEN_ID

        cur_len += turn_len

        if not tokenizer.legacy:
            cur_len -= 1

        target[cur_len:] = IGNORE_TOKEN_ID


    return dict(
        input_ids=input_ids,
        labels=targets,
        attention_mask=input_ids.ne(tokenizer.pad_token_id),
    )


class SupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer):
        super(SupervisedDataset, self).__init__()

        rank0_print("Formatting inputs...")
        sources = raw_data
        data_dict = preprocess(sources, tokenizer)

        self.raw_data = raw_data
        self.input_ids = data_dict["input_ids"]
        self.labels = data_dict["labels"]
        self.attention_mask = data_dict["attention_mask"]

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        return dict(
            input_ids=self.input_ids[i],
            labels=self.labels[i],
            attention_mask=self.attention_mask[i],
            input_text=self.raw_data[i]['input'],
            target_text=self.raw_data[i]['output']
        )


class LazySupervisedDataset(Dataset):
    """Dataset for supervised fine-tuning."""

    def __init__(self, raw_data, tokenizer: transformers.PreTrainedTokenizer):
        super(LazySupervisedDataset, self).__init__()
        self.tokenizer = tokenizer

        rank0_print("Formatting inputs...Skip in lazy mode")
        self.tokenizer = tokenizer
        self.raw_data = raw_data
        self.cached_data_dict = {}

    def __len__(self):
        return len(self.raw_data)

    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
        if i in self.cached_data_dict:
            return self.cached_data_dict[i]

        ret = preprocess([self.raw_data[i]], self.tokenizer)
        ret = dict(
            input_ids=ret["input_ids"][0],
            labels=ret["labels"][0],
            attention_mask=ret["attention_mask"][0],
            input_text=self.raw_data[i]['input'],
            target_text=self.raw_data[i]['output']
        )
        self.cached_data_dict[i] = ret

        return ret


def make_supervised_data_module(
    tokenizer: transformers.PreTrainedTokenizer, data_args
) -> Dict:
    """Make dataset and collator for supervised fine-tuning."""
    dataset_cls = (
        LazySupervisedDataset if data_args.lazy_preprocess else SupervisedDataset
    )
    rank0_print("Loading data...")

    with open(data_args.data_path, "r") as f:
        train_json = f.readlines()
        train_json = [json.loads(e) for e in train_json]

    train_dataset = dataset_cls(train_json, tokenizer=tokenizer)

    if data_args.eval_data_path:
        with open(data_args.eval_data_path, "r") as f:
            eval_json = f.readlines()
            eval_json = [json.loads(e) for e in eval_json][:100]
        eval_dataset = dataset_cls(eval_json, tokenizer=tokenizer)
    else:
        eval_dataset = None

    return dict(train_dataset=train_dataset, eval_dataset=eval_dataset)




import tempfile
import os
from my_hdfs_io import hcopy,hmkdir
def load_hdfs_path(ckpt_path):
    if ckpt_path.startswith("hdfs"):
        tmp_dir = os.path.join(
            tempfile.gettempdir(), os.path.basename(ckpt_path)
        )
        local_dir = tmp_dir
        hcopy(ckpt_path, local_dir)
    else:
        local_dir = ckpt_path
    return local_dir


from transformers import Trainer
import torch
from tqdm import tqdm

class CustomTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def evaluate(self, eval_dataset=None, ignore_keys=None, metric_key_prefix="eval"):
        metrics = super().evaluate(eval_dataset, ignore_keys, metric_key_prefix)

        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        self.model.eval()
        predictions = []
        references = []

        for batch in tqdm(eval_dataset):
            inputs = self.tokenizer(batch['input_text'], return_tensors='pt', padding=True, truncation=True, max_length=16).to(self.model.device)
            with torch.no_grad():
                outputs = self.model.generate(**inputs, max_new_tokens=10)
            decoded_outputs = self.tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True)
            predictions.extend(decoded_outputs)
            references.append(batch['target_text'])

        em_score, predictions, references = self.calculate_EM(predictions, references)
        print(f"Eval EM score: {em_score}")
        print("predict:", predictions[:5])
        print("ref:", references[:5])

        metrics[f"{metric_key_prefix}_em"] = em_score
        print("metrics:",metrics)
        self.log(metrics)
        return metrics

    @staticmethod
    def calculate_EM(predictions, references):
        predictions = [e.strip() for e in predictions]
        STOP_WORDS = ['\n']
        processed_predictions = []
        for prediction in predictions:
            for stop_word in STOP_WORDS:
                if stop_word in prediction:
                    prediction = prediction.split(stop_word)[0]
                    break
            processed_predictions.append(prediction)

        predictions = processed_predictions
        references = [e.strip() for e in references]
        same = [e1 == e2 for e1, e2 in zip(predictions, references)]
        return float(sum(same)) / float(len(same)), predictions, references

def train():
    global local_rank

    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    local_rank = training_args.local_rank

    rank0_init_wandb()

    config = transformers.AutoConfig.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        trust_remote_code=model_args.trust_remote_code,
    )
    orig_ctx_len = getattr(config, "max_position_embeddings", None)
    if orig_ctx_len and training_args.model_max_length > orig_ctx_len:
        scaling_factor = float(math.ceil(training_args.model_max_length / orig_ctx_len))
        config.rope_scaling = {"type": "linear", "factor": scaling_factor}
    config.use_cache = False

    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        config=config,
        cache_dir=training_args.cache_dir,
        trust_remote_code=model_args.trust_remote_code,
    )
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side=model_args.padding_side,
        use_fast=False,
        trust_remote_code=model_args.trust_remote_code,
    )

    if tokenizer.pad_token != tokenizer.unk_token:
        tokenizer.pad_token = tokenizer.unk_token

    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)














    trainer = CustomTrainer(
        model=model, tokenizer=tokenizer, args=training_args, **data_module
    )
    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()

    model.config.use_cache = True
    trainer.save_state()
    if trainer.is_deepspeed_enabled:
        trainer.save_model()
    else:
        trainer_save_model_safe(trainer)


if __name__ == "__main__":
    train()
