import dataclasses
from enum import Enum
from typing import Any, Dict, List, Literal, Optional

from cehrbert_data.decorators.patient_event_decorator_base import AttType

from cehrbert.data_generators.hf_data_generator.meds_to_cehrbert_conversion_rules import (
    MedsToBertMimic4,
    MedsToCehrBertConversion,
    MedsToCehrbertOMOP,
)

# Create an enum dynamically from the list
MedsToCehrBertConversionType = Enum(
    "MedsToCehrBertConversionType",
    [cls.__name__ for cls in MedsToCehrBertConversion.__subclasses__()],
)


class FineTuneModelType(Enum):
    POOLING = "pooling"
    LSTM = "lstm"


@dataclasses.dataclass
class DataTrainingArguments:
    """Arguments pertaining to what data we are going to input our model for training and eval."""

    data_folder: Optional[str] = dataclasses.field(
        metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_prepared_path: Optional[str] = dataclasses.field(
        metadata={"help": "The folder in which the prepared dataset is cached"}
    )
    test_data_folder: Optional[str] = dataclasses.field(
        default=None,
        metadata={"help": "The name of the test dataset to use (via the datasets library)."},
    )
    cohort_folder: Optional[str] = dataclasses.field(
        default=None,
        metadata={"help": "The name of the cohort generated by ACES or OHDSI cohort builder."},
    )
    chronological_split: Optional[bool] = dataclasses.field(
        default=False,
        metadata={
            "help": "A flag to indicate whether the data will be split chronologically, "
            "where the historical data is used for training "
            "and the future data is used for validation adn testing"
        },
    )
    split_by_patient: Optional[bool] = dataclasses.field(
        default=False,
        metadata={
            "help": "A flag to indicate whether the records associated with the same person_id "
            "should end up in the same split"
        },
    )
    validation_split_percentage: Optional[float] = dataclasses.field(
        default=0.05,
        metadata={"help": "The percentage of the train set used as validation set in case there's no validation split"},
    )
    validation_split_num: Optional[int] = dataclasses.field(
        default=1000,
        metadata={"help": "The number of the train set used as validation set in case there's no validation split"},
    )
    test_eval_ratio: Optional[float] = dataclasses.field(
        default=0.5,
        metadata={"help": "The percentage of the train set used as validation set in case there's no validation split"},
    )
    preprocessing_num_workers: Optional[int] = dataclasses.field(
        default=4,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    preprocessing_batch_size: Optional[int] = dataclasses.field(
        default=10000,
        metadata={"help": "The batch size to use for preprocessing a streaming dataset"},
    )
    att_function_type: Literal[
        AttType.CEHR_BERT.value,
        AttType.DAY.value,
        AttType.WEEK.value,
        AttType.MONTH.value,
        AttType.MIX.value,
        AttType.NONE.value,
    ] = dataclasses.field(
        default=AttType.CEHR_BERT.value,
        metadata={
            "help": "The ATT type to choose the level of granularity to use for creating the "
            "artificial time tokens between visits",
            "choices": f"choices={[e.value for e in AttType]}",
        },
    )
    is_data_in_meds: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "The boolean indicator to indicate whether the data is in the MEDS format"},
    )
    inpatient_att_function_type: Literal[
        AttType.CEHR_BERT.value,
        AttType.DAY.value,
        AttType.WEEK.value,
        AttType.MONTH.value,
        AttType.MIX.value,
        AttType.NONE.value,
    ] = dataclasses.field(
        default=AttType.NONE,
        metadata={
            "help": "The ATT type to choose the level of granularity to use for creating the "
            "artificial time tokens between neighboring events within inpatient visits."
            "Default to None, meaning the inpatient artificial time tokens are not created.",
            "choices": f"choices={[e.value for e in AttType]}",
        },
    )
    meds_exclude_tables: Optional[List[str]] = dataclasses.field(
        default_factory=list,
        metadata={"help": "The tables to exclude in the conversion e.g. measurement"},
    )
    # TODO: Python 3.9/10 do not support dynamic unpacking, we have to manually provide the entire
    #  list right now.
    meds_to_cehrbert_conversion_type: Literal[
        MedsToCehrBertConversionType[MedsToBertMimic4.__name__],
        MedsToCehrBertConversionType[MedsToCehrbertOMOP.__name__],
    ] = dataclasses.field(
        default=MedsToCehrBertConversionType[MedsToBertMimic4.__name__],
        metadata={
            "help": "The MEDS to CEHRBERT conversion type e.g. MedsToBertMimic4",
            "choices": f"choices={[e for e in MedsToCehrBertConversionType.__members__]}",
        },
    )
    include_auxiliary_token: Optional[bool] = dataclasses.field(
        default=False,
        metadata={
            "help": "The boolean indicator to indicate whether visit type should be included "
            "at the beginning of the visit and discharge facility should be included at "
            "the end of the visit"
        },
    )
    include_demographic_prompt: Optional[bool] = dataclasses.field(
        default=False,
        metadata={
            "help": "The boolean indicator to indicate whether the demographic tokens should be "
            "added at the beginning of the sequence "
            "including start_year, start_age, gender, race."
        },
    )
    streaming: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "The boolean indicator to indicate whether the data should be streamed"},
    )
    vocab_size: Optional[int] = dataclasses.field(
        default=50_000,
        metadata={"help": "The maximum vocab size allowed for the tokenizer trainer to use"},
    )
    min_frequency: Optional[int] = dataclasses.field(
        default=0,
        metadata={"help": "The minimum frequency for concepts to be kept by the tokenizer"},
    )
    min_num_tokens: Optional[int] = dataclasses.field(
        default=20,
        metadata={"help": "The minimum num of tokens required in a sequences to be included for training"},
    )
    shuffle_records: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Indicates whether to randomly shuffle the records that have the same rank"},
    )
    offline_stats_capacity: Optional[int] = dataclasses.field(
        default=100,
        metadata={
            "help": "The minimum num of lab values to collect for the truncated offline statistics before switching to the online statistics calculation"
        },
    )
    value_outlier_std: Optional[float] = dataclasses.field(
        default=3.0,
        metadata={"help": "The lower quantile for excluding the extreme lower lab values"},
    )


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

    or train from scratch.
    """

    model_name_or_path: Optional[str] = dataclasses.field(
        metadata={
            "help": (
                "The model checkpoint for weights initialization. "
                "Don't set if you want to train a model from scratch."
            )
        },
    )
    tokenizer_name_or_path: Optional[str] = dataclasses.field(
        metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    early_stopping_patience: Optional[int] = dataclasses.field(
        default=1,
        metadata={
            "help": "stop training when the specified metric worsens " "for `early_stopping_patience` evaluation calls."
        },
    )
    cache_dir: Optional[str] = dataclasses.field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded " "from huggingface.co"},
    )
    use_auth_token: bool = dataclasses.field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. "
            "Please use `token` instead."
        },
    )
    trust_remote_code: bool = dataclasses.field(
        default=False,
        metadata={
            "help": (
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
                "execute code present on the Hub on your local machine."
            )
        },
    )
    torch_dtype: Optional[str] = dataclasses.field(
        default=None,
        metadata={
            "help": (
                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
                "dtype will be automatically derived from the model's weights."
            ),
            "choices": ["auto", "bfloat16", "float16", "float32"],
        },
    )
    hidden_size: Optional[int] = dataclasses.field(
        default=128,
        metadata={"help": "The embedding and hidden size for the transformer block"},
    )
    num_hidden_layers: Optional[int] = dataclasses.field(
        default=6,
        metadata={"help": "The number of layers used in the transformer model"},
    )
    n_head: Optional[int] = dataclasses.field(
        default=8, metadata={"help": "The number of heads in Multi-Head Attention"}
    )
    max_position_embeddings: Optional[int] = dataclasses.field(
        default=512,
        metadata={"help": "The maximum length of the sequence allowed for the transformer model"},
    )
    finetune_model_type: Literal[FineTuneModelType.POOLING.value, FineTuneModelType.LSTM.value] = dataclasses.field(
        default=FineTuneModelType.POOLING.value,
        metadata={
            "help": "The finetune model type to choose from",
            "choices": f"choices={[e.value for e in FineTuneModelType]}",
        },
    )
    use_lora: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "The flag to indicate whether or not to use the Lora adapter for finetuning"},
    )
    lora_rank: Optional[int] = dataclasses.field(
        default=16, metadata={"help": "Lora attention dimension (the “rank”)."}
    )
    lora_alpha: Optional[int] = dataclasses.field(
        default=16, metadata={"help": "The alpha parameter for Lora scaling."}
    )
    target_modules: Optional[List[str]] = dataclasses.field(
        default_factory=lambda: ["query", "value"],
        metadata={
            "help": "The names of the modules to apply the adapter to. If this is specified, only the modules with the "
            "specified names will be replaced. When passing a string, a regex match will be performed. When "
            "passing a list of strings, either an exact match will be performed or it is checked if the name "
            "of the module ends with any of the passed strings. If this is specified as ‘all-linear’, "
            "then all linear/Conv1D modules are chosen, excluding the output layer. If this is not specified, "
            "modules will be chosen according to the model architecture. If the architecture is not known, "
            "an error will be raised — in this case, you should specify the target modules manually."
        },
    )
    lora_dropout: Optional[float] = dataclasses.field(
        default=0.1, metadata={"help": "The dropout probability for Lora layers"}
    )
    exclude_position_ids: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Whether or not to exclude position ids from the transformer model"},
    )
    include_values: Optional[bool] = dataclasses.field(
        default=False,
        metadata={"help": "Whether or not to include values into the model"},
    )
    use_sub_time_tokenization: Optional[bool] = dataclasses.field(
        default=True,
        metadata={"help": "Whether or not to decompose the time interval into year/month/day"},
    )
    include_value_prediction: Optional[bool] = dataclasses.field(
        default=True,
        metadata={"help": "Whether or not to include value prediction head for cehrbert"},
    )
    include_ttv_prediction: Optional[bool] = dataclasses.field(
        default=True,
        metadata={"help": "Whether or not to include the time to visit prediction"},
    )
    time_token_loss_weight: Optional[float] = dataclasses.field(
        default=1.0, metadata={"help": "The weight of the time token loss"}
    )
    time_to_visit_loss_weight: Optional[float] = dataclasses.field(
        default=1.0, metadata={"help": "The weight of the time to visit loss"}
    )

    def as_dict(self) -> Dict[str, Any]:
        return dataclasses.asdict(self)
