"""
 Copyright 2024 Intel Corporation

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      https://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """

from typing import List
from dataclasses import dataclass, field
from transformers import AutoTokenizer

@dataclass
class DeltaTunerArguments:
    search_engine: str = field(default="EvolutionarySearchEngine", metadata={"help": "Which search strategy used in the DE-NAS. Can be 'EvolutionarySearchEngine', 'RandomSearchEngine' and 'SigoptSearchEngine'"})
    model_id: str = field(default=None, metadata={"help": "The model id in the hugging face or the local model path"})
    batch_size: int = field(default=4, metadata={"help": "The batch size used in the search process to calculate the inference"})
    max_epochs: int = field(default=10, metadata={"help": "The max epochs of the search processing"})
    scale_factor: int = field(default=10, metadata={"help": "The scaling factor used in the EA search"})
    select_num: int = field(default=50, metadata={"help": "The num of top candidates stored in the EA search process"})
    population_num: int = field(default=50, metadata={"help": "The num of population candidates used in the EA/Random search process"})
    m_prob: float = field(default=0.2, metadata={"help": "The m probability to produce new mutation network depth"})
    s_prob: float = field(default=0.4, metadata={"help": "The s probability to produce new other mutation network structures"})
    crossover_num: int = field(default=25, metadata={"help": "The number of crossover in the EA search process"})
    mutation_num: int = field(default=25, metadata={"help": "The number of mutation in the EA search process"})
    img_size: int = field(default=32, metadata={"help": "The input or output sequence length of the model."})
    expressivity_weight: float = field(default=0, metadata={"help": "The weight of expressivity score in the de-score"})
    complexity_weight: float = field(default=0, metadata={"help": "The weight of complexity score in the de-score"})
    diversity_weight: float = field(default=1, metadata={"help": "The weight of diversity score in the de-score"})
    saliency_weight: float = field(default=1, metadata={"help":"The weight of saliency score in the de-score"})
    latency_weight: float = field(default=0, metadata={"help":"The weight of latency score in the de-score"})
    max_param_limits: int = field(default=None, metadata={"help": "The max M number of trainable parameter size"})
    min_param_limits: int = field(default=None, metadata={"help": "The min M number of trainable parameter size"})
    budget_latency_max: int = field(default=None, metadata={"help": "The max latency limit of the candidate network (etc., '3*10**6' second for your estimation of max latency of LLM forward)"})
    budget_latency_min: int = field(default=None, metadata={"help": "The min latency limit of the candidate network (etc., '1*10**6' second for your estimation of min latency of LLM forward)"})
    layer_name: str = field(default="num_hidden_layers", metadata={"help": "The layer name for a specific model search space"})
    best_model_structure: str = field(default=None, metadata={"help": "The file path to load the best model structure"})
    search_space_name: List[str] = field(default=None, metadata={"help": "The search space name of delta layers"})
    random_seed: int = field(default=12345, metadata={"help": "The random seed of the deltatuner search process"})
    denas: bool = field(default=True, metadata={"help": "Whether to use the denas"})
    ssf_target_module: List[str] = field(default_factory=lambda: None, metadata={"help": "Target modules for the SSF method."},)
    task_type: str = field(default="CAUSAL_LM", metadata={"help": "The task type of the deltatuner model"})
    algo: str = field(default="auto", metadata={"help": "The base algorithm of the deltatuner model, current support algorithms including: lora, ssf, auto"})