Scripts/finetune.py

"""

Kurdish Handwritten Paragraph Recognition - Fine-tuning Script

DenseNet121-Transformer Architecture



Fine-tunes a pre-trained model on real handwritten paragraph images.

Loads weights from pretrain.py output checkpoint.



Usage:

    python finetune.py --data_dir ./data/UniqueHandwrittenParagraphs \

                       --vocab_path ./vocab.json \

                       --pretrained_path ./output/pretrained_model.pth



    python finetune.py --data_dir ./data/DASNUS-Paragraphs \

                       --vocab_path ./vocab.json \

                       --pretrained_path ./output/pretrained_model.pth \

                       --freeze_epochs 10

"""

import os
import glob
import time
import argparse
import json
import math
import random
import re
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 - Fine-tuning")

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

    # 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=5e-5)
    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 (must match pre-trained model)
    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.2)
    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.05,
                        help="Teacher forcing noise rate (default: 0.05)")

    # Encoder freezing
    parser.add_argument("--freeze_epochs", type=int, default=10,
                        help="Number of epochs to freeze CNN encoder (default: 10)")
    parser.add_argument("--encoder_lr_mult", type=float, default=0.1,
                        help="Learning rate multiplier for encoder (default: 0.1)")

    # LR scheduler
    parser.add_argument("--lr_patience", type=int, default=5,
                        help="ReduceLROnPlateau patience")
    parser.add_argument("--lr_factor", type=float, default=0.5,
                        help="ReduceLROnPlateau 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")
    parser.add_argument("--clean_text", action="store_true", default=True,
                        help="Clean invisible Unicode characters from labels")
    parser.add_argument("--no_clean_text", action="store_true")

    # 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="finetuned_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
PAD_TOKEN = 0
SOS_TOKEN = 1
EOS_TOKEN = 2


# ===============================
# Text Cleaning
# ===============================

INVISIBLE_CHARS = [
    '\u200e', '\u200f', '\u200b', '\u200d', '\ufeff', '\u00ad',
    '\u2060', '\u2061', '\u2062', '\u2063', '\u2064',
    '\u206a', '\u206b', '\u206c', '\u206d', '\u206e', '\u206f',
    '\u2028', '\u2029',
]


def clean_text(text):
    """Remove invisible Unicode characters that inflate CER.

    Preserves ZWNJ (U+200C) which is used in Kurdish."""
    for char in INVISIBLE_CHARS:
        if char != '\u200c':  # Keep ZWNJ
            text = text.replace(char, '')
    text = re.sub(r' +', ' ', text)
    lines = text.split('\n')
    lines = [line.strip() for line in lines]
    return '\n'.join(lines)


