Scripts/pretrain.py

"""
Kurdish Handwritten Paragraph Recognition - Pre-training Script
DenseNet121-Transformer Architecture with Curriculum Learning

Pre-trains the model on synthetic paragraph images before fine-tuning
on real handwritten paragraphs.

Usage:
    python pretrain.py --data_dir ./data/SyntheticParagraphs_12000 --vocab_path ./vocab.json
    python pretrain.py --data_dir ./data/SyntheticParagraphs_12000 --vocab_path ./vocab.json --no_curriculum
"""

import os
import glob
import time
import argparse
import json
import math
import random
import numpy as np
from PIL import Image
from datetime import datetime

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.transforms import InterpolationMode
from torch.nn import functional as F
from torch.amp import autocast, GradScaler
from tqdm import tqdm
import gc


# ===============================
# Argument Parser
# ===============================

def parse_args():
    parser = argparse.ArgumentParser(
        description="Kurdish Handwritten Paragraph Recognition - Pre-training")

    # Data paths
    parser.add_argument("--data_dir", type=str, required=True,
                        help="Root directory with Training/ and Validation/ subfolders")
    parser.add_argument("--vocab_path", type=str, required=True,
                        help="Path to vocabulary JSON file (vocab.json)")

    # Image dimensions
    parser.add_argument("--img_height", type=int, default=600)
    parser.add_argument("--img_width", type=int, default=1235)
    parser.add_argument("--max_seq_len", type=int, default=555)

    # Training hyperparameters
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_epochs", type=int, default=80)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--grad_clip", type=float, default=5.0)
    parser.add_argument("--weight_decay", type=float, default=1e-4)
    parser.add_argument("--seed", type=int, default=42)

    # Model architecture
    parser.add_argument("--hidden_size", type=int, default=256)
    parser.add_argument("--encoder_layers", type=int, default=3)
    parser.add_argument("--decoder_layers", type=int, default=6)
    parser.add_argument("--num_heads", type=int, default=8)
    parser.add_argument("--ff_dim", type=int, default=2048)
    parser.add_argument("--dropout", type=float, default=0.3)
    parser.add_argument("--use_upsample", action="store_true", default=True,
                        help="Enable horizontal upsampling layer (default: True)")
    parser.add_argument("--no_upsample", action="store_true",
                        help="Disable horizontal upsampling layer")

    # Teacher forcing
    parser.add_argument("--tf_noise_rate", type=float, default=0.15,
                        help="Teacher forcing noise rate (default: 0.15)")

    # Curriculum learning
    parser.add_argument("--no_curriculum", action="store_true",
                        help="Disable curriculum learning (train on all data from start)")

    # LR scheduler
    parser.add_argument("--lr_step_size", type=int, default=15,
                        help="StepLR step size in epochs")
    parser.add_argument("--lr_gamma", type=float, default=0.5,
                        help="StepLR decay factor")

    # Early stopping
    parser.add_argument("--patience", type=int, default=15)

    # Training options
    parser.add_argument("--mixed_precision", action="store_true", default=True)
    parser.add_argument("--no_mixed_precision", action="store_true")
    parser.add_argument("--no_aug", action="store_true",
                        help="Disable data augmentation")

    # CER computation
    parser.add_argument("--cer_every", type=int, default=5,
                        help="Compute train CER every N epochs (0 to disable)")
    parser.add_argument("--cer_max_samples", type=int, default=256,
                        help="Max samples for train CER computation")

    # Output
    parser.add_argument("--output_dir", type=str, default="./output",
                        help="Directory to save model and logs")
    parser.add_argument("--model_name", type=str, default="pretrained_model",
                        help="Base name for saved model file")

    return parser.parse_args()


# ===============================
# Vocabulary Loader
# ===============================

