import os
import sys
from dataclasses import dataclass, field, asdict
from typing import Optional
from transformers import TrainingArguments, HfArgumentParser
from enum import Enum


task_to_keys = {
    "cola": ("sentence", None),
    "mnli": ("premise", "hypothesis"),
    "mrpc": ("sentence1", "sentence2"),
    "qnli": ("question", "sentence"),
    "qqp": ("question1", "question2"),
    "rte": ("sentence1", "sentence2"),
    "sst2": ("sentence", None),
    "stsb": ("sentence1", "sentence2"),
    "wnli": ("sentence1", "sentence2"),
}


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
            "with private models)."
        },
    )
    apply_lora: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply LoRA or not."},
    )
    lora_kind: Optional[str] = field(
        default="LoRA",
        metadata={"help": "The kind of lora layer to instantiate"}
    )
    train_lora_B: Optional[bool] = field(
        default=False,
        metadata={"help": "Train lora_B matrix"}
    )
    peft_type: Optional[str] = field(
        default="none",
        metadata={'help': 'Finetuning via PEFT, can be either AdaLoRA or BOFT'},
    )
    adaptive_ranks: Optional[bool] = field(
        default=False,
        metadata={"help": "Redistribute ranks according to explained variances of PCA."}
    )
    redistribute_from_scratch: Optional[bool] = field(
        default=False,
        metadata={"help": "Redistribute ranks according to explained variances of PCA from scratch."}
    )
    redist_metric: Optional[str] = field(
        default="ratio",
        metadata={"help": "One of ratio, raw, sum, or max"}
    )
    whiten_pca: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to whiten PCA so that components have unit variance."}
    )
    rand_rotate: Optional[bool] = field(
        default=False,
        metadata={"help": "Apply a random rotation matrix to each PCA component to isolate effect of directions."}
    )
    permute_rows: Optional[bool] = field(
        default=False,
        metadata={"help": "Apply a row-wise random permutation matrix to PCA init to isolate effect of scale."}
    )
    init_cls_head: Optional[bool] = field(
        default=False,
        metadata={"help": "Initialize classifier head with PCA"}
    )
    exp_var_threshold: Optional[float] = field(
        default=0.9,
        metadata={"help": "explained variance threshold for adaptive ranks"}
    )
    lora_alpha: Optional[int] = field(
        default=1,
        metadata={"help": "LoRA alpha"},
    )
    apply_ssf: Optional[bool] = field(
        default=False,
        metadata={"help": "Perform SSF on fetures"},
    )
    lora_r: Optional[int] = field(
        default=0,
        metadata={"help": "LoRA r"},
    )
    lora_path: Optional[str] = field(
        default=None,
        metadata={"help": "The file path of LoRA parameters."},
    )
    apply_adapter: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply adapter or not."},
    )
    adapter_path: Optional[str] = field(
        default=None,
        metadata={"help": "The file path of adapter parameters."},
    )
    adapter_type: Optional[str] = field(
        default='houlsby',
        metadata={"help": "houlsby or pfeiffer"},
    )
    adapter_size: Optional[int] = field(
        default=64,
        metadata={"help": "8, 16, 32, 64"},
    )
    apply_bitfit: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply bitfit or not."},
    )
    reg_loss_wgt: Optional[float] = field(
        default=0.0,
        metadata={"help": "Regularization Loss Weight"},
    )
    masking_prob: Optional[float] = field(
        default=0.0,
        metadata={"help": "Token Masking Probability"},
    )
    post_layer_norm: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply post layer norm in LoRA or not."},
    )
    pre_batch_norm: Optional[bool] = field(
        default=False,
        metadata={"help": "Whether to apply pre batch norm in LoRA or not."},
    )

    def to_dict(self):
        """
        Serializes this instance while replace `Enum` by their values (for JSON serialization support).
        """
        d = asdict(self)
        for k, v in d.items():
            if isinstance(v, Enum):
                d[k] = v.value
            if isinstance(v, list) and len(v) > 0 and isinstance(v[0], Enum):
                d[k] = [x.value for x in v]
        return d


@dataclass
class CustomTrainingArgs(TrainingArguments):
    """
    Arguments pertaining to what data we are going to input our model for training and eval.

    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """
    experiment_name: Optional[str] = field(default=None, metadata={"help": "Experiment name for wandb tracking."})

    task_name: Optional[str] = field(
        default=None,
        metadata={"help": "The name of the task to train on: " + ", ".join(task_to_keys.keys())},
    )
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
            "help": "Whether to pad all samples to `max_seq_length`. "
            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
    max_val_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of validation examples to this "
            "value if set."
        },
    )
    max_test_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of test examples to this "
            "value if set."
        },
    )
    train_file: Optional[str] = field(
        default=None, metadata={"help": "A csv or a json file containing the training data."}
    )
    validation_file: Optional[str] = field(
        default=None, metadata={"help": "A csv or a json file containing the validation data."}
    )
    test_file: Optional[str] = field(default=None, metadata={"help": "A csv or a json file containing the test data."})
    
    data_cache_dir: Optional[str] = field(default=None, metadata={"help": "Directory path to save cache files."})

    def __post_init__(self):
        super().__post_init__()
        if self.task_name is not None:
            self.task_name = self.task_name.lower()
            if self.task_name not in task_to_keys.keys():
                raise ValueError("Unknown task, you should pick one in " + ",".join(task_to_keys.keys()))
        elif self.train_file is None or self.validation_file is None:
            raise ValueError("Need either a GLUE task or a training/validation file.")
        else:
            train_extension = self.train_file.split(".")[-1]
            assert train_extension in ["csv", "json"], "`train_file` should be a csv or a json file."
            validation_extension = self.validation_file.split(".")[-1]
            assert (
                validation_extension == train_extension
            ), "`validation_file` should have the same extension (csv or json) as `train_file`."


def setup_args(): 
    parser = HfArgumentParser((ModelArguments, CustomTrainingArgs))
    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, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, training_args = parser.parse_args_into_dataclasses()
    # force to None, to avoid having paths as runnames
    training_args.run_name = None
    # add arguments in data_args to training_args
    return model_args, training_args, training_args