#!/usr/bin/env python3
"""ByT5 Hittite transliteration fine-tune.

Input:  cuneiform Unicode (or transliteration tokens)
Output: transliteration string
Data:   TLHdig 280K Hittite parallel pairs

Sequence packing (block-diagonal attention) için HF native DataCollator kullanıyoruz.
"""
import os, sys, json, argparse, time
from pathlib import Path
import yaml

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def log(msg): print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/transliteration_hitit_only.yaml'))
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/hitit_byt5/'))
    args = ap.parse_args()

    cfg = yaml.safe_load(open(args.config))
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)

    try:
        import torch
        from transformers import (AutoTokenizer, T5ForConditionalGeneration,
                                   Seq2SeqTrainer, Seq2SeqTrainingArguments,
                                   DataCollatorForSeq2Seq)
        from datasets import Dataset
    except ImportError as e:
        log(f"Missing library: {e}")
        sys.exit(1)

    model_name = cfg.get('model', {}).get('name', 'google/byt5-small')
    log(f"Loading model {model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)

    # Data: hittite_parallel (cuneiform → transliteration)
    data_path = ROOT / 'datasets/processed/tlhdig_corpora/hittite_parallel.jsonl'
    log(f"Loading data: {data_path}")
    examples = []
    with open(data_path) as f:
        for line in f:
            r = json.loads(line)
            cune = r.get('cuneiform') or ''
            translit = r.get('transliteration') or ''
            if cune and translit and len(cune) < 200 and len(translit) < 400:
                examples.append({
                    'input_text': cune,
                    'target_text': translit,
                })

    log(f"Total examples: {len(examples)}")
    val_frac = cfg.get('data', {}).get('val_fraction', 0.1)
    n_val = int(len(examples) * val_frac)
    val_examples = examples[:n_val]
    train_examples = examples[n_val:]

    log(f"Train: {len(train_examples)}, Val: {len(val_examples)}")

    def preprocess(batch):
        inputs = tokenizer(batch['input_text'], max_length=512, truncation=True, padding=False)
        targets = tokenizer(batch['target_text'], max_length=256, truncation=True, padding=False)
        inputs['labels'] = targets['input_ids']
        return inputs

    train_ds = Dataset.from_list(train_examples).map(preprocess, batched=True, remove_columns=['input_text', 'target_text'])
    val_ds = Dataset.from_list(val_examples).map(preprocess, batched=True, remove_columns=['input_text', 'target_text'])

    # Training
    training = cfg.get('training', {})
    # BF16 availability check (torchrun sometimes loses this)
    import torch
    bf16_supported = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    use_bf16 = training.get('bf16', True) and bf16_supported

    targs = Seq2SeqTrainingArguments(
        output_dir=str(output),
        num_train_epochs=training.get('epochs', 30),
        per_device_train_batch_size=training.get('batch_size', 32),
        per_device_eval_batch_size=training.get('batch_size', 32),
        learning_rate=training.get('lr', 5e-5),
        weight_decay=training.get('weight_decay', 0.01),
        warmup_steps=training.get('warmup_steps', 1000),
        bf16=use_bf16,
        fp16=False,
        logging_steps=100,
        eval_strategy='epoch',
        save_strategy='epoch',
        save_total_limit=2,
        load_best_model_at_end=True,
        metric_for_best_model='eval_loss',
        greater_is_better=False,
        gradient_accumulation_steps=1,
        dataloader_num_workers=4,
        report_to='none',
    )

    collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding='longest')

    trainer = Seq2SeqTrainer(
        model=model,
        args=targs,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        data_collator=collator,
        tokenizer=tokenizer,
    )

    log("Starting training...")
    trainer.train()

    log("Saving final model...")
    trainer.save_model(str(output / 'best'))
    tokenizer.save_pretrained(str(output / 'best'))
    log(f"DONE: {output}")

if __name__ == '__main__':
    main()