def load_vocabulary(vocab_path):
    """Load vocabulary from JSON file."""
    with open(vocab_path, "r", encoding="utf-8") as f:
        vocab_data = json.load(f)

    if "vocab_list" in vocab_data:
        char_list = vocab_data["vocab_list"]
    elif "char_to_idx" in vocab_data:
        mapping = vocab_data["char_to_idx"]
        char_list = [None] * len(mapping)
        for char, idx in mapping.items():
            char_list[idx] = char
    else:
        raise ValueError("Vocabulary JSON must contain 'vocab_list' or 'char_to_idx'")

    char_to_idx = {char: idx for idx, char in enumerate(char_list)}
    idx_to_char = {idx: char for idx, char in enumerate(char_list)}

    return char_list, char_to_idx, idx_to_char


# Special token indices (fixed by convention)
PAD_TOKEN = 0
SOS_TOKEN = 1
EOS_TOKEN = 2


# ===============================
# Helper Functions
# ===============================

def tensor_to_text(tensor, idx_to_char):
    """Convert a tensor of character indices to text."""
    if isinstance(tensor, torch.Tensor):
        tensor = tensor.cpu().tolist()
    text = ""
    for idx in tensor:
        if idx == PAD_TOKEN or idx == SOS_TOKEN:
            continue
        if idx == EOS_TOKEN:
            break
        if idx in idx_to_char:
            text += idx_to_char[idx]
    return text


def count_lines_in_text(text):
    """Count the number of lines in a paragraph text."""
    if not text:
        return 0
    return text.count('\n') + 1


# ===============================
# Curriculum Learning
# ===============================

# Default schedule: progressive difficulty over 80 epochs
DEFAULT_CURRICULUM = [
    (1,  8,  1, 1),   # Epochs 1-8:   1 line only
    (9,  16, 1, 2),   # Epochs 9-16:  1-2 lines
    (17, 28, 2, 3),   # Epochs 17-28: 2-3 lines
    (29, 40, 2, 4),   # Epochs 29-40: 2-4 lines
    (41, 52, 3, 5),   # Epochs 41-52: 3-5 lines
    (53, 64, 3, 6),   # Epochs 53-64: 3-6 lines
    (65, 80, 4, 7),   # Epochs 65-80: 4-7 lines (full complexity)
]


def categorize_paragraphs_by_lines(data_dir):
    """Group paragraph samples by their line count."""
    categories = {}

    image_files = []
    for ext in ["*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg"]:
        image_files.extend(glob.glob(os.path.join(data_dir, ext)))
        image_files.extend(glob.glob(os.path.join(data_dir, ext.upper())))
    image_files = sorted(list(set(image_files)))

    for img_path in image_files:
        label_path = os.path.splitext(img_path)[0] + ".txt"
        if not os.path.exists(label_path):
            continue

        try:
            with open(label_path, "r", encoding="utf-8") as f:
                text = f.read().strip()
        except Exception:
            try:
                with open(label_path, "r", encoding="utf-8-sig") as f:
                    text = f.read().strip()
            except Exception:
                continue

        num_lines = count_lines_in_text(text)
        if num_lines not in categories:
            categories[num_lines] = []
        categories[num_lines].append((img_path, text))

    return categories


def get_curriculum_stage(epoch, schedule):
    """Get the min/max line range for the current epoch."""
    for start_epoch, end_epoch, min_lines, max_lines in schedule:
        if start_epoch <= epoch <= end_epoch:
            return min_lines, max_lines
    return 1, 7


def filter_paragraphs_by_lines(categories, min_lines, max_lines):
    """Filter paragraphs to include only those within the line range."""
    filtered = []
    for num_lines, paragraphs in categories.items():
        if min_lines <= num_lines <= max_lines:
            filtered.extend(paragraphs)
    return filtered


# ===============================
# Dataset
# ===============================

