"""Default hyperparameters shared across training experiments."""

# ── HuggingFace org & auth ───────────────────────────────────────────────────
HF_ORG   = "AITeamUIT"

# Collection slugs — fill in after creating collections on HF Hub UI:
#   https://huggingface.co/AITeamUIT → Collections → New collection
MODEL_COLLECTION_SLUG = "AITeamUIT/securepii-model"  # e.g. "AITeamUIT/securepii-<id>"
BENCH_COLLECTION_SLUG ="AITeamUIT/securepii-bench"   # e.g. "AITeamUIT/securepii-bench-<id>"

# ── Dataset ──────────────────────────────────────────────────────────────────
TRAIN_DATASET_ID  = "quynong/rosenxt-v7-15-6-full"
EVAL_DATASET_ID   = "trinhtrantran122/gen_aug_dataset"
EVAL_SPLIT        = "train"
LANGUAGES         = ["en", "de", "vi"]
MAX_TRAIN_SAMPLES = 1_000_000
MAX_VAL_SAMPLES   = 5000

# ── K-Fold Cross-Validation ──────────────────────────────────────────────────
FOLD_DATASET_ID = "quynong/rosenxt-v7-15-6"   # dataset with fold_k_train / fold_k_test splits
N_FOLDS         = 5

# Fine-tune validation set (secondary eval source)
FINETUNE_VAL_DATASET  = "quynong/rosenxt-v7-15-6-full"
FINETUNE_VAL_SPLIT    = "test"
FINETUNE_VAL_MAX      = 5000
EXTRA_LANGS_FROM_FT   = ["en", "de", "vi"]
USE_EXTRA_FT_VAL      = True

# ── Training defaults (Baseline & DataAug) ───────────────────────────────────
USE_LABEL_DESCRIPTIONS = True
MAX_WIDTH              = 30

DEFAULT_TRAINING = dict(
    num_epochs=3,
    max_steps=-1,
    batch_size=8,
    eval_batch_size=8,
    max_len=1024,          # word-level tokens (WhitespaceTokenSplitter units)
    encoder_lr=1e-5,
    task_lr=5e-4,
    weight_decay=0.01,
    adam_beta1=0.9,
    adam_beta2=0.999,
    adam_epsilon=1e-8,
    max_grad_norm=1.0,
    scheduler_type="linear",
    warmup_ratio=0.1,
    warmup_steps=0,
    num_cycles=0.5,
    fp16=True,
    bf16=False,
    eval_strategy="epoch",
    save_total_limit=2,
    save_best=True,
    metric_for_best="eval_loss",
    greater_is_better=False,
    logging_steps=50,
    logging_first_step=True,
    report_to_wandb=False,
    early_stopping=False,
    early_stopping_patience=3,
    early_stopping_threshold=0.0,
    num_workers=4,
    pin_memory=True,
    prefetch_factor=2,
    seed=42,
    deterministic=False,
    validate_data=True,
    use_lora=False,
    lora_r=16,
    lora_alpha=32.0,
    lora_dropout=0.0,
    lora_target_modules=["encoder", "span_rep", "classifier", "count_embed", "count_pred"],
    save_adapter_only=True,
)

# ── mmBERT overrides (applied on top of DEFAULT_TRAINING) ───────────────────
MMBERT_OVERRIDES = dict(
    num_epochs=6,         # cold-start task heads need more epochs
    max_len=2048,
    task_lr=2e-4,
    adam_epsilon=1e-6,
    fp16=False,
    bf16=True,            # H100 native bfloat16
    eval_strategy="steps",
    eval_steps=2000,
    early_stopping=True,
    early_stopping_patience=2,
)

# ── Data augmentation defaults ───────────────────────────────────────────────
AUG_FACTOR = 1     # augmented copies per selected example
AUG_RATIO  = 0.5   # fraction of train examples to augment
AUG_SEED   = 42

# ── Evaluation defaults ──────────────────────────────────────────────────────
THRESHOLD              = 0.7
BENCHMARK_THRESHOLD    = 0.7   # stricter threshold used in full benchmark runs

