# kiri_ocr/model_transformer.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import List, Dict, Optional, Tuple, NamedTuple
import json
import math
import numpy as np
from PIL import Image


# ========== CONFIGURATION ==========
@dataclass
class CFG:
    # --- Model Architecture ---
    IMG_H: int = 48
    IMG_W: int = 640
    MAX_DEC_LEN: int = 260
    UNK_TOKEN: str = "<unk>"
    COLLAPSE_WHITESPACE: bool = True
    UNICODE_NFC: bool = True

    ENC_DIM: int = 256
    ENC_LAYERS: int = 4
    ENC_HEADS: int = 8
    ENC_FF: int = 1024
    DROPOUT: float = 0.15

    USE_DECODER: bool = True
    DEC_DIM: int = 256
    DEC_LAYERS: int = 3
    DEC_HEADS: int = 8
    DEC_FF: int = 1024

    USE_CTC: bool = True
    USE_LM: bool = True
    USE_LM_FUSION_EVAL: bool = True
    LM_FUSION_ALPHA: float = 0.35
    USE_FP16: bool = True
    USE_AUTOCAST: bool = True

    # --- Inference Params ---
    CTC_FUSION_ALPHA: float = 0.5
    BEAM: int = 4
    BEAM_LENP: float = 0.6  # FIXED: was 0.0

    EOS_LOGP_BIAS: float = 5.0
    EOS_LOGP_BOOST: float = 5.0
    EOS_BIAS_UNTIL_LEN: int = 3

    REPEAT_LAST_PENALTY: float = 3.0
    UNK_LOGP_PENALTY: float = 2.0

    DEC_MAX_LEN_RATIO: float = 1.5
    DEC_MAX_LEN_PAD: int = 10
    MEM_MAX_LEN_RATIO: float = 0.75


# ========== RESULT TYPE ==========
class OCRResult(NamedTuple):
    """Structured result from OCR inference"""

    text: str
    confidence: float
    ctc_confidence: Optional[float] = None
    decoder_confidence: Optional[float] = None


# ========== TOKENIZER ==========
class CharTokenizer:
    def __init__(self, vocab_path: str, cfg: CFG):
        with open(vocab_path, "r", encoding="utf-8") as f:
            vocab_raw: Dict[str, int] = json.load(f)

        if cfg.UNK_TOKEN not in vocab_raw:
            vocab_raw[cfg.UNK_TOKEN] = max(vocab_raw.values(), default=-1) + 1

        items = sorted(vocab_raw.items(), key=lambda kv: kv[1])
        self.token_to_id = {tok: i for i, (tok, _) in enumerate(items)}
        self.id_to_token = {i: tok for i, (tok, _) in enumerate(items)}

        self.unk_token = cfg.UNK_TOKEN
        self.unk_id = self.token_to_id[cfg.UNK_TOKEN]
        self.blank_id = 0
        self.pad_id = 1
        self.ctc_offset = 2
        self.vocab_size = len(self.token_to_id)
        self.ctc_classes = self.vocab_size + self.ctc_offset

        self.dec_pad = 0
        self.dec_bos = 1
        self.dec_eos = 2
        self.dec_offset = 3
        self.dec_vocab = self.vocab_size + self.dec_offset

    def decode_ctc(self, ids: List[int]) -> str:
        """Decode CTC output with deduplication"""
        chars = []
        prev_id = None
        for idx in ids:
            if idx == prev_id:
                continue
            prev_id = idx
            if idx < self.ctc_offset:
                continue
            raw_id = idx - self.ctc_offset
            if 0 <= raw_id < self.vocab_size:
                char = self.id_to_token.get(raw_id, "")
                if char != self.unk_token:
                    chars.append(char)
        return "".join(chars)

    def decode_dec(self, ids: List[int]) -> str:
        out = []
        for x in ids:
            if x in (self.dec_pad, self.dec_bos, self.dec_eos):
                continue
            y = x - self.dec_offset
            if 0 <= y < self.vocab_size:
                t = self.id_to_token.get(y, self.unk_token)
                out.append("" if t == self.unk_token else t)
        return "".join(out)

    def dec_to_ctc_id(self, dec_id: int) -> int:
        """Convert decoder token ID to CTC token ID"""
        if dec_id in (self.dec_pad, self.dec_bos, self.dec_eos):
            return self.blank_id
        raw_id = dec_id - self.dec_offset
        if 0 <= raw_id < self.vocab_size:
            return raw_id + self.ctc_offset
        return self.unk_id + self.ctc_offset