class KurdishParagraphDataset(data.Dataset):
    """Dataset for Kurdish handwritten paragraph images."""

    def __init__(self, root_dir=None, transform=None, max_samples=None,
                 max_seq_len=555, filtered_data=None, img_height=600,
                 img_width=1235, char_to_idx=None):
        self.transform = transform
        self.max_seq_len = max_seq_len
        self.img_height = img_height
        self.img_width = img_width
        self.char_to_idx = char_to_idx

        if filtered_data is not None:
            self.data = filtered_data
        else:
            self.data = []
            image_files = []
            for ext in ["*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg"]:
                image_files.extend(glob.glob(os.path.join(root_dir, ext)))
                image_files.extend(glob.glob(os.path.join(root_dir, ext.upper())))
            image_files = sorted(list(set(image_files)))

            for img_path in image_files:
                label_path = os.path.splitext(img_path)[0] + ".txt"
                if not os.path.exists(label_path):
                    continue
                try:
                    with open(label_path, "r", encoding="utf-8") as f:
                        text = f.read().strip()
                except Exception:
                    try:
                        with open(label_path, "r", encoding="utf-8-sig") as f:
                            text = f.read().strip()
                    except Exception:
                        continue
                if len(text) > 0:
                    self.data.append((img_path, text))

        if max_samples and max_samples < len(self.data):
            random.shuffle(self.data)
            self.data = self.data[:max_samples]

        label = "filtered" if filtered_data else root_dir
        print(f"  Loaded {len(self.data)} paragraph images ({label})")

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

    def __getitem__(self, idx):
        img_path, text = self.data[idx]

        image = Image.open(img_path).convert("RGB")
        orig_width, orig_height = image.size

        # Aspect-ratio-preserving resize
        scale = min(self.img_width / orig_width, self.img_height / orig_height)
        new_width = int(orig_width * scale)
        new_height = int(orig_height * scale)
        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

        # Right-aligned on white canvas (RTL script)
        canvas = Image.new('RGB', (self.img_width, self.img_height), (255, 255, 255))
        x_offset = self.img_width - new_width
        canvas.paste(image, (x_offset, 0))

        if self.transform:
            canvas = self.transform(canvas)

        # Encode text to indices
        indices = ([SOS_TOKEN] +
                   [self.char_to_idx.get(c, self.char_to_idx.get(" ", 0)) for c in text] +
                   [EOS_TOKEN])
        if len(indices) > self.max_seq_len:
            indices = indices[:self.max_seq_len - 1] + [EOS_TOKEN]

        target = torch.LongTensor(indices)
        return canvas, target, len(indices), text


def collate_fn(batch):
    """Collate function with padding for variable-length targets."""
    batch.sort(key=lambda x: x[2], reverse=True)
    images, targets, lengths, texts = zip(*batch)

    images = torch.stack(images, 0)
    max_length = max(lengths)

    padded = torch.ones(len(targets), max_length).long() * PAD_TOKEN
    for i, target in enumerate(targets):
        padded[i, :lengths[i]] = target[:lengths[i]]

    return images, padded, torch.LongTensor(lengths), texts


# ===============================
# Augmentation
# ===============================

def build_train_transform():
    """Training augmentation pipeline for paragraph images."""
    class ParagraphTransform:
        def __call__(self, img):
            if random.random() < 0.5:
                img = transforms.ColorJitter(brightness=0.15, contrast=0.15)(img)

            if random.random() < 0.4:
                img = transforms.RandomAffine(
                    degrees=2, translate=(0.02, 0.02),
                    scale=(0.98, 1.02), shear=(-3, 3),
                    interpolation=InterpolationMode.BILINEAR, fill=255)(img)

            if random.random() < 0.2:
                img = transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))(img)

            img = transforms.ToTensor()(img)

            if random.random() < 0.3:
                noise = torch.randn_like(img) * 0.01
                img = torch.clamp(img + noise, 0.0, 1.0)

            img = transforms.Normalize(
                (0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(img)
            return img

    return ParagraphTransform()


def build_eval_transform():
    """Evaluation transform (normalisation only)."""
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    ])


# ===============================
# Positional Encodings
# ===============================