# ===============================
# 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


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

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

    def __init__(self, root_dir, transform=None, max_seq_len=555,

                 img_height=600, img_width=1235, char_to_idx=None,

                 clean_text_enabled=True):
        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
        self.clean_text_enabled = clean_text_enabled

        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 self.clean_text_enabled:
                text = clean_text(text)
            if len(text) > 0:
                self.data.append((img_path, text))

        print(f"  Loaded {len(self.data)} paragraph images from {root_dir}")

    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

        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)

        canvas = Image.new('RGB', (self.img_width, self.img_height), (255, 255, 255))
        x_offset = self.img_width - new_width  # Right-align for RTL
        canvas.paste(image, (x_offset, 0))

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

        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():
    """Standard augmentation for fine-tuning."""
    class FinetuneTransform:
        def __call__(self, img):
            if random.random() < 0.5:
                img = transforms.ColorJitter(
                    brightness=0.15, contrast=0.15,
                    saturation=0.05, hue=0.01)(img)

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

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

            img = transforms.ToTensor()(img)

            if random.random() < 0.2:
                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 FinetuneTransform()


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."""

    def __init__(self, vocab_size, hidden_size=256, nhead=8,

                 num_encoder_layers=3, num_decoder_layers=6,

                 dim_feedforward=2048, dropout=0.2,

                 use_upsample=True, max_seq_len=555,

                 tf_noise_rate=0.05):
        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):
        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):
        memory, feat_h, feat_w = self.feature_extractor(src)
        memory = self.pos_encoder_2d(memory, feat_h, feat_w)
        memory = self.transformer_encoder(memory)

        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 ys

    def freeze_encoder(self):
        """Freeze CNN backbone parameters."""
        for param in self.feature_extractor.parameters():
            param.requires_grad = False
        print("  Encoder (CNN) frozen")

    def unfreeze_encoder(self):
        """Unfreeze CNN backbone parameters."""
        for param in self.feature_extractor.parameters():
            param.requires_grad = True
        print("  Encoder (CNN) unfrozen")


# ===============================
# Weight Loading
# ===============================

def load_pretrained_weights(model, pretrained_path, device):
    """Load pre-trained weights, handling PE size mismatches gracefully."""
    print(f"\n  Loading pre-trained model: {pretrained_path}")

    if not os.path.exists(pretrained_path):
        raise FileNotFoundError(f"Checkpoint not found: {pretrained_path}")

    ckpt = torch.load(pretrained_path, map_location=device)

    if 'epoch' in ckpt:
        print(f"  Pre-trained epoch: {ckpt['epoch']}")
    if 'val_cer' in ckpt:
        print(f"  Pre-trained Val CER: {ckpt['val_cer']:.4f}")

    state_dict = ckpt.get('model_state_dict', ckpt)
    model_state = model.state_dict()

    loaded, skipped = {}, []
    for key, value in state_dict.items():
        if key in model_state:
            if value.shape == model_state[key].shape:
                loaded[key] = value
            else:
                skipped.append((key, f"{value.shape} vs {model_state[key].shape}"))
        else:
            skipped.append((key, "not in model"))

    model.load_state_dict(loaded, strict=False)

    print(f"  Loaded: {len(loaded)}/{len(model_state)} parameters")
    if skipped:
        print(f"  Skipped: {len(skipped)} (PE buffers regenerated)")

    return model


# ===============================
# 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 calculate_line_accuracy(preds, targets):
    total, correct = 0, 0
    for pred, true in zip(preds, targets):
        pred_lines = pred.split('\n')
        true_lines = true.split('\n')
        total += len(true_lines)
        for pl, tl in zip(pred_lines, true_lines):
            if pl.strip() == tl.strip():
                correct += 1
    return correct / max(1, total)


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 dataloader:
            images = images.to(device)
            if max_samples and count + images.size(0) > max_samples:
                images = images[:max_samples - count]
                texts = texts[:max_samples - count]

            batch_output = model.generate_batch(images)
            preds = [tensor_to_text(seq, idx_to_char) for seq in batch_output]
            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)


# ===============================
# Comprehensive Test Evaluation
# ===============================

def comprehensive_evaluation(model, dataloader, device, idx_to_char):
    """Full evaluation with CER, WER, line accuracy, and timing."""
    model.eval()
    all_preds, all_targets = [], []
    inference_times = []

    # Warmup
    with torch.no_grad():
        for images, _, _, _ in dataloader:
            images = images.to(device)
            _ = model.generate_batch(images[:min(3, images.size(0))])
            break

    if torch.cuda.is_available():
        torch.cuda.synchronize()

    with torch.no_grad():
        for images, _, _, texts in tqdm(dataloader, desc="Evaluating"):
            images = images.to(device)
            batch_size = images.size(0)

            if torch.cuda.is_available():
                torch.cuda.synchronize()
            start = time.perf_counter()

            batch_output = model.generate_batch(images)

            if torch.cuda.is_available():
                torch.cuda.synchronize()
            elapsed = time.perf_counter() - start

            per_sample = elapsed / batch_size
            inference_times.extend([per_sample] * batch_size)

            preds = [tensor_to_text(seq, idx_to_char) for seq in batch_output]
            all_preds.extend(preds)
            all_targets.extend(texts)

    cer = calculate_cer(all_preds, all_targets)
    wer = calculate_wer(all_preds, all_targets)
    line_acc = calculate_line_accuracy(all_preds, all_targets)

    total_params = sum(p.numel() for p in model.parameters())

    return {
        'cer': cer, 'wer': wer, 'line_accuracy': line_acc,
        'avg_inference_ms': np.mean(inference_times) * 1000,
        'std_inference_ms': np.std(inference_times) * 1000,
        'fps': len(inference_times) / sum(inference_times),
        'total_params': total_params,
        'predictions': all_preds, 'targets': 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.")


# ===============================
# 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():
    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_clean_text = args.clean_text and not args.no_clean_text

    # 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}")

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

    # Dataset 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,
        clean_text_enabled=use_clean_text)

    # Datasets
    train_dir = os.path.join(args.data_dir, "Training")
    val_dir = os.path.join(args.data_dir, "Validation")
    test_dir = os.path.join(args.data_dir, "Testing")

    train_dataset = KurdishParagraphDataset(train_dir, transform=train_transform, **ds_kwargs)
    val_dataset = KurdishParagraphDataset(val_dir, transform=eval_transform, **ds_kwargs)
    test_dataset = KurdishParagraphDataset(test_dir, transform=eval_transform, **ds_kwargs)

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

    print(f"\n  Training: {len(train_dataset)} | Validation: {len(val_dataset)} | Testing: {len(test_dataset)}")

    # 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)

    # Load pre-trained weights
    model = load_pretrained_weights(model, args.pretrained_path, device)

    # Freeze encoder
    if args.freeze_epochs > 0:
        model.freeze_encoder()

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  Total parameters: {total_params:,}")
    print(f"  Trainable parameters: {trainable_params:,}")

    # Optimizer with differential learning rates
    encoder_params = list(model.feature_extractor.parameters())
    other_params = [p for n, p in model.named_parameters() if 'feature_extractor' not in n]

    optimizer = optim.AdamW([
        {'params': encoder_params, 'lr': args.learning_rate * args.encoder_lr_mult},
        {'params': other_params, 'lr': args.learning_rate}
    ], weight_decay=args.weight_decay)

    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=args.lr_factor,
        patience=args.lr_patience, min_lr=1e-7)

    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"\nFine-tuning started: {datetime.now():%Y-%m-%d %H:%M:%S}")
    log(f"Pre-trained model: {args.pretrained_path}")
    log(f"Config: {vars(args)}")

    # Initial evaluation
    initial_cer = evaluate_cer_batch(model, val_loader, device, idx_to_char)
    log(f"\n  Initial Val CER (pre-trained): {initial_cer:.4f}")

    # Fine-tuning loop
    best_val_cer = float('inf')

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

        # Unfreeze encoder after freeze period
        if epoch == args.freeze_epochs + 1 and args.freeze_epochs > 0:
            model.unfreeze_encoder()
            trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
            log(f"\n  Epoch {epoch}: Encoder unfrozen ({trainable:,} trainable params)")

        # 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(val_cer)
        elapsed = time.time() - start_time
        mins, secs = divmod(elapsed, 60)

        lr_enc = optimizer.param_groups[0]['lr']
        lr_dec = optimizer.param_groups[1]['lr']

        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} | "
            f"LR: Enc={lr_enc:.2e}, Dec={lr_dec:.2e}")

        if val_cer < best_val_cer:
            best_val_cer = val_cer

        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 comprehensive evaluation
    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'])
    log(f"  Best epoch: {ckpt['epoch']}, Best Val CER: {ckpt['val_cer']:.4f}")

    # Validation results
    log(f"\n--- Validation Set ---")
    val_results = comprehensive_evaluation(model, val_loader, device, idx_to_char)
    log(f"  CER: {val_results['cer']:.4f} | WER: {val_results['wer']:.4f} | "
        f"Line Acc: {val_results['line_accuracy']:.4f}")
    log(f"  Inference: {val_results['avg_inference_ms']:.2f} ms | FPS: {val_results['fps']:.2f}")

    # Test results
    log(f"\n--- Test Set ---")
    test_results = comprehensive_evaluation(model, test_loader, device, idx_to_char)
    log(f"  CER: {test_results['cer']:.4f} ({(1-test_results['cer'])*100:.2f}% accuracy)")
    log(f"  WER: {test_results['wer']:.4f} ({(1-test_results['wer'])*100:.2f}% accuracy)")
    log(f"  Line Accuracy: {test_results['line_accuracy']:.4f}")
    log(f"  Inference: {test_results['avg_inference_ms']:.2f} ± {test_results['std_inference_ms']:.2f} ms")
    log(f"  FPS: {test_results['fps']:.2f}")
    log(f"  Parameters: {test_results['total_params']:,}")

    # Sample predictions
    log(f"\n--- Sample Predictions ---")
    for i in range(min(5, len(test_results['predictions']))):
        log(f"\nSample {i + 1}:")
        pred_preview = test_results['predictions'][i][:200]
        true_preview = test_results['targets'][i][:200]
        log(f"  Predicted: {pred_preview}")
        log(f"  Actual:    {true_preview}")

    log(f"\nFine-tuning complete: {datetime.now():%Y-%m-%d %H:%M:%S}")
    log(f"Best model: {best_model_path}")
    log(f"Improvement: {initial_cer:.4f} -> {best_val_cer:.4f} "
        f"({(initial_cer - best_val_cer)*100:.2f}% absolute)")

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


if __name__ == "__main__":
    main()