# ========== MODEL COMPONENTS ==========
class PosEnc2D(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def _make_pe(
        self, length: int, dim: int, device: torch.device, dtype: torch.dtype
    ) -> torch.Tensor:
        pos = torch.arange(length, dtype=dtype, device=device).unsqueeze(1)
        div = torch.exp(
            torch.arange(0, dim, 2, dtype=dtype, device=device)
            * (-math.log(10000.0) / dim)
        )
        pe = torch.zeros((length, dim), dtype=dtype, device=device)
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        return pe

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b, c, h, w = x.shape
        num_feats = c // 2
        if num_feats == 0:
            return x
        pe_y = self._make_pe(h, num_feats, x.device, x.dtype)
        pe_x = self._make_pe(w, num_feats, x.device, x.dtype)
        pe_y = pe_y.unsqueeze(2).repeat(1, 1, w)
        pe_x = pe_x.transpose(0, 1).unsqueeze(0).repeat(h, 1, 1)
        pe = torch.cat([pe_y, pe_x], dim=1)
        pe = pe.permute(1, 0, 2)
        if pe.size(0) < c:
            pad = torch.zeros((c - pe.size(0), h, w), device=x.device, dtype=x.dtype)
            pe = torch.cat([pe, pad], dim=0)
        return x + pe.unsqueeze(0)


class ConvStem(nn.Module):
    def __init__(self, dim: int, dropout: float):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(1, 48, 3, 1, 1, bias=False),
            nn.BatchNorm2d(48),
            nn.SiLU(inplace=True),
            nn.Conv2d(48, 96, 3, (2, 2), 1, bias=False),
            nn.BatchNorm2d(96),
            nn.SiLU(inplace=True),
            nn.Conv2d(96, 160, 3, (2, 2), 1, bias=False),
            nn.BatchNorm2d(160),
            nn.SiLU(inplace=True),
            nn.Conv2d(160, dim, 3, (2, 1), 1, bias=False),
            nn.BatchNorm2d(dim),
            nn.SiLU(inplace=True),
            nn.Dropout2d(dropout),
        )

    def forward(self, x):
        return self.net(x)


