"""
inference.py
============
Inference script for the Kurdish handwritten word recognition models in the
Karez/KHWR repository.

The script auto-detects the model family from the config.json located next to
the model file, so the same command works for any of the four architectures
(Baseline, Luong, MHSA, FAA) and for the fine-tuned FAA checkpoints.

Examples:
    # Single image
    python Scripts/inference.py \
        --image Sample/sample_word.tif \
        --model_path FAA-Word-Model/model.safetensors \
        --vocab_path FAA-Word-Model/vocab.json

    # Directory of images
    python Scripts/inference.py \
        --image_dir ./test_words \
        --model_path FAA-Word-Model/model.safetensors \
        --vocab_path FAA-Word-Model/vocab.json
"""

import argparse
import os
import sys
import json
import math
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image


# ===============================
# Argument Parser
# ===============================
def parse_args():
    parser = argparse.ArgumentParser(
        description="Kurdish Handwritten Word Recognition Inference"
    )
    # Input
    parser.add_argument("--image", type=str, default=None,
                        help="Path to a single image (.tif/.png/.jpg)")
    parser.add_argument("--image_dir", type=str, default=None,
                        help="Directory of images to process")
    # Model
    parser.add_argument("--model_path", type=str, required=True,
                        help="Path to model.safetensors (or .pth checkpoint)")
    parser.add_argument("--vocab_path", type=str, required=True,
                        help="Path to vocab.json")
    parser.add_argument("--config_path", type=str, default=None,
                        help="Path to config.json. If omitted, the script "
                             "looks for it next to --model_path.")
    # Image dimensions
    parser.add_argument("--img_height", type=int, default=64)
    parser.add_argument("--img_width", type=int, default=164)
    # Hardware
    parser.add_argument("--device", type=str, default="auto",
                        choices=["auto", "cuda", "cpu"])
    # Output
    parser.add_argument("--output_file", type=str, default=None,
                        help="Optional path to save predictions as TSV "
                             "(filename<TAB>prediction).")
    return parser.parse_args()


# ===============================
# Vocabulary
# ===============================
def load_vocabulary(vocab_path):
    with open(vocab_path, "r", encoding="utf-8") as f:
        vocab = json.load(f)
    blank_idx = vocab.get("<BLANK>", 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):
    out = []
    for i in indices:
        i = int(i)
        if i == blank_idx:
            continue
        if i in idx_to_char:
            out.append(idx_to_char[i])
    return "".join(out)


# ===============================
# Model Architectures (mirror train.py)
# ===============================
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):
    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):
    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, keys_bth = x.permute(1, 0, 2), 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)
        output  = torch.tanh(self.out_proj(torch.cat([x, context], dim=-1)))
        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(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=4, ff_dim=320, dropout=0.1):
        super().__init__()
        assert hidden_size % num_heads == 0
        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, self.norm2 = nn.LayerNorm(hidden_size), 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(t): return t.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        Q, K, V = split(self.q_proj(x_btn)), split(self.k_proj(x_btn)), split(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().view(B, T, H)
        attn_out = self.out_proj(attn_out)
        x_btn = self.norm1(x_btn + self.dropout(attn_out))
        x_btn = self.norm2(x_btn + self.dropout(self.ff(x_btn)))
        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=4, ff_dim=320, mhsa_dropout=0.1):
        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, ff_dim, 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):
    def __init__(self, hidden_size, vocab_size, freq_weights):
        super().__init__()
        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


# ===============================
# Config + State Dict Loading
# ===============================
def load_state_dict(model_path, device):
    """Load weights from .safetensors or .pth."""
    if model_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        return load_file(model_path, device=str(device))
    ckpt = torch.load(model_path, map_location=device)
    return ckpt.get("model_state_dict", ckpt)


def auto_config_path(model_path, override=None):
    if override is not None:
        return override
    return os.path.join(os.path.dirname(model_path), "config.json")


