#!/usr/bin/env python3
"""Image → Transliteration end-to-end seq2seq fine-tune.

ViT encoder + ByT5 decoder (cross-attention).
Data: 83 Hitit tablet with full transliteration sequences.
"""
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)

class ImageTransliterationDataset:
    def __init__(self, jsonl_path, processor, tokenizer, max_tgt=512):
        import torch
        self.records = []
        with open(jsonl_path) as f:
            for line in f:
                r = json.loads(line)
                if r.get('image_path') and Path(r['image_path']).exists():
                    self.records.append(r)
        self.processor = processor
        self.tokenizer = tokenizer
        self.max_tgt = max_tgt

    def __len__(self): return len(self.records)

    def __getitem__(self, idx):
        from PIL import Image
        r = self.records[idx]
        img = Image.open(r['image_path']).convert('RGB').resize((384, 384))
        pixel_values = self.processor(img, return_tensors='pt')['pixel_values'][0]
        target = r.get('sequence_plain', '')[:1000]  # truncate
        labels = self.tokenizer(target, max_length=self.max_tgt,
                                 truncation=True, return_tensors='pt')['input_ids'][0]
        return {'pixel_values': pixel_values, 'labels': labels}

def _remap_timm_to_hf_vit(state):
    """Map timm/DINO ViT state_dict keys to HuggingFace ViTModel keys."""
    import torch, re
    out = {}
    for k, v in state.items():
        if k == 'cls_token':
            out['embeddings.cls_token'] = v
        elif k == 'pos_embed':
            out['embeddings.position_embeddings'] = v
        elif k == 'patch_embed.proj.weight':
            out['embeddings.patch_embeddings.projection.weight'] = v
        elif k == 'patch_embed.proj.bias':
            out['embeddings.patch_embeddings.projection.bias'] = v
        elif k == 'norm.weight':
            out['layernorm.weight'] = v
        elif k == 'norm.bias':
            out['layernorm.bias'] = v
        else:
            m = re.match(r'blocks\.(\d+)\.(.+)', k)
            if not m:
                continue
            i, sub = m.group(1), m.group(2)
            base = f'encoder.layer.{i}'
            if sub == 'norm1.weight': out[f'{base}.layernorm_before.weight'] = v
            elif sub == 'norm1.bias': out[f'{base}.layernorm_before.bias'] = v
            elif sub == 'norm2.weight': out[f'{base}.layernorm_after.weight'] = v
            elif sub == 'norm2.bias': out[f'{base}.layernorm_after.bias'] = v
            elif sub == 'attn.qkv.weight':
                d = v.size(0) // 3
                out[f'{base}.attention.attention.query.weight'] = v[:d]
                out[f'{base}.attention.attention.key.weight'] = v[d:2*d]
                out[f'{base}.attention.attention.value.weight'] = v[2*d:]
            elif sub == 'attn.qkv.bias':
                d = v.size(0) // 3
                out[f'{base}.attention.attention.query.bias'] = v[:d]
                out[f'{base}.attention.attention.key.bias'] = v[d:2*d]
                out[f'{base}.attention.attention.value.bias'] = v[2*d:]
            elif sub == 'attn.proj.weight': out[f'{base}.attention.output.dense.weight'] = v
            elif sub == 'attn.proj.bias': out[f'{base}.attention.output.dense.bias'] = v
            elif sub == 'mlp.fc1.weight': out[f'{base}.intermediate.dense.weight'] = v
            elif sub == 'mlp.fc1.bias': out[f'{base}.intermediate.dense.bias'] = v
            elif sub == 'mlp.fc2.weight': out[f'{base}.output.dense.weight'] = v
            elif sub == 'mlp.fc2.bias': out[f'{base}.output.dense.bias'] = v
    return out

