import os
import torch

MODE = os.environ.get("TRANSLIT_MODE", "sentence")   # "sentence" | "ime"


class Config:
    # ---------------- common ----------------
    d_model    = 384
    nhead      = 6
    num_layers = 4
    dim_ff     = 1536
    dropout    = 0.1

    batch_size   = 1024
    lr           = 7e-4
    warmup_steps = 2000
    weight_decay = 0.01
    label_smooth = 0.1
    grad_clip    = 1.0
    num_workers  = 24
    amp_dtype    = torch.bfloat16

    max_src_len = 256
    max_tgt_len = 320

    dry_run_samples = 20_000
    dry_run_steps   = 100

    # ---------------- mode specific ----------------
    if MODE == "ime":
        train_path = "data/word_train.jsonl"
        val_path   = "data/word_val.jsonl"
        test_path  = "data/word_test.jsonl"
        ckpt_dir   = "checkpoints_ime"
        max_len    = 64
        batch_size = 1024         # big dataset again -> big batch
        epochs     = 10           # millions of words -> few epochs suffice
        lr         = 7e-4
        warmup_steps = 2000
        dropout    = 0.1          # plenty of data -> normal dropout
    else:
        # sentence-level Malayalam -> Manglish (original)
        train_path = "data/train.jsonl"
        val_path   = "data/val.jsonl"
        test_path  = "data/test.jsonl"
        ckpt_dir   = "checkpoints"
        max_len    = 512
        epochs     = 8


CFG = Config()
print(f"[config] mode={MODE}  data={CFG.train_path}  ckpt={CFG.ckpt_dir}")