#!/usr/bin/env python3
"""Image → Transliteration seq2seq model.

Reference: SumTablets Simmons 2024, HATFormer 2024.
Architecture: ViT encoder + ByT5 decoder + cross-attention.
"""
import torch
import torch.nn as nn
from pathlib import Path

class VisionByT5(nn.Module):
    """ViT encoder + ByT5 decoder for cuneiform transliteration."""
    
    def __init__(self, encoder_name="facebook/dinov2-large", 
                 decoder_name="google/byt5-small",
                 image_size=384):
        super().__init__()
        
        # Encoder (lazy import — transformers library)
        from transformers import AutoModel, T5ForConditionalGeneration, AutoTokenizer
        self.encoder = AutoModel.from_pretrained(encoder_name)
        self.decoder = T5ForConditionalGeneration.from_pretrained(decoder_name)
        self.tokenizer = AutoTokenizer.from_pretrained(decoder_name)
        
        # Projection: encoder hidden → decoder hidden
        enc_dim = self.encoder.config.hidden_size
        dec_dim = self.decoder.config.d_model
        self.enc_to_dec_proj = nn.Linear(enc_dim, dec_dim)
        
        # LoRA adapt encoder
        self._apply_lora(self.encoder, r=32, alpha=64)
    
    def _apply_lora(self, module, r=32, alpha=64):
        try:
            from peft import LoraConfig, get_peft_model
            config = LoraConfig(
                r=r, lora_alpha=alpha,
                target_modules=["query", "key", "value", "dense"],
                lora_dropout=0.1, bias="none"
            )
            module = get_peft_model(module, config)
        except ImportError:
            print("peft not installed; LoRA skipped", flush=True)
    
    def forward(self, pixel_values, labels=None):
        """
        pixel_values: (B, 3, H, W)
        labels: (B, T) — target byte sequence
        """
        # Encode image
        enc_out = self.encoder(pixel_values=pixel_values)
        enc_hidden = enc_out.last_hidden_state  # (B, N_patches, enc_dim)
        
        # Project to decoder dim
        enc_hidden = self.enc_to_dec_proj(enc_hidden)
        
        # Decode
        outputs = self.decoder(
            encoder_outputs=(enc_hidden,),
            labels=labels,
        )
        return outputs  # loss + logits
    
    @torch.no_grad()
    def generate(self, pixel_values, max_length=512, num_beams=5):
        enc_out = self.encoder(pixel_values=pixel_values)
        enc_hidden = self.enc_to_dec_proj(enc_out.last_hidden_state)
        
        return self.decoder.generate(
            encoder_outputs=(enc_hidden,),
            max_length=max_length,
            num_beams=num_beams,
            length_penalty=0.6,
            no_repeat_ngram_size=3,
            early_stopping=True,
        )

if __name__ == '__main__':
    print("Vision-ByT5 seq2seq model architecture template hazır.")
    print("Training data: tablet image + transliteration pair")
    print("Expected chrF (SumTablets baseline): 0.9755")