"""32M checkpoint consumer for RegionalPipeline's RTL bank; raw CTC retained."""
import json,math
from pathlib import Path
import numpy as np
import torch
from recognizer import Recognizer
from geometry import crop_line64,pad_batch
from visual_token import RTLVisualChar

class RTLReader:
    def __init__(self, checkpoint, vocabulary, device='cpu', frame_budget=8192):
        self.device = torch.device(device)
        data = json.loads(Path(vocabulary).read_text())
        self.alphabet = data['chars']
        self.model = Recognizer(classes=len(self.alphabet) + 1, gradient_checkpointing=False)
        state = torch.load(checkpoint, map_location='cpu', weights_only=True)
        self.model.load_state_dict(state['model'])
        self.model.to(self.device).eval()
        self.frame_budget = frame_budget

    @torch.inference_mode()
    def __call__(self, gray, boxes):
        from indra_ocr.pipeline import MIN_CONF
        crops, receipts = [], []
        for box in boxes:
            crop, receipt = crop_line64(gray, box)
            crops.append(crop)
            receipts.append(receipt)
        order = sorted(range(len(crops)), key=lambda i: crops[i].shape[1])
        parts = []
        for i in order:
            frames = math.ceil(crops[i].shape[1] / 8) * 2
            if parts and len(parts[-1]) < 32 and (len(parts[-1]) + 1) * frames <= self.frame_budget:
                parts[-1].append(i)
            else:
                parts.append([i])
        readings = [None] * len(crops)
        for ids in parts:
            x, widths = pad_batch([crops[i] for i in ids])
            x = torch.from_numpy(x).to(self.device)
            logits, lengths = self.model(x, torch.from_numpy(widths))
            probs, best = logits.float().softmax(-1).max(-1)
            best, probs = best.cpu().numpy(), probs.cpu().numpy()
            for k, i in enumerate(ids):
                n = int(lengths[k])
                seq, pm = best[k, :n], probs[k, :n]
                keep = (seq != 0) & np.concatenate(([True], seq[1:] != seq[:-1]))
                steps = np.nonzero(keep)[0]
                token_ids = seq[steps].tolist()
                visual = ''.join(self.alphabet[t - 1] for t in token_ids)
                conf = float(pm[keep].mean()) if keep.any() else 0.
                g = receipts[i]
                left, _, right, _ = g['padded_source_extent']
                xs = left + (steps * 4 + 3.5) * (right - left) / crops[i].shape[1]
                toks = list(zip([RTLVisualChar(self.alphabet[t - 1]) for t in token_ids], xs.tolist()))
                receipts[i] = dict(g, raw_visual_text=visual, raw_token_ids=token_ids, confidence=conf, frame_centers=xs.tolist(), confidence_threshold=MIN_CONF, confidence_filtered=conf < MIN_CONF)
                readings[i] = (visual, conf, toks) if conf >= MIN_CONF else ('', conf, [])
        return readings, receipts