class PositionalEncoding2D(nn.Module):
    """2D sinusoidal positional encoding for visual feature maps."""

    def __init__(self, d_model, max_h=100, max_w=300):
        super().__init__()
        pe = torch.zeros(max_h, max_w, d_model)
        d_half = d_model // 2

        pos_h = torch.arange(0, max_h, dtype=torch.float).unsqueeze(1)
        div_h = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half))
        pe_h = torch.zeros(max_h, d_half)
        pe_h[:, 0::2] = torch.sin(pos_h * div_h)
        pe_h[:, 1::2] = torch.cos(pos_h * div_h)

        pos_w = torch.arange(0, max_w, dtype=torch.float).unsqueeze(1)
        div_w = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half))
        pe_w = torch.zeros(max_w, d_half)
        pe_w[:, 0::2] = torch.sin(pos_w * div_w)
        pe_w[:, 1::2] = torch.cos(pos_w * div_w)

        for h in range(max_h):
            for w in range(max_w):
                pe[h, w, :d_half] = pe_h[h]
                pe[h, w, d_half:] = pe_w[w]

        self.register_buffer('pe', pe)

    def forward(self, x, height, width):
        _, seq_len, d_model = x.shape
        pe_2d = self.pe[:height, :width, :].reshape(height * width, d_model)
        if seq_len <= pe_2d.size(0):
            pe_2d = pe_2d[:seq_len]
        else:
            pad = torch.zeros(seq_len - pe_2d.size(0), d_model, device=x.device)
            pe_2d = torch.cat([pe_2d, pad], dim=0)
        return x + pe_2d.unsqueeze(0)


class PositionalEncoding1D(nn.Module):
    """1D sinusoidal positional encoding for decoder sequences."""

    def __init__(self, d_model, max_len=1000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1), :]


# ===============================
# CNN Feature Extractor
# ===============================

class CNNFeatureExtractor(nn.Module):
    """DenseNet-121 backbone with optional horizontal upsampling."""

    def __init__(self, output_dim=256, use_upsample=True):
        super().__init__()
        densenet = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT)
        self.features = densenet.features
        backbone_channels = 1024

        if use_upsample:
            self.upsample = nn.Sequential(
                nn.ConvTranspose2d(backbone_channels, 512,
                                   kernel_size=(1, 4), stride=(1, 2), padding=(0, 1)),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True))
            adapt_in = 512
        else:
            self.upsample = None
            adapt_in = backbone_channels

        self.adaptation = nn.Sequential(
            nn.Conv2d(adapt_in, output_dim, kernel_size=1),
            nn.BatchNorm2d(output_dim),
            nn.ReLU(inplace=True))

    def forward(self, x):
        features = F.relu(self.features(x), inplace=True)
        if self.upsample is not None:
            features = self.upsample(features)
        features = self.adaptation(features)
        b, c, h, w = features.shape
        return features.view(b, c, h * w).permute(0, 2, 1), h, w


# ===============================
# Transformer OCR Model
# ===============================

