#!/usr/bin/env python3
"""ByT5 text-only pretrain on Hittite parallel corpus.

Cuneiform Unicode → Transliteration, 280K pairs.
Reference: SumTablets ACL ML4AL 2024.
"""
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/seq2seq_hitit.yaml'))
    ap.add_argument('--phase', default='pretrain')
    ap.add_argument('--data', default=str(ROOT / 'datasets/processed/tlhdig_corpora/hittite_parallel.jsonl'))
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/seq2seq_text_pretrain/'))
    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('decoder', {}).get('pretrained', 'google/byt5-small')
    log(f"Phase: {args.phase}, Model: {model_name}")

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)

    # Data
    log(f"Loading {args.data}")
    examples = []
    with open(args.data) 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: {len(examples)} pairs")

    n_val = int(len(examples) * 0.05)
    val_examples = examples[:n_val]
    train_examples = examples[n_val:]

    def preprocess(batch):
        inputs = tokenizer(batch['input_text'], max_length=512, truncation=True)
        targets = tokenizer(batch['target_text'], max_length=512, truncation=True)
        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 = cfg.get('training', {})
    targs = Seq2SeqTrainingArguments(
        output_dir=str(output),
        num_train_epochs=3,  # pretrain epochs (shortened for 1-day target)
        per_device_train_batch_size=training.get('batch_size', 8),
        per_device_eval_batch_size=16,
        learning_rate=1e-4,
        weight_decay=0.01,
        warmup_steps=2000,
        bf16=True,
        fp16=False,
        logging_steps=200,
        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,
        dataloader_num_workers=4,
        report_to='none',
    )

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

    log("Training text-only pretrain...")
    trainer.train()

    trainer.save_model(str(output / 'checkpoint'))
    tokenizer.save_pretrained(str(output / 'checkpoint'))
    # Also save symlink/copy as checkpoint.pt for seq2seq_ft
    import shutil
    # Save in trainer format; train_image loads via HF from_pretrained
    log(f"DONE: {output}/checkpoint")

if __name__ == '__main__':
    main()