def collate_seq2seq(batch, tokenizer):
    import torch
    pixel_values = torch.stack([b['pixel_values'] for b in batch])
    labels = [b['labels'] for b in batch]
    max_len = max(l.size(0) for l in labels)
    padded = torch.full((len(labels), max_len), -100, dtype=torch.long)
    for i, l in enumerate(labels):
        padded[i, :l.size(0)] = l
    return {'pixel_values': pixel_values, 'labels': padded}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/seq2seq_hitit.yaml'))
    ap.add_argument('--phase', default='finetune')
    ap.add_argument('--encoder', help='SSL DINOv3 checkpoint')
    ap.add_argument('--decoder', help='ByT5 pretrained checkpoint')
    ap.add_argument('--data', default=str(ROOT / 'datasets/processed/seq2seq_pairs.jsonl'))
    ap.add_argument('--output', default=str(ROOT / 'hitit_ocr/runs/seq2seq_image_ft/'))
    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
        import torch.nn as nn
        from transformers import (AutoTokenizer,
                                   ViTImageProcessor, ViTModel,
                                   T5ForConditionalGeneration,
                                   Seq2SeqTrainer,
                                   Seq2SeqTrainingArguments)
        from transformers.modeling_outputs import BaseModelOutput
    except ImportError as e:
        log(f"Missing: {e}"); sys.exit(1)

    # Encoder: generic ViT (DINOv3 checkpoint olsa da HF interface için ViT yüklüyoruz)
    enc_name = 'google/vit-base-patch16-384'
    dec_name = args.decoder if args.decoder and Path(args.decoder).exists() else 'google/byt5-small'

    log(f"Encoder: {enc_name}, Decoder: {dec_name}")

    encoder = ViTModel.from_pretrained(enc_name, add_pooling_layer=False)
    t5 = T5ForConditionalGeneration.from_pretrained(dec_name)
    image_processor = ViTImageProcessor.from_pretrained(enc_name)
    tokenizer = AutoTokenizer.from_pretrained(dec_name)

    # ViT hidden size genellikle T5 d_model ile eşleşmez; projeksiyon katmanı ekle
    enc_dim = encoder.config.hidden_size
    dec_dim = t5.config.d_model
    proj = nn.Linear(enc_dim, dec_dim) if enc_dim != dec_dim else nn.Identity()

    class ViTT5Seq2Seq(nn.Module):
        """ViT encoder + T5 decoder; T5'in kendi encoder'ı atlanır, encoder_outputs= verilir."""
        def __init__(self, encoder, t5, proj, decoder_start_token_id, pad_token_id):
            super().__init__()
            self.encoder = encoder
            self.t5 = t5
            self.proj = proj
            self.config = t5.config
            self.config.decoder_start_token_id = decoder_start_token_id
            self.config.pad_token_id = pad_token_id

        def _shift_right(self, input_ids):
            shifted = input_ids.new_full(input_ids.shape, self.config.pad_token_id)
            shifted[:, 1:] = input_ids[:, :-1].clone()
            shifted[:, 0] = self.config.decoder_start_token_id
            shifted.masked_fill_(shifted == -100, self.config.pad_token_id)
            return shifted

        def forward(self, pixel_values=None, labels=None, **kw):
            enc_out = self.encoder(pixel_values=pixel_values).last_hidden_state
            enc_out = self.proj(enc_out)
            encoder_outputs = BaseModelOutput(last_hidden_state=enc_out)
            decoder_input_ids = self._shift_right(labels) if labels is not None else None
            out = self.t5(encoder_outputs=encoder_outputs,
                          decoder_input_ids=decoder_input_ids,
                          labels=labels, return_dict=True)
            return out

        def gradient_checkpointing_enable(self, **kw):
            self.encoder.gradient_checkpointing_enable(**kw)
            self.t5.gradient_checkpointing_enable(**kw)

    model = ViTT5Seq2Seq(encoder, t5, proj,
                         decoder_start_token_id=tokenizer.pad_token_id or 0,
                         pad_token_id=tokenizer.pad_token_id or 0)

    # Load SSL encoder if provided (timm/DINO -> HF ViT key remap)
    if args.encoder and Path(args.encoder).exists():
        log(f"Loading SSL encoder from {args.encoder}")
        ck = torch.load(args.encoder, map_location='cpu')
        bb_state = ck.get('backbone', ck)
        remapped = _remap_timm_to_hf_vit(bb_state)
        own = model.encoder.state_dict()
        # Interpolate position embeddings if grid size differs (e.g. 224->384)
        pe_key = 'embeddings.position_embeddings'
        if pe_key in remapped and pe_key in own and remapped[pe_key].shape != own[pe_key].shape:
            src, tgt = remapped[pe_key], own[pe_key]
            cls_pe, grid_pe = src[:, :1], src[:, 1:]
            n_src, n_tgt = grid_pe.size(1), tgt.size(1) - 1
            gs_src, gs_tgt = int(n_src ** 0.5), int(n_tgt ** 0.5)
            grid_pe = grid_pe.reshape(1, gs_src, gs_src, -1).permute(0, 3, 1, 2)
            grid_pe = torch.nn.functional.interpolate(
                grid_pe, size=(gs_tgt, gs_tgt), mode='bicubic', align_corners=False)
            grid_pe = grid_pe.permute(0, 2, 3, 1).reshape(1, n_tgt, -1)
            remapped[pe_key] = torch.cat([cls_pe, grid_pe], dim=1)
            log(f"  Interpolated pos_embed {gs_src}x{gs_src} -> {gs_tgt}x{gs_tgt}")
        matched = {k: v for k, v in remapped.items() if k in own and own[k].shape == v.shape}
        missing = [k for k in own if k not in matched]
        model.encoder.load_state_dict(matched, strict=False)
        log(f"Matched {len(matched)}/{len(own)} tensors, missing {len(missing)}")
        if missing and len(missing) < 10:
            log(f"  Missing: {missing}")

    # Data
    ds = ImageTransliterationDataset(args.data, image_processor, tokenizer)
    log(f"Data: {len(ds)} tablet")
    n_val = max(1, len(ds) // 10)
    train_ds = [ds[i] for i in range(n_val, len(ds))]
    val_ds = [ds[i] for i in range(n_val)]

    def collate(batch):
        return collate_seq2seq(batch, tokenizer)

    training = cfg.get('training', {})
    targs = Seq2SeqTrainingArguments(
        output_dir=str(output),
        num_train_epochs=training.get('epochs', 60),
        per_device_train_batch_size=training.get('batch_size', 8),
        per_device_eval_batch_size=4,
        learning_rate=training.get('lr_decoder', 1e-4),
        weight_decay=0.01,
        warmup_steps=training.get('warmup_steps', 200),
        bf16=True,
        fp16=False,
        gradient_accumulation_steps=training.get('gradient_accumulation', 4),
        gradient_checkpointing=False,  # static_graph DDP ile uyumsuz; bf16 + H100 80GB zaten yeterli
        logging_steps=20,
        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=2,
        predict_with_generate=False,
        report_to='none',
        ddp_find_unused_parameters=True,
        ddp_broadcast_buffers=False,
        save_safetensors=False,  # T5 shared embed_tokens (encoder/decoder/shared aynı tensor)
    )

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

    log("Training image→transliteration seq2seq...")
    trainer.train()

    best_dir = output / 'best'
    best_dir.mkdir(parents=True, exist_ok=True)
    torch.save(model.state_dict(), best_dir / 'pytorch_model.bin')
    tokenizer.save_pretrained(str(best_dir))
    image_processor.save_pretrained(str(best_dir))
    # best.pt symlink for orchestrator
    torch.save({'model_dir': str(best_dir)}, output / 'best.pt')
    log(f"DONE: {output}")

if __name__ == '__main__':
    main()