# Per-label threshold overrides for schema building.
# Labels listed here use the given threshold; all others use the global
# threshold passed to model.extract() (BENCHMARK_THRESHOLD = 0.7).
PER_LABEL_THRESHOLDS = {
    # Nhóm cần giảm threshold để tăng recall
    "API_KEY": 0.35,
    "PLATE": 0.50,
    "ACCOUNT_ID": 0.55,
    "BANK_ACCOUNT": 0.55,
    "LICENSE": 0.55,
    "MEDICAL_INFO": 0.55,
    "TIME": 0.55,
    "PREFIX": 0.55,
    "ETHNICITY": 0.60,
    "TIN": 0.60,
    "NATIONAL_ID": 0.65,

    # Nhóm cần tăng threshold để giảm false positive
    "TICKET_ID": 0.65,
    "CARD_NUMBER": 0.85,
    "EMPLOYEE_ID": 0.6,
    "JOB_TITLE": 0.85,
    "PIN": 0.88,

    # Nhóm giữ ổn định
    "ADDRESS": 0.70,
    "AGE": 0.70,
    "BIRTHDATE": 0.70,
    "COORDINATE": 0.70,
    "DATE": 0.70,
    "EMAIL": 0.70,
    "IBAN": 0.70,
    "INSURANCE_ID": 0.70,
    "IP": 0.70,
    "LOCATION": 0.70,
    "MARITAL": 0.70,
    "MONEY": 0.70,
    "NATIONALITY": 0.70,
    "ORGANIZATION": 0.70,
    "PASSPORT": 0.70,
    "PASSWORD": 0.70,
    "PERSON": 0.70,
    "PHONE": 0.70,
    "RELIGION": 0.70,
    "SWIFT": 0.70,
    "TRADE_UNION": 0.70,
    "USERNAME": 0.70,

    # Nhóm hơi nghiêng về recall cao, có thể giữ hoặc tăng nhẹ
    "CARD_ISSUER": 0.75,
    "CVV": 0.75,
    "GENDER": 0.75,
    "URL": 0.75,
    "WALLET": 0.75,
    "ZIP_CODE": 0.75,
}
IOU_THRESH             = 0.5
MAX_TOKENS_PER_CHUNK   = 1024
CHUNK_SAFETY_MARGIN    = 20
MAX_CHARS_PER_CHUNK    = 1000
RANDOM_SEED            = 123
PARALLEL_BATCH_SIZE    = 8

# ── Benchmark model registry ─────────────────────────────────────────────────
BENCHMARK_MODELS = [
    {
        "name":         "glinerv2",
        "display":      "glinerv2",
        "repo_id":      "AITeamUIT/gliner2-multi-v1-3e-20260514",
        "subfolder":    None,
        "max_width":    30,
        "chunk":        True,
        "chunk_tokens": 768,
    },
    {
        "name":         "glinerv2_data_aug",
        "display":      "glinerv2 data aug",
        "repo_id":      "AITeamUIT/gliner2-multi-v1-3e-20260514-dataaug",
        "subfolder":    None,
        "max_width":    30,
        "chunk":        True,
        "chunk_tokens": 768,
    },
    {
        "name":         "glinerv2_mmbert_backbone",
        "display":      "glinerv2 mmbert backbone",
        "repo_id":      "AITeamUIT/gliner2-mmbert-base-rosenxt-6e-20260515",
        "subfolder":    None,
        "max_width":    30,
        "chunk":        True,
        "chunk_tokens": 2048,
    },
    {
        "name":         "glinerv2_quy",
        "display":      "glinerv2 quy",
        "repo_id":      "AITeamUIT/gliner2-rosenxt-full-20260515-quy",
        "subfolder":    "best",
        "max_width":    None,
        "chunk":        False,
        "chunk_tokens": None,
    },
    {
        "name":         "gliner_v2_quy_code_sai",
        "display":      "gliner v2 quy code sai",
        "repo_id":      "AITeamUIT/gliner2-rosenxt-multi-v7",
        "subfolder":    "best",
        "max_width":    None,
        "chunk":        False,
        "chunk_tokens": None,
    },
]