class TransformerOCRParagraphModel(nn.Module):
    """
    DenseNet121-Transformer for end-to-end paragraph recognition.

    Architecture:
        1. DenseNet-121 CNN + optional horizontal upsample
        2. 2D positional encoding + Transformer encoder
        3. Transformer decoder with 1D positional encoding
        4. Linear output projection
    """

    def __init__(self, vocab_size, hidden_size=256, nhead=8,
                 num_encoder_layers=3, num_decoder_layers=6,
                 dim_feedforward=2048, dropout=0.3,
                 use_upsample=True, max_seq_len=555,
                 tf_noise_rate=0.15):
        super().__init__()

        self.max_seq_len = max_seq_len
        self.vocab_size = vocab_size
        self.tf_noise_rate = tf_noise_rate

        self.feature_extractor = CNNFeatureExtractor(
            output_dim=hidden_size, use_upsample=use_upsample)

        self.pos_encoder_2d = PositionalEncoding2D(hidden_size)
        self.pos_decoder_1d = PositionalEncoding1D(hidden_size, max_len=max_seq_len)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size, nhead=nhead,
            dim_feedforward=dim_feedforward, dropout=dropout,
            batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer, num_layers=num_encoder_layers)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_size, nhead=nhead,
            dim_feedforward=dim_feedforward, dropout=dropout,
            batch_first=True)
        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer, num_layers=num_decoder_layers)

        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.output_projection = nn.Linear(hidden_size, vocab_size)
        self.hidden_size = hidden_size

        nn.init.xavier_uniform_(self.token_embedding.weight)
        nn.init.xavier_uniform_(self.output_projection.weight)

    def _generate_square_subsequent_mask(self, sz):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        return mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, 0.0)

    def _add_teacher_forcing_noise(self, tgt_input):
        """Replace random tokens to build decoder robustness."""
        if self.tf_noise_rate <= 0 or not self.training:
            return tgt_input
        noise_mask = (torch.rand_like(tgt_input.float()) < self.tf_noise_rate)
        noise_mask = noise_mask & (tgt_input != PAD_TOKEN) & (tgt_input != SOS_TOKEN)
        random_tokens = torch.randint(3, self.vocab_size, tgt_input.shape, device=tgt_input.device)
        return torch.where(noise_mask, random_tokens, tgt_input)

    def forward(self, src, tgt, tgt_key_padding_mask=None):
        # Encode
        memory, feat_h, feat_w = self.feature_extractor(src)
        memory = self.pos_encoder_2d(memory, feat_h, feat_w)
        memory = self.transformer_encoder(memory)

        # Decode with teacher forcing
        tgt_input = self._add_teacher_forcing_noise(tgt[:, :-1])
        tgt_embedded = self.pos_decoder_1d(self.token_embedding(tgt_input))

        tgt_mask = self._generate_square_subsequent_mask(tgt_embedded.size(1)).to(src.device)
        tgt_pad_mask = tgt_key_padding_mask[:, :-1] if tgt_key_padding_mask is not None else None

        output = self.transformer_decoder(
            tgt_embedded, memory,
            tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_pad_mask)

        return self.output_projection(output)

    def generate_batch(self, imgs, max_length=None):
        """Auto-regressive greedy batch generation."""
        if max_length is None:
            max_length = self.max_seq_len
        self.eval()
        batch_size = imgs.size(0)

        with torch.no_grad():
            memory, feat_h, feat_w = self.feature_extractor(imgs)
            memory = self.pos_encoder_2d(memory, feat_h, feat_w)
            memory = self.transformer_encoder(memory)

            ys = torch.ones(batch_size, 1).fill_(SOS_TOKEN).long().to(imgs.device)
            finished = torch.zeros(batch_size, dtype=torch.bool, device=imgs.device)

            for _ in range(max_length - 1):
                tgt_embedded = self.pos_decoder_1d(self.token_embedding(ys))
                tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(imgs.device)
                out = self.transformer_decoder(tgt_embedded, memory, tgt_mask=tgt_mask)
                out = self.output_projection(out)

                next_tokens = out[:, -1].argmax(dim=-1)
                next_tokens[finished] = PAD_TOKEN
                ys = torch.cat([ys, next_tokens.unsqueeze(1)], dim=1)
                finished = finished | (next_tokens == EOS_TOKEN)
                if finished.all():
                    break

        return [tensor_to_text(seq, idx_to_char) for seq in ys]


# ===============================
# Metrics
# ===============================

def levenshtein_distance(s1, s2):
    if len(s1) < len(s2):
        return levenshtein_distance(s2, s1)
    if len(s2) == 0:
        return len(s1)
    prev = range(len(s2) + 1)
    for c1 in s1:
        curr = [prev[0] + 1]
        for j, c2 in enumerate(s2):
            curr.append(min(prev[j + 1] + 1, curr[j] + 1, prev[j] + (c1 != c2)))
        prev = curr
    return prev[-1]


def calculate_cer(preds, targets):
    total_dist = sum(levenshtein_distance(p, t) for p, t in zip(preds, targets))
    total_chars = sum(len(t) for t in targets)
    return total_dist / max(1, total_chars)


def calculate_wer(preds, targets):
    total_dist = sum(levenshtein_distance(p.split(), t.split()) for p, t in zip(preds, targets))
    total_words = sum(len(t.split()) for t in targets)
    return total_dist / max(1, total_words)