def build_model_from_config(config, vocab_size, state_dict, device):
    """Read config.json + state_dict, instantiate the right architecture."""
    family = config.get("model_family")
    rnn_cfg = config.get("rnn", {})
    nh         = rnn_cfg.get("hidden_size", 160)
    n_layers   = rnn_cfg.get("num_layers",  3)
    lstm_drop  = rnn_cfg.get("dropout",     0.3)
    cnn_drop   = config.get("cnn_dropout",  0.2)

    if family == "FAA":
        # Pull freq_weights from state_dict
        fw = state_dict.get("freq_attention.freq_weights")
        if fw is None:
            fw = torch.ones(vocab_size) / vocab_size
        return FAACRNN(vocab_size, nh, n_layers, lstm_drop, cnn_drop,
                       freq_weights=fw.to(device)).to(device)

    if family == "MHSA":
        atten = config.get("attention", {})
        return MHSACRNN(
            vocab_size, nh, n_layers, lstm_drop, cnn_drop,
            num_heads=atten.get("num_heads", 4),
            ff_dim=atten.get("feed_forward_dim", 320),
            mhsa_dropout=atten.get("dropout", 0.1),
        ).to(device)

    if family == "Luong":
        return LuongCRNN(vocab_size, nh, n_layers, lstm_drop, cnn_drop).to(device)

    # Baseline (or unknown -> baseline)
    return BaselineCRNN(vocab_size, nh, n_layers, lstm_drop, cnn_drop).to(device)


# ===============================
# Preprocessing + CTC Decode
# ===============================
def preprocess_image(image_path, img_height, img_width, device):
    image = Image.open(image_path).convert("L")
    ow, oh = image.size
    new_h = img_height
    new_w = min(int(new_h * ow / oh), img_width)
    image = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
    canvas = Image.new("L", (img_width, img_height), color=255)
    canvas.paste(image, (0, 0))
    # to tensor, normalize
    import numpy as np
    arr = np.asarray(canvas, dtype="float32") / 255.0
    arr = (arr - 0.5) / 0.5
    tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0)   # [1, 1, H, W]
    return tensor.to(device)


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 i in seq:
            if i != blank_idx and i != prev:
                out.append(int(i))
            prev = i
        decoded.append(out)
    return decoded


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

    if args.image is None and args.image_dir is None:
        print("ERROR: provide either --image or --image_dir")
        sys.exit(1)

    # Device
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    print(f"Device: {device}")

    # Vocabulary
    vocab, idx_to_char, blank_idx = load_vocabulary(args.vocab_path)
    vocab_size = len(vocab)

    # Config (auto-detect path)
    config_path = auto_config_path(args.model_path, args.config_path)
    if not os.path.exists(config_path):
        print(f"ERROR: config.json not found at {config_path}")
        print("Pass --config_path explicitly.")
        sys.exit(1)
    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    # Load state_dict
    state_dict = load_state_dict(args.model_path, device)

    # Build the right architecture
    model = build_model_from_config(config, vocab_size, state_dict, device)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing or unexpected:
        print(f"  Missing keys   : {len(missing)}")
        print(f"  Unexpected keys: {len(unexpected)}")
    model.eval()

    family = config.get("model_family", "Baseline")
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model family: {family}  |  Parameters: {n_params:,}")

    # Collect images
    if args.image is not None:
        image_paths = [args.image]
    else:
        exts = ("*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg")
        image_paths = []
        for ext in exts:
            image_paths.extend(sorted(glob.glob(os.path.join(args.image_dir, ext))))
        if not image_paths:
            print(f"No images found in {args.image_dir}")
            sys.exit(1)

    # Inference
    results = []
    print("")
    print(f"{'File':<40} | Prediction")
    print("-" * 70)
    with torch.no_grad():
        for path in image_paths:
            x = preprocess_image(path, args.img_height, args.img_width, device)
            logits = model(x)
            decoded = ctc_greedy_decode(logits, blank_idx=blank_idx)[0]
            pred = indices_to_text(decoded, idx_to_char, blank_idx)
            results.append((os.path.basename(path), pred))
            print(f"{os.path.basename(path):<40} | {pred}")

    # Optional TSV output
    if args.output_file is not None:
        with open(args.output_file, "w", encoding="utf-8") as f:
            for name, pred in results:
                f.write(f"{name}\t{pred}\n")
        print(f"\nPredictions saved: {args.output_file}")


if __name__ == "__main__":
    main()