"""
train.py
========
Training script for the Kurdish handwritten word recognition models in the
Karez/KHWR repository. Supports four model families through a single
--model_type flag:

    baseline  -- CRNN without attention
    luong     -- CRNN with Luong multiplicative attention
    mhsa      -- CRNN with Multi-Head Self-Attention
    faa       -- CRNN with Frequency-Adaptive Attention (proposed)

Example:
    python Scripts/train.py \
        --model_type faa \
        --data_dir ./data/DASTNUS/Unique-Words \
        --vocab_path FAA-Word-Model/vocab.json \
        --output_dir ./output/faa_seed42 \
        --seed 42
"""

import argparse
import os
import sys
import json
import glob
import math
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data as data
import torchvision.transforms as T
from PIL import Image
from collections import Counter
from datetime import datetime
from tqdm import tqdm


# ===============================
# Argument Parser
# ===============================
def parse_args():
    parser = argparse.ArgumentParser(
        description="Kurdish Handwritten Word Recognition Training"
    )
    # Model selection
    parser.add_argument("--model_type", type=str, required=True,
                        choices=["baseline", "luong", "mhsa", "faa"],
                        help="Which model family to train")
    # Data paths
    parser.add_argument("--data_dir", type=str, required=True,
                        help="Root directory containing Training/, "
                             "Validation/, and Testing/ 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=64)
    parser.add_argument("--img_width", type=int, default=164)
    # Training hyperparameters
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--num_epochs", type=int, default=80)
    parser.add_argument("--learning_rate", type=float, default=5e-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 parameters
    parser.add_argument("--hidden_size", type=int, default=160,
                        help="LSTM hidden size per direction")
    parser.add_argument("--lstm_layers", type=int, default=3)
    parser.add_argument("--lstm_dropout", type=float, default=0.3)
    parser.add_argument("--cnn_dropout", type=float, default=0.2)
    # MHSA specific
    parser.add_argument("--num_heads", type=int, default=4)
    parser.add_argument("--ff_dim", type=int, default=320)
    parser.add_argument("--mhsa_dropout", type=float, default=0.1)
    # Early stopping
    parser.add_argument("--patience", type=int, default=10)
    # Augmentation
    parser.add_argument("--no_aug", action="store_true",
                        help="Disable adaptive augmentation")
    # Fine-tuning
    parser.add_argument("--init_checkpoint", type=str, default=None,
                        help="Optional .pth checkpoint to initialize from "
                             "(e.g. for few-shot fine-tuning)")
    # Output
    parser.add_argument("--output_dir", type=str, default="./output",
                        help="Directory to save best model and logs")
    return parser.parse_args()


# ===============================
# Vocabulary Loader
# ===============================
def load_vocabulary(vocab_path):
    """Load vocab.json; expects char-to-index mapping with <BLANK> at 0."""
    with open(vocab_path, "r", encoding="utf-8") as f:
        vocab = json.load(f)
    if "<BLANK>" in vocab:
        blank_idx = vocab["<BLANK>"]
    else:
        blank_idx = 0
    idx_to_char = {v: k for k, v in vocab.items()}
    return vocab, idx_to_char, blank_idx


def indices_to_text(indices, idx_to_char, blank_idx=0):
    """Map a list of integer indices back to a string (skips blank)."""
    return "".join(idx_to_char.get(int(i), "")
                   for i in indices if int(i) != blank_idx)


def text_to_indices(text, vocab):
    return [vocab[c] for c in text if c in vocab]


# ===============================
# Dataset with Adaptive Augmentation
# ===============================
class KurdishWordDataset(data.Dataset):
    """
    Loads (.tif, .txt) image-label pairs from a folder. Each .tif has a
    matching .txt with one line of Kurdish ground truth.
    """
    def __init__(self, root_dir, img_height, img_width, augment=False,
                 aug_strength=0.0):
        self.root_dir    = root_dir
        self.img_height  = img_height
        self.img_width   = img_width
        self.augment     = augment
        self.aug_strength = aug_strength

        self.image_files = sorted(glob.glob(os.path.join(root_dir, "*.tif")))
        # Keep only those that have a matching label
        self.image_files = [
            p for p in self.image_files
            if os.path.exists(os.path.splitext(p)[0] + ".txt")
        ]
        print(f"  Loaded {len(self.image_files)} samples from {root_dir}")

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

    def set_aug_strength(self, strength):
        """Updated each epoch by the trainer (adaptive augmentation)."""
        self.aug_strength = max(0.0, min(1.0, strength))

    def __getitem__(self, idx):
        img_path   = self.image_files[idx]
        label_path = os.path.splitext(img_path)[0] + ".txt"

        # Load and resize keeping aspect ratio
        image = Image.open(img_path).convert("L")
        ow, oh = image.size
        new_h = self.img_height
        new_w = min(int(new_h * ow / oh), self.img_width)
        image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
        canvas = Image.new("L", (self.img_width, self.img_height), color=255)
        canvas.paste(image, (0, 0))

        # Apply adaptive augmentation (training only)
        if self.augment and self.aug_strength > 0.0:
            canvas = self._apply_augmentation(canvas)

        # Normalize and convert to tensor
        canvas = T.functional.to_tensor(canvas)
        canvas = T.functional.normalize(canvas, (0.5,), (0.5,))

        # Load label
        try:
            with open(label_path, "r", encoding="utf-8") as f:
                text = f.readline().strip()
        except UnicodeDecodeError:
            with open(label_path, "r", encoding="utf-8-sig") as f:
                text = f.readline().strip()

        return canvas, text

    def _apply_augmentation(self, img):
        """Light geometric and intensity perturbations scaled by aug_strength."""
        s = self.aug_strength
        # Random rotation
        if random.random() < 0.5:
            angle = random.uniform(-3.0 * s, 3.0 * s)
            img = img.rotate(angle, resample=Image.BILINEAR, fillcolor=255)
        # Random translation
        if random.random() < 0.5:
            dx = int(random.uniform(-4 * s, 4 * s))
            dy = int(random.uniform(-2 * s, 2 * s))
            img = img.transform(img.size, Image.AFFINE,
                                (1, 0, dx, 0, 1, dy), fillcolor=255)
        # Light Gaussian noise
        if random.random() < 0.3:
            arr = np.asarray(img, dtype=np.float32)
            arr += np.random.normal(0, 5.0 * s, arr.shape)
            arr = np.clip(arr, 0, 255).astype(np.uint8)
            img = Image.fromarray(arr, mode="L")
        return img


def collate_fn(batch, vocab):
    imgs, texts = zip(*batch)
    imgs = torch.stack(imgs, 0)
    target_indices = []
    target_lengths = []
    for t in texts:
        idx_seq = text_to_indices(t, vocab)
        target_indices.extend(idx_seq)
        target_lengths.append(len(idx_seq))
    targets = torch.tensor(target_indices, dtype=torch.long)
    lengths = torch.tensor(target_lengths, dtype=torch.long)
    return imgs, targets, lengths, list(texts)


# ===============================
# Model Architectures
# ===============================
class BidirectionalLSTM(nn.Module):
    def __init__(self, nIn, nHidden, nOut, dropout=0.0):
        super().__init__()
        self.rnn       = nn.LSTM(nIn, nHidden, bidirectional=True)
        self.embedding = nn.Linear(nHidden * 2, nOut)
        self.dropout   = nn.Dropout(dropout) if dropout > 0 else None

    def forward(self, x):
        recurrent, _ = self.rnn(x)
        if self.dropout:
            recurrent = self.dropout(recurrent)
        T, b, h = recurrent.size()
        return self.embedding(recurrent.view(T * b, h)).view(T, b, -1)


def _build_cnn(nc, cnn_dropout):
    """Shared 6-block CNN backbone (max 256 channels)."""
    return nn.Sequential(
        nn.Conv2d(nc, 64,  3, 1, 1), nn.BatchNorm2d(64),  nn.ReLU(True),
        nn.MaxPool2d(2, 2), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(64,  128, 3, 1, 1), nn.BatchNorm2d(128), nn.ReLU(True),
        nn.MaxPool2d(2, 2), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(128, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(True),
        nn.MaxPool2d((2, 1), (2, 1)), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(256, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(True),
        nn.MaxPool2d((2, 1), (2, 1)), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(256, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(True),
        nn.MaxPool2d((2, 1), (2, 1)), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(256, 256, 3, 1, 1), nn.BatchNorm2d(256), nn.ReLU(True),
        nn.MaxPool2d((2, 1), (2, 1)), nn.Dropout2d(cnn_dropout),

        nn.Conv2d(256, 256, (1, 3), 1, (0, 1)), nn.BatchNorm2d(256), nn.ReLU(True),
    )


class BaselineCRNN(nn.Module):
    """CRNN baseline without attention."""
    def __init__(self, nclass, nh, num_lstm_layers, lstm_dropout, cnn_dropout):
        super().__init__()
        self.cnn = _build_cnn(1, cnn_dropout)
        layers = []
        input_size = 256
        for i in range(num_lstm_layers):
            out_size = nclass if i == num_lstm_layers - 1 else nh
            drop     = 0      if i == num_lstm_layers - 1 else lstm_dropout
            layers.append(BidirectionalLSTM(input_size, nh, out_size, dropout=drop))
            input_size = nh
        self.rnn = nn.Sequential(*layers)

    def forward(self, x):
        conv = self.cnn(x)
        b, c, h, w = conv.size()
        if h != 1:
            conv = F.adaptive_avg_pool2d(conv, (1, w))
        conv = conv.squeeze(2).permute(2, 0, 1)
        return self.rnn(conv)


class LuongAttention(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.W_a      = nn.Linear(hidden_size, hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size * 2, hidden_size)
        self.norm     = nn.LayerNorm(hidden_size)

    def forward(self, x):
        T, B, H  = x.size()
        keys     = self.W_a(x)
        x_bth    = x.permute(1, 0, 2)
        keys_bth = keys.permute(1, 0, 2)
        scores   = torch.bmm(x_bth, keys_bth.transpose(1, 2)) / (H ** 0.5)
        weights  = torch.softmax(scores, dim=-1)
        context  = torch.bmm(weights, x_bth).permute(1, 0, 2)
        combined = torch.cat([x, context], dim=-1)
        output   = torch.tanh(self.out_proj(combined))
        return self.norm(output + x)


class LuongCRNN(nn.Module):
    def __init__(self, nclass, nh, num_lstm_layers, lstm_dropout, cnn_dropout):
        super().__init__()
        self.cnn = _build_cnn(1, cnn_dropout)
        self.lstm_layers = nn.ModuleList()
        input_size = 256
        for i in range(num_lstm_layers):
            out_size = nclass if i == num_lstm_layers - 1 else nh
            drop     = 0      if i == num_lstm_layers - 1 else lstm_dropout
            self.lstm_layers.append(
                BidirectionalLSTM(input_size, nh, out_size, dropout=drop)
            )
            input_size = nh
        self.attention = LuongAttention(hidden_size=nh)

    def forward(self, x):
        conv = self.cnn(x)
        b, c, h, w = conv.size()
        if h != 1:
            conv = F.adaptive_avg_pool2d(conv, (1, w))
        out = conv.squeeze(2).permute(2, 0, 1)
        for i, layer in enumerate(self.lstm_layers):
            out = layer(out)
            if i == 1:
                out = self.attention(out)
        return out


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, ff_dim, dropout):
        super().__init__()
        assert hidden_size % num_heads == 0
        self.hidden_size = hidden_size
        self.num_heads   = num_heads
        self.head_dim    = hidden_size // num_heads
        self.q_proj   = nn.Linear(hidden_size, hidden_size)
        self.k_proj   = nn.Linear(hidden_size, hidden_size)
        self.v_proj   = nn.Linear(hidden_size, hidden_size)
        self.out_proj = nn.Linear(hidden_size, hidden_size)
        self.ff = nn.Sequential(
            nn.Linear(hidden_size, ff_dim),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(ff_dim, hidden_size),
        )
        self.norm1   = nn.LayerNorm(hidden_size)
        self.norm2   = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        T, B, H = x.size()
        x_btn   = x.permute(1, 0, 2)
        def split_heads(t):
            return t.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        Q = split_heads(self.q_proj(x_btn))
        K = split_heads(self.k_proj(x_btn))
        V = split_heads(self.v_proj(x_btn))
        scores  = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        weights = self.dropout(torch.softmax(scores, dim=-1))
        attn_out = torch.matmul(weights, V).transpose(1, 2).contiguous()
        attn_out = self.out_proj(attn_out.view(B, T, H))
        x_btn = self.norm1(x_btn + self.dropout(attn_out))
        ff_out = self.ff(x_btn)
        x_btn  = self.norm2(x_btn + self.dropout(ff_out))
        return x_btn.permute(1, 0, 2)


class MHSACRNN(nn.Module):
    def __init__(self, nclass, nh, num_lstm_layers, lstm_dropout, cnn_dropout,
                 num_heads, ff_dim, mhsa_dropout):
        super().__init__()
        self.cnn = _build_cnn(1, cnn_dropout)
        self.lstm_layers = nn.ModuleList()
        input_size = 256
        for i in range(num_lstm_layers):
            out_size = nclass if i == num_lstm_layers - 1 else nh
            drop     = 0      if i == num_lstm_layers - 1 else lstm_dropout
            self.lstm_layers.append(
                BidirectionalLSTM(input_size, nh, out_size, dropout=drop)
            )
            input_size = nh
        self.attention = MultiHeadSelfAttention(
            nh, num_heads=num_heads, ff_dim=ff_dim, dropout=mhsa_dropout
        )

    def forward(self, x):
        conv = self.cnn(x)
        b, c, h, w = conv.size()
        if h != 1:
            conv = F.adaptive_avg_pool2d(conv, (1, w))
        out = conv.squeeze(2).permute(2, 0, 1)
        for i, layer in enumerate(self.lstm_layers):
            out = layer(out)
            if i == 1:
                out = self.attention(out)
        return out


class FrequencyAdaptiveAttention(nn.Module):
    """
    Proposed Frequency-Adaptive Attention. freq_weights is a non-trainable
    buffer; surrounding modules learn how to use it.
    """
    def __init__(self, hidden_size, vocab_size, freq_weights):
        super().__init__()
        self.hidden_size = hidden_size
        self.register_buffer("freq_weights", freq_weights)
        self.attention = nn.Sequential(
            nn.Linear(hidden_size, hidden_size // 2),
            nn.Tanh(),
            nn.Linear(hidden_size // 2, 1),
        )
        self.char_predictor = nn.Linear(hidden_size, vocab_size)
        self.freq_adapter   = nn.Sequential(
            nn.Linear(hidden_size + 1, hidden_size // 2),
            nn.ReLU(),
            nn.Linear(hidden_size // 2, hidden_size),
            nn.Sigmoid(),
        )
        self.out_proj = nn.Linear(hidden_size * 2, hidden_size)
        self.gate     = nn.Sequential(
            nn.Linear(hidden_size * 2, hidden_size),
            nn.Sigmoid(),
        )

    def forward(self, x):
        T, batch, _ = x.size()
        attn_scores     = self.attention(x)
        attn_weights    = torch.softmax(attn_scores, dim=0)
        char_logits     = self.char_predictor(x)
        char_probs      = torch.softmax(char_logits, dim=-1)
        expected_rarity = (
            char_probs * self.freq_weights.unsqueeze(0).unsqueeze(0)
        ).sum(dim=-1, keepdim=True)
        attn_boosted = attn_weights * (1.0 + expected_rarity)
        attn_boosted = attn_boosted / (attn_boosted.sum(dim=0, keepdim=True) + 1e-8)
        context = (x * attn_boosted).sum(dim=0, keepdim=True).expand(T, -1, -1)
        freq_input      = torch.cat([x, expected_rarity], dim=-1)
        freq_adapt_gate = self.freq_adapter(freq_input)
        combined  = torch.cat([x, context], dim=-1)
        projected = self.out_proj(combined)
        gate      = self.gate(combined)
        return gate * (freq_adapt_gate * projected) + (1 - gate) * x


class FAACRNN(nn.Module):
    def __init__(self, nclass, nh, num_lstm_layers, lstm_dropout, cnn_dropout,
                 freq_weights):
        super().__init__()
        self.cnn = _build_cnn(1, cnn_dropout)
        self.lstm_layers = nn.ModuleList()
        input_size = 256
        for i in range(num_lstm_layers):
            out_size = nclass if i == num_lstm_layers - 1 else nh
            drop     = 0      if i == num_lstm_layers - 1 else lstm_dropout
            self.lstm_layers.append(
                BidirectionalLSTM(input_size, nh, out_size, dropout=drop)
            )
            input_size = nh
        self.freq_attention = FrequencyAdaptiveAttention(
            hidden_size=nh, vocab_size=nclass, freq_weights=freq_weights
        )

    def forward(self, x):
        conv = self.cnn(x)
        b, c, h, w = conv.size()
        if h != 1:
            conv = F.adaptive_avg_pool2d(conv, (1, w))
        out = conv.squeeze(2).permute(2, 0, 1)
        for i, layer in enumerate(self.lstm_layers):
            out = layer(out)
            if i == 1:
                out = self.freq_attention(out)
        return out


def build_model(args, vocab_size, freq_weights, device):
    """Instantiate the architecture selected by --model_type."""
    if args.model_type == "baseline":
        return BaselineCRNN(
            nclass=vocab_size, nh=args.hidden_size,
            num_lstm_layers=args.lstm_layers,
            lstm_dropout=args.lstm_dropout, cnn_dropout=args.cnn_dropout
        ).to(device)
    if args.model_type == "luong":
        return LuongCRNN(
            nclass=vocab_size, nh=args.hidden_size,
            num_lstm_layers=args.lstm_layers,
            lstm_dropout=args.lstm_dropout, cnn_dropout=args.cnn_dropout
        ).to(device)
    if args.model_type == "mhsa":
        return MHSACRNN(
            nclass=vocab_size, nh=args.hidden_size,
            num_lstm_layers=args.lstm_layers,
            lstm_dropout=args.lstm_dropout, cnn_dropout=args.cnn_dropout,
            num_heads=args.num_heads, ff_dim=args.ff_dim,
            mhsa_dropout=args.mhsa_dropout
        ).to(device)
    if args.model_type == "faa":
        return FAACRNN(
            nclass=vocab_size, nh=args.hidden_size,
            num_lstm_layers=args.lstm_layers,
            lstm_dropout=args.lstm_dropout, cnn_dropout=args.cnn_dropout,
            freq_weights=freq_weights.to(device)
        ).to(device)
    raise ValueError(f"Unknown model_type: {args.model_type}")


# ===============================
# Frequency Weights for FAA
# ===============================
def compute_freq_weights(train_dir, vocab):
    """
    Compute per-character rarity weights from the training labels:
        weight[i] = log(N / count_i),  normalized to [0, 1]
    The CTC blank (index 0) is assigned weight 0.
    """
    counts = Counter()
    for label_path in glob.glob(os.path.join(train_dir, "*.txt")):
        try:
            with open(label_path, "r", encoding="utf-8") as f:
                text = f.readline().strip()
        except UnicodeDecodeError:
            with open(label_path, "r", encoding="utf-8-sig") as f:
                text = f.readline().strip()
        counts.update(text)

    vocab_size = len(vocab)
    weights = torch.zeros(vocab_size, dtype=torch.float32)
    total = sum(counts.values()) + len(vocab)  # Laplace smoothing
    for char, idx in vocab.items():
        if char == "<BLANK>":
            continue
        c = counts.get(char, 0) + 1
        weights[idx] = math.log(total / c)
    # Normalize to [0, 1]
    if weights.max() > 0:
        weights = weights / weights.max()
    return weights


# ===============================
# CTC Decode + Metrics
# ===============================
def ctc_greedy_decode(logits, blank_idx=0):
    _, max_idx = torch.max(logits, dim=2)
    decoded = []
    for b in range(max_idx.size(1)):
        seq = max_idx[:, b].cpu().numpy()
        out, prev = [], None
        for idx in seq:
            if idx != blank_idx and idx != prev:
                out.append(int(idx))
            prev = idx
        decoded.append(out)
    return decoded


def edit_distance(s1, s2):
    if len(s1) < len(s2):
        return edit_distance(s2, s1)
    if not s2:
        return len(s1)
    prev = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        curr = [i + 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 compute_cer(preds, gts):
    d = sum(edit_distance(p, t) for p, t in zip(preds, gts))
    c = sum(len(t) for t in gts)
    return d / max(c, 1)


def compute_wer(preds, gts):
    if not preds:
        return 0.0
    return sum(1 for p, t in zip(preds, gts) if p != t) / len(preds)


# ===============================
# Train / Validate
# ===============================
def train_one_epoch(model, loader, optimizer, scheduler, ctc_loss,
                    device, grad_clip, scaler):
    model.train()
    running_loss = 0.0
    n_batches = 0
    use_amp = (device.type == "cuda")

    for imgs, targets, lengths, _ in tqdm(loader, desc="  Train", leave=False):
        imgs    = imgs.to(device)
        targets = targets.to(device)
        lengths = lengths.to(device)

        optimizer.zero_grad()
        with torch.cuda.amp.autocast(enabled=use_amp):
            logits = model(imgs)                            # [T, B, V]
            log_probs = F.log_softmax(logits, dim=2)
            T, B, _ = log_probs.size()
            input_lengths = torch.full((B,), T, dtype=torch.long, device=device)
            loss = ctc_loss(log_probs, targets, input_lengths, lengths)

        if use_amp:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            optimizer.step()

        if scheduler is not None:
            scheduler.step()

        running_loss += loss.item()
        n_batches += 1

    return running_loss / max(n_batches, 1)


@torch.no_grad()
def validate(model, loader, idx_to_char, device, blank_idx=0):
    model.eval()
    preds_all, gts_all = [], []
    for imgs, _, _, texts in tqdm(loader, desc="  Val", leave=False):
        imgs = imgs.to(device)
        out  = model(imgs)
        decoded = ctc_greedy_decode(out, blank_idx=blank_idx)
        for seq in decoded:
            preds_all.append(indices_to_text(seq, idx_to_char, blank_idx))
        gts_all.extend(texts)
    return compute_cer(preds_all, gts_all), compute_wer(preds_all, gts_all)


# ===============================
# Main
# ===============================
def main():
    args = parse_args()

    # Reproducibility
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    os.makedirs(args.output_dir, exist_ok=True)
    log_path = os.path.join(args.output_dir, "training_log.txt")
    log_file = open(log_path, "w", encoding="utf-8")

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

    log("=" * 70)
    log(f"KHWR Training | model_type = {args.model_type}")
    log("=" * 70)
    log(f"Timestamp : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    log(f"Data dir  : {args.data_dir}")
    log(f"Vocab     : {args.vocab_path}")
    log(f"Output    : {args.output_dir}")
    log(f"Seed      : {args.seed}")

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

    # Load vocabulary
    vocab, idx_to_char, blank_idx = load_vocabulary(args.vocab_path)
    vocab_size = len(vocab)
    log(f"Vocab size: {vocab_size} (blank index = {blank_idx})")

    # Frequency weights (used only by FAA, computed for transparency anyway)
    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")
    freq_weights = compute_freq_weights(train_dir, vocab)

    # Datasets and loaders
    train_ds = KurdishWordDataset(train_dir, args.img_height, args.img_width,
                                  augment=not args.no_aug, aug_strength=0.3)
    val_ds   = KurdishWordDataset(val_dir,   args.img_height, args.img_width,
                                  augment=False)
    test_ds  = KurdishWordDataset(test_dir,  args.img_height, args.img_width,
                                  augment=False)

    def make_loader(ds, shuffle):
        return data.DataLoader(
            ds, batch_size=args.batch_size, shuffle=shuffle,
            num_workers=0, pin_memory=True,
            collate_fn=lambda b: collate_fn(b, vocab)
        )

    train_loader = make_loader(train_ds, shuffle=True)
    val_loader   = make_loader(val_ds, shuffle=False)
    test_loader  = make_loader(test_ds, shuffle=False)

    # Model
    model = build_model(args, vocab_size, freq_weights, device)
    total_params = sum(p.numel() for p in model.parameters())
    log(f"Parameters: {total_params:,}")

    # Optional warm-start (used for few-shot fine-tuning)
    if args.init_checkpoint is not None and os.path.exists(args.init_checkpoint):
        ckpt = torch.load(args.init_checkpoint, map_location=device)
        sd = ckpt.get("model_state_dict", ckpt)
        missing, unexpected = model.load_state_dict(sd, strict=False)
        log(f"Loaded init checkpoint: {args.init_checkpoint}")
        if missing:    log(f"  Missing keys   : {len(missing)}")
        if unexpected: log(f"  Unexpected keys: {len(unexpected)}")

    # Optimizer, scheduler, CTC loss
    optimizer = torch.optim.AdamW(model.parameters(),
                                  lr=args.learning_rate,
                                  weight_decay=args.weight_decay)
    steps_per_epoch = max(1, len(train_loader))
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=args.learning_rate,
        total_steps=args.num_epochs * steps_per_epoch
    )
    ctc_loss = nn.CTCLoss(blank=blank_idx, reduction="mean", zero_infinity=True)
    scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))

    # Training loop
    best_val_cer = float("inf")
    best_epoch   = 0
    patience_ctr = 0

    for epoch in range(1, args.num_epochs + 1):
        # Adaptive augmentation: ramp up strength to 1.0 across the schedule
        if not args.no_aug:
            train_ds.set_aug_strength(min(1.0, 0.3 + 0.7 * (epoch - 1) / max(1, args.num_epochs - 1)))

        loss = train_one_epoch(model, train_loader, optimizer, scheduler,
                               ctc_loss, device, args.grad_clip, scaler)
        val_cer, val_wer = validate(model, val_loader, idx_to_char,
                                    device, blank_idx)

        improved = val_cer < best_val_cer
        if improved:
            best_val_cer = val_cer
            best_epoch   = epoch
            patience_ctr = 0
            ckpt_path = os.path.join(args.output_dir, "best_model.pth")
            torch.save({
                "epoch":            epoch,
                "model_state_dict": model.state_dict(),
                "val_cer":          val_cer,
                "val_wer":          val_wer,
                "model_config": {
                    "model_type":       args.model_type,
                    "vocab_size":       vocab_size,
                    "lstm_hidden_size": args.hidden_size,
                    "lstm_layers":      args.lstm_layers,
                    "lstm_dropout":     args.lstm_dropout,
                    "cnn_dropout":      args.cnn_dropout,
                },
            }, ckpt_path)
        else:
            patience_ctr += 1

        marker = " *" if improved else ""
        log(f"Epoch {epoch:>3}/{args.num_epochs}  "
            f"train_loss={loss:.4f}  val_CER={val_cer:.4f}  "
            f"val_WER={val_wer:.4f}{marker}")

        if patience_ctr >= args.patience:
            log(f"Early stopping at epoch {epoch} "
                f"(patience={args.patience}, best={best_val_cer:.4f} @ epoch {best_epoch})")
            break

    # Final test evaluation on best checkpoint
    ckpt_path = os.path.join(args.output_dir, "best_model.pth")
    if os.path.exists(ckpt_path):
        ckpt = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(ckpt["model_state_dict"])

    test_cer, test_wer = validate(model, test_loader, idx_to_char,
                                  device, blank_idx)
    log("")
    log("=" * 70)
    log(f"FINAL TEST RESULTS (best epoch = {best_epoch})")
    log("=" * 70)
    log(f"Test CER : {test_cer:.4f}")
    log(f"Test WER : {test_wer:.4f}")
    log(f"Best Val : {best_val_cer:.4f}")
    log(f"Checkpoint: {ckpt_path}")

    log_file.close()


if __name__ == "__main__":
    main()