def evaluate_cer_batch(model, dataloader, device, idx_to_char, max_samples=None):
    """Compute CER using batch generation."""
    model.eval()
    all_preds, all_targets = [], []
    count = 0

    with torch.no_grad():
        for images, _, _, texts in tqdm(dataloader, desc="Computing CER"):
            images = images.to(device)
            if max_samples and count + images.size(0) > max_samples:
                images = images[:max_samples - count]
                texts = texts[:max_samples - count]

            preds = model.generate_batch(images)
            all_preds.extend(preds)
            all_targets.extend(texts)
            count += len(preds)

            if max_samples and count >= max_samples:
                break

    return calculate_cer(all_preds, all_targets)


def evaluate_full(model, dataloader, device, idx_to_char):
    """Full evaluation returning CER, WER, predictions, and targets."""
    model.eval()
    all_preds, all_targets = [], []

    with torch.no_grad():
        for images, _, _, texts in tqdm(dataloader, desc="Evaluating"):
            images = images.to(device)
            preds = model.generate_batch(images)
            all_preds.extend(preds)
            all_targets.extend(texts)

    cer = calculate_cer(all_preds, all_targets)
    wer = calculate_wer(all_preds, all_targets)
    return cer, wer, all_preds, all_targets


# ===============================
# Early Stopping
# ===============================

class EarlyStopping:
    def __init__(self, patience=15):
        self.patience = patience
        self.counter = 0
        self.best_cer = float('inf')
        self.early_stop = False

    def __call__(self, val_cer, model, epoch, path):
        if val_cer < self.best_cer:
            self.best_cer = val_cer
            self.counter = 0
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'val_cer': val_cer
            }, path)
            print(f"  Model saved (Val CER: {val_cer:.4f})")
        else:
            self.counter += 1
            print(f"  Early stopping: {self.counter}/{self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True
                print("  Early stopping triggered.")

    def reset(self):
        self.counter = 0


# ===============================
# Training Functions
# ===============================

def train_epoch(model, dataloader, optimizer, criterion, device, scaler,
                use_mixed_precision=True, grad_clip=5.0):
    """Train for one epoch."""
    model.train()
    epoch_loss = 0

    for images, targets, _, _ in tqdm(dataloader, desc="Training"):
        images, targets = images.to(device), targets.to(device)
        tgt_pad_mask = (targets == PAD_TOKEN).to(device)

        optimizer.zero_grad()

        if use_mixed_precision:
            with autocast(device_type='cuda'):
                outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask)
                loss = criterion(outputs.reshape(-1, outputs.shape[-1]),
                                 targets[:, 1:].reshape(-1))
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            scaler.step(optimizer)
            scaler.update()
        else:
            outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask)
            loss = criterion(outputs.reshape(-1, outputs.shape[-1]),
                             targets[:, 1:].reshape(-1))
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()

        epoch_loss += loss.item()

    return epoch_loss / len(dataloader)


def evaluate_loss(model, dataloader, criterion, device, use_mixed_precision=True):
    """Evaluate model loss."""
    model.eval()
    epoch_loss = 0

    with torch.no_grad():
        for images, targets, _, _ in dataloader:
            images, targets = images.to(device), targets.to(device)
            tgt_pad_mask = (targets == PAD_TOKEN).to(device)

            if use_mixed_precision:
                with autocast(device_type='cuda'):
                    outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask)
                    loss = criterion(outputs.reshape(-1, outputs.shape[-1]),
                                     targets[:, 1:].reshape(-1))
            else:
                outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask)
                loss = criterion(outputs.reshape(-1, outputs.shape[-1]),
                                 targets[:, 1:].reshape(-1))

            epoch_loss += loss.item()

    return epoch_loss / len(dataloader)


# ===============================
# Main
# ===============================