# ========== MAIN MODEL ==========
class KiriOCR(nn.Module):
    def __init__(self, cfg: CFG, tok: CharTokenizer):
        super().__init__()
        self.cfg = cfg
        self.tok = tok
        d = cfg.DROPOUT

        self.stem = ConvStem(cfg.ENC_DIM, d)
        self.pos2d = PosEnc2D(cfg.ENC_DIM)

        self.enc_ln_in = nn.LayerNorm(cfg.ENC_DIM)
        enc_layer = nn.TransformerEncoderLayer(
            d_model=cfg.ENC_DIM,
            nhead=cfg.ENC_HEADS,
            dim_feedforward=cfg.ENC_FF,
            dropout=d,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.enc = nn.TransformerEncoder(
            enc_layer,
            num_layers=cfg.ENC_LAYERS,
            enable_nested_tensor=False,
        )
        self.enc_ln = nn.LayerNorm(cfg.ENC_DIM)

        if cfg.USE_CTC:
            self.ctc_head = nn.Sequential(
                nn.LayerNorm(cfg.ENC_DIM),
                nn.Dropout(d),
                nn.Linear(cfg.ENC_DIM, tok.ctc_classes),
            )

        self.mem_proj = nn.Linear(cfg.ENC_DIM, cfg.DEC_DIM, bias=False)
        self.dec_emb = nn.Embedding(tok.dec_vocab, cfg.DEC_DIM)

        dec_layer = nn.TransformerDecoderLayer(
            d_model=cfg.DEC_DIM,
            nhead=cfg.DEC_HEADS,
            dim_feedforward=cfg.DEC_FF,
            dropout=d,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.dec = nn.TransformerDecoder(dec_layer, num_layers=cfg.DEC_LAYERS)
        self.dec_ln = nn.LayerNorm(cfg.DEC_DIM)
        self.dec_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)

        if cfg.USE_LM:
            self.lm_head = nn.Linear(cfg.DEC_DIM, tok.dec_vocab)

    def encode(self, imgs: torch.Tensor) -> torch.Tensor:
        x = self.stem(imgs)
        x = self.pos2d(x)
        x = F.adaptive_avg_pool2d(x, (1, x.size(-1)))
        x = x.squeeze(2).permute(0, 2, 1)
        x = self.enc_ln_in(x)
        x = self.enc(x)
        x = self.enc_ln(x)
        return x


# ========== PREPROCESSING ==========
class ResizeKeepRatioPadNoCrop:
    def __init__(self, h: int, w: int):
        self.h = h
        self.w = w

    def __call__(self, img: Image.Image) -> Image.Image:
        iw, ih = img.size
        if ih <= 0 or iw <= 0:
            return img.resize((self.w, self.h), Image.BILINEAR)

        scale = self.h / float(ih)
        nw = max(1, int(round(iw * scale)))
        img = img.resize((nw, self.h), Image.BILINEAR)

        if nw >= self.w:
            return img.crop((0, 0, self.w, self.h))

        # Left-aligned padding with gray (128)
        new_img = Image.new("L", (self.w, self.h), 128)
        new_img.paste(img, (0, 0))
        return new_img


def preprocess_pil(cfg: CFG, pil: Image.Image) -> torch.Tensor:
    img = pil.convert("L")
    img = ResizeKeepRatioPadNoCrop(cfg.IMG_H, cfg.IMG_W)(img)
    img_tensor = torch.from_numpy(np.array(img)).float() / 255.0
    img_tensor = (img_tensor - 0.5) / 0.5
    return img_tensor.unsqueeze(0).unsqueeze(0)


# ========== CONFIDENCE HELPERS ==========
def compute_ctc_confidence(
    ctc_logits: torch.Tensor, tok: CharTokenizer
) -> Tuple[float, str, int]:
    """
    Compute CTC confidence and decode text.

    Returns:
        (confidence, decoded_text, estimated_length)
    """
    if ctc_logits.dim() == 3:
        ctc_logits = ctc_logits.squeeze(0)  # [T, C]

    probs = F.softmax(ctc_logits, dim=-1)
    best_ids = ctc_logits.argmax(dim=-1).tolist()

    # Decode with deduplication
    text = tok.decode_ctc(best_ids)

    # Confidence: average max probability at each timestep
    max_probs = probs.max(dim=-1).values
    confidence = max_probs.mean().item()

    # Estimate length (count non-blank, non-repeated tokens)
    length = 0
    prev = None
    for idx in best_ids:
        if idx != prev and idx >= tok.ctc_offset:
            length += 1
        prev = idx

    return confidence, text, length


def compute_sequence_confidence(log_probs: List[float]) -> float:
    """Convert sequence log probabilities to confidence score."""
    if not log_probs:
        return 0.0

    # Average log prob
    avg_logp = sum(log_probs) / len(log_probs)

    # Convert to probability (clamped to [0, 1])
    confidence = math.exp(avg_logp)
    return min(1.0, max(0.0, confidence))


# ========== MAIN DECODING FUNCTION ==========
@torch.inference_mode()
def beam_decode_one_batched(
    model: KiriOCR,
    mem_proj_1: torch.Tensor,
    tok: CharTokenizer,
    cfg: CFG,
    ctc_logits_1: Optional[torch.Tensor] = None,
) -> Tuple[str, float]:
    """
    Beam search decoder with confidence score.

    Returns:
        (decoded_text, confidence_score)
    """
    device = mem_proj_1.device
    is_cuda = device.type == "cuda"

    # Get CTC info for length estimation and fusion
    ctc_confidence = None
    target_len = None

    if ctc_logits_1 is not None:
        ctc_confidence, ctc_text, target_len = compute_ctc_confidence(ctc_logits_1, tok)

    # Calculate max decoding steps
    if target_len and target_len > 0:
        max_steps = min(
            cfg.MAX_DEC_LEN,
            int(target_len * cfg.DEC_MAX_LEN_RATIO) + cfg.DEC_MAX_LEN_PAD,
        )
    else:
        mem_len = mem_proj_1.size(1)
        max_steps = min(
            cfg.MAX_DEC_LEN, int(mem_len * cfg.MEM_MAX_LEN_RATIO) + cfg.DEC_MAX_LEN_PAD
        )

    # Beam state: (score, token_ids, token_logprobs, finished)
    beams: List[Tuple[float, List[int], List[float], bool]] = [
        (0.0, [tok.dec_bos], [], False)
    ]

    full_causal = torch.triu(
        torch.ones(
            (cfg.MAX_DEC_LEN + 2, cfg.MAX_DEC_LEN + 2), device=device, dtype=torch.bool
        ),
        diagonal=1,
    )

    use_amp = cfg.USE_AUTOCAST and is_cuda

    for step in range(max_steps):
        if all(b[3] for b in beams):
            break

        alive = [b for b in beams if not b[3]]
        done = [b for b in beams if b[3]]

        if not alive:
            beams = done
            break

        maxL = max(len(b[1]) for b in alive)
        B = len(alive)

        inp = torch.full((B, maxL), tok.dec_pad, device=device, dtype=torch.long)
        for i, (_, seq, _, _) in enumerate(alive):
            inp[i, : len(seq)] = torch.tensor(seq, device=device, dtype=torch.long)

        tgt = model.dec_emb(inp)
        causal = full_causal[:maxL, :maxL]

        with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
            out = model.dec(
                tgt=tgt, memory=mem_proj_1.expand(B, -1, -1), tgt_mask=causal
            )
            out = model.dec_ln(out)
            logits = model.dec_head(out)[:, -1, :]
            logp = F.log_softmax(logits, dim=-1)

            if cfg.USE_LM and cfg.USE_LM_FUSION_EVAL and hasattr(model, "lm_head"):
                lm_logits = model.lm_head(out)[:, -1, :]
                logp = logp + cfg.LM_FUSION_ALPHA * F.log_softmax(lm_logits, dim=-1)

        # Apply penalties
        unk_id = tok.unk_id + tok.dec_offset

        for i, (_, seq, _, _) in enumerate(alive):
            cur_len = len(seq) - 1

            # EOS bias based on target length
            if target_len and target_len > 0:
                min_len = min(cfg.EOS_BIAS_UNTIL_LEN, max(1, int(target_len * 0.5)))
                if cur_len < min_len:
                    logp[i, tok.dec_eos] -= cfg.EOS_LOGP_BIAS
                elif cur_len >= target_len:
                    logp[i, tok.dec_eos] += cfg.EOS_LOGP_BOOST
            else:
                if cur_len < cfg.EOS_BIAS_UNTIL_LEN:
                    logp[i, tok.dec_eos] -= cfg.EOS_LOGP_BIAS

            # Repeat penalty
            if len(seq) >= 4 and seq[-1] == seq[-2] == seq[-3]:
                logp[i, seq[-1]] -= cfg.REPEAT_LAST_PENALTY

            # UNK penalty
            logp[i, unk_id] -= cfg.UNK_LOGP_PENALTY

        topv, topi = torch.topk(logp, k=cfg.BEAM, dim=-1)

        # Expand beams
        new_beams: List[Tuple[float, List[int], List[float], bool]] = list(done)

        for bi, (base_score, seq, logprobs, _) in enumerate(alive):
            for v, tid in zip(topv[bi].tolist(), topi[bi].tolist()):
                new_seq = seq + [int(tid)]
                new_logprobs = logprobs + [float(v)]
                is_finished = int(tid) == tok.dec_eos
                new_score = base_score + float(v)
                new_beams.append((new_score, new_seq, new_logprobs, is_finished))

        # Length-normalized scoring
        def normed(entry):
            score, seq, _, _ = entry
            L = max(1, len(seq) - 1)
            return score / (L**cfg.BEAM_LENP)

        new_beams.sort(key=normed, reverse=True)
        beams = new_beams[: cfg.BEAM]

    # ========== FINAL SCORING WITH CTC FUSION ==========
    def final_score_and_confidence(entry):
        score, seq, logprobs, _ = entry
        length = max(1, len(seq) - 1)

        dec_score = score / (length**cfg.BEAM_LENP)
        dec_conf = compute_sequence_confidence(logprobs)

        # CTC fusion
        if ctc_logits_1 is not None and cfg.CTC_FUSION_ALPHA > 0:
            ctc_score = compute_ctc_alignment_score(ctc_logits_1, seq, tok)
            combined_score = dec_score + cfg.CTC_FUSION_ALPHA * ctc_score
            return combined_score, dec_conf

        return dec_score, dec_conf

    scored = [(final_score_and_confidence(b), b) for b in beams]
    scored.sort(key=lambda x: x[0][0], reverse=True)

    (_, best_dec_conf), best_beam = scored[0]
    _, best_seq, _, _ = best_beam

    # Decode text
    ids = []
    for x in best_seq[1:]:
        if x == tok.dec_eos:
            break
        ids.append(x)
    text = tok.decode_dec(ids)

    # ========== COMPUTE FINAL CONFIDENCE ==========
    # Combine decoder confidence with CTC confidence (if available)
    if ctc_confidence is not None:
        # Weighted average: favor decoder slightly
        final_confidence = 0.6 * best_dec_conf + 0.4 * ctc_confidence
    else:
        final_confidence = best_dec_conf

    return text, final_confidence


def compute_ctc_alignment_score(
    ctc_logits: torch.Tensor, dec_seq: List[int], tok: CharTokenizer
) -> float:
    """Compute CTC alignment score using forward algorithm."""
    if ctc_logits.dim() == 3:
        ctc_logits = ctc_logits.squeeze(0)

    log_probs = F.log_softmax(ctc_logits, dim=-1)

    # Convert decoder sequence to CTC labels
    labels = []
    for x in dec_seq[1:]:
        if x == tok.dec_eos:
            break
        if x in (tok.dec_pad, tok.dec_bos):
            continue
        ctc_id = tok.dec_to_ctc_id(x)
        labels.append(ctc_id)

    if not labels:
        return log_probs[:, tok.blank_id].sum().item() / max(1, log_probs.size(0))

    # Forward algorithm
    T = log_probs.size(0)
    blank = tok.blank_id

    # Extend: [b, l0, b, l1, b, ...]
    ext = [blank]
    for lid in labels:
        ext.append(lid)
        ext.append(blank)
    S = len(ext)

    alpha = log_probs.new_full((S,), float("-inf"))
    alpha[0] = log_probs[0, blank]
    if S > 1:
        alpha[1] = log_probs[0, ext[1]]

    for t in range(1, T):
        new_alpha = log_probs.new_full((S,), float("-inf"))
        for s in range(S):
            candidates = [alpha[s]]
            if s > 0:
                candidates.append(alpha[s - 1])
            if s > 1 and ext[s] != blank and ext[s] != ext[s - 2]:
                candidates.append(alpha[s - 2])

            stacked = torch.stack(
                [
                    (
                        c
                        if isinstance(c, torch.Tensor)
                        else torch.tensor(c, device=log_probs.device)
                    )
                    for c in candidates
                ]
            )
            new_alpha[s] = torch.logsumexp(stacked, dim=0) + log_probs[t, ext[s]]
        alpha = new_alpha

    if S == 1:
        total = alpha[0]
    else:
        total = torch.logsumexp(torch.stack([alpha[S - 1], alpha[S - 2]]), dim=0)

    return total.item() / max(1, len(labels))


# ========== GREEDY CTC DECODE (FAST ALTERNATIVE) ==========
@torch.inference_mode()
def greedy_ctc_decode(
    model: KiriOCR, imgs: torch.Tensor, tok: CharTokenizer, cfg: CFG
) -> Tuple[str, float]:
    """
    Fast greedy CTC decoding (no beam search).

    Returns:
        (text, confidence)
    """
    mem = model.encode(imgs[:1])
    ctc_logits = model.ctc_head(mem)

    confidence, text, _ = compute_ctc_confidence(ctc_logits, tok)
    return text, confidence