def main():
    global idx_to_char  # Used by generate_batch -> tensor_to_text

    args = parse_args()

    # Handle flag conflicts
    use_upsample = args.use_upsample and not args.no_upsample
    use_mixed_precision = args.mixed_precision and not args.no_mixed_precision
    use_curriculum = not args.no_curriculum

    # Seeds
    torch.manual_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)

    # Device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    # Output directory
    os.makedirs(args.output_dir, exist_ok=True)

    # Vocabulary
    char_list, char_to_idx, idx_to_char = load_vocabulary(args.vocab_path)
    vocab_size = len(char_list)
    print(f"Vocabulary size: {vocab_size}")

    # Data directories
    train_dir = os.path.join(args.data_dir, "Training")
    val_dir = os.path.join(args.data_dir, "Validation")

    # Categorize paragraphs by line count
    print("\nCategorizing paragraphs by line count...")
    train_categories = categorize_paragraphs_by_lines(train_dir)
    val_categories = categorize_paragraphs_by_lines(val_dir)

    total_train = sum(len(v) for v in train_categories.values())
    total_val = sum(len(v) for v in val_categories.values())
    print(f"  Training: {total_train} paragraphs")
    print(f"  Validation: {total_val} paragraphs")

    if use_curriculum:
        print("\n  Curriculum schedule:")
        for s, e, mn, mx in DEFAULT_CURRICULUM:
            label = f"{mn} line only" if mn == mx else f"{mn}-{mx} lines"
            print(f"    Epochs {s:2d}-{e:2d}: {label}")

    # Transforms
    train_transform = build_eval_transform() if args.no_aug else build_train_transform()
    eval_transform = build_eval_transform()

    # Dataset common kwargs
    ds_kwargs = dict(
        max_seq_len=args.max_seq_len,
        img_height=args.img_height,
        img_width=args.img_width,
        char_to_idx=char_to_idx)

    # Model
    print("\nInitializing model...")
    model = TransformerOCRParagraphModel(
        vocab_size=vocab_size,
        hidden_size=args.hidden_size,
        nhead=args.num_heads,
        num_encoder_layers=args.encoder_layers,
        num_decoder_layers=args.decoder_layers,
        dim_feedforward=args.ff_dim,
        dropout=args.dropout,
        use_upsample=use_upsample,
        max_seq_len=args.max_seq_len,
        tf_noise_rate=args.tf_noise_rate
    ).to(device)

    total_params = sum(p.numel() for p in model.parameters())
    print(f"  Parameters: {total_params:,}")
    print(f"  Upsample: {'ON' if use_upsample else 'OFF'}")
    print(f"  Curriculum: {'ON' if use_curriculum else 'OFF'}")
    print(f"  Teacher forcing noise: {args.tf_noise_rate * 100:.0f}%")

    # Optimizer, scheduler, criterion
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate,
                            weight_decay=args.weight_decay)
    scheduler = optim.lr_scheduler.StepLR(optimizer,
                                          step_size=args.lr_step_size,
                                          gamma=args.lr_gamma)
    criterion = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN)
    scaler = GradScaler('cuda') if use_mixed_precision else None
    early_stopping = EarlyStopping(patience=args.patience)

    best_model_path = os.path.join(args.output_dir, f"{args.model_name}.pth")

    # Log file
    log_path = os.path.join(args.output_dir,
                            f"{args.model_name}_LOG_{datetime.now():%Y%m%d_%H%M%S}.txt")
    log_file = open(log_path, 'w', encoding='utf-8')

    def log(msg):
        print(msg)
        log_file.write(msg + '\n')
        log_file.flush()

    log(f"\nPre-training started: {datetime.now():%Y-%m-%d %H:%M:%S}")
    log(f"Config: {vars(args)}")

    # Training loop
    current_min_lines = None
    current_max_lines = None
    train_loader = None
    val_loader = None

    for epoch in range(1, args.num_epochs + 1):
        start_time = time.time()

        # Curriculum stage management
        if use_curriculum:
            new_min, new_max = get_curriculum_stage(epoch, DEFAULT_CURRICULUM)

            if new_min != current_min_lines or new_max != current_max_lines:
                current_min_lines, current_max_lines = new_min, new_max

                train_filtered = filter_paragraphs_by_lines(
                    train_categories, current_min_lines, current_max_lines)
                val_filtered = filter_paragraphs_by_lines(
                    val_categories, current_min_lines, current_max_lines)

                label = (f"{current_min_lines} line only" if current_min_lines == current_max_lines
                         else f"{current_min_lines}-{current_max_lines} lines")
                log(f"\n  Curriculum stage: {label} "
                    f"(train={len(train_filtered)}, val={len(val_filtered)})")

                train_dataset = KurdishParagraphDataset(
                    transform=train_transform, filtered_data=train_filtered, **ds_kwargs)
                val_dataset = KurdishParagraphDataset(
                    transform=eval_transform, filtered_data=val_filtered, **ds_kwargs)

                train_loader = data.DataLoader(
                    train_dataset, batch_size=args.batch_size, shuffle=True,
                    num_workers=0, collate_fn=collate_fn, pin_memory=True)
                val_loader = data.DataLoader(
                    val_dataset, batch_size=args.batch_size, shuffle=False,
                    num_workers=0, collate_fn=collate_fn, pin_memory=True)

                early_stopping.reset()
        else:
            if train_loader is None:
                all_train = [p for ps in train_categories.values() for p in ps]
                all_val = [p for ps in val_categories.values() for p in ps]

                train_dataset = KurdishParagraphDataset(
                    transform=train_transform, filtered_data=all_train, **ds_kwargs)
                val_dataset = KurdishParagraphDataset(
                    transform=eval_transform, filtered_data=all_val, **ds_kwargs)

                train_loader = data.DataLoader(
                    train_dataset, batch_size=args.batch_size, shuffle=True,
                    num_workers=0, collate_fn=collate_fn, pin_memory=True)
                val_loader = data.DataLoader(
                    val_dataset, batch_size=args.batch_size, shuffle=False,
                    num_workers=0, collate_fn=collate_fn, pin_memory=True)

        # Train
        train_loss = train_epoch(model, train_loader, optimizer, criterion,
                                 device, scaler, use_mixed_precision, args.grad_clip)

        # Train CER (periodic)
        train_cer = None
        if args.cer_every > 0 and epoch % args.cer_every == 0:
            train_cer = evaluate_cer_batch(model, train_loader, device,
                                           idx_to_char, args.cer_max_samples)

        # Validation
        val_loss = evaluate_loss(model, val_loader, criterion, device, use_mixed_precision)
        val_cer = evaluate_cer_batch(model, val_loader, device, idx_to_char)

        scheduler.step()
        elapsed = time.time() - start_time
        mins, secs = divmod(elapsed, 60)

        # Log
        cer_str = f", Train CER: {train_cer:.4f}" if train_cer is not None else ""
        log(f"Epoch {epoch}/{args.num_epochs} ({mins:.0f}m {secs:.0f}s) | "
            f"Train Loss: {train_loss:.4f}{cer_str} | "
            f"Val Loss: {val_loss:.4f} | Val CER: {val_cer:.4f}")

        # Early stopping and model saving
        early_stopping(val_cer, model, epoch, best_model_path)
        if early_stopping.early_stop:
            break

        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    # Final evaluation on full validation set
    log(f"\nLoading best model for final evaluation...")
    ckpt = torch.load(best_model_path, map_location=device)
    model.load_state_dict(ckpt['model_state_dict'])

    all_val = [p for ps in val_categories.values() for p in ps]
    full_val_dataset = KurdishParagraphDataset(
        transform=eval_transform, filtered_data=all_val, **ds_kwargs)
    full_val_loader = data.DataLoader(
        full_val_dataset, batch_size=args.batch_size, shuffle=False,
        num_workers=0, collate_fn=collate_fn, pin_memory=True)

    final_cer, final_wer, preds, targets = evaluate_full(
        model, full_val_loader, device, idx_to_char)

    log(f"\nFinal Validation Results (Full Set, {len(full_val_dataset)} paragraphs):")
    log(f"  CER: {final_cer:.4f}")
    log(f"  WER: {final_wer:.4f}")

    log(f"\nSample Predictions:")
    for i in range(min(3, len(preds))):
        log(f"\n--- Sample {i + 1} ---")
        log(f"Predicted: {preds[i][:200]}")
        log(f"Actual:    {targets[i][:200]}")

    log(f"\nPre-training complete: {datetime.now():%Y-%m-%d %H:%M:%S}")
    log(f"Best model saved to: {best_model_path}")

    log_file.close()
    print(f"Log saved to: {log_path}")


if __name__ == "__main__":
    main()