"""End-to-end document processing: file -> pages -> text segments -> labelled fields -> JSON."""

import json
import time
from pathlib import Path

import cv2
import numpy as np

from ..charset import NUM_CLASSES, ctc_greedy_decode
from ..data.datasets import line_to_input, orientation_page, text_patches
from ..features import MAX_SEGMENTS, segment_features
from ..geometry import crop_quad, expand_quad, quad_bbox, quad_size, transform_quads
from ..models.orientation import ANGLES, PATCH
from ..models.recognizer import WIDTH_STRIDE
from ..schema import DOC_TYPES, FIELD_LABELS, REL_TYPES
from .assemble import (assemble_fields, assemble_page_structure, block_groups, block_order, groups_from_block_map,
                       structure_from_groups, xy_cut)
from .detect import DET_MAX_SIDE, detector_input, find_rules, kernels_to_quads, split_gutter_crossers
from .inputs import estimate_skew, load_pages, rot90_matrix, rotate_image

COMPONENTS = ("detector", "recognizer", "orientation", "extractor")

# Frequent words used to check a proposed page rotation by reading a few lines both ways (upside-down text
# still yields plausible-looking letters, so OCR confidence alone does not separate the two).
COMMON_WORDS = frozenset("""
the of and to in is for on that with as by this be are from at or an it was will not have all can has
but which their more one were been other its into also than about after these such may when only over
new they we you he she his her our your there two first time most some would what so if no up out who
total tax subtotal cash change qty amount price date invoice item items paid due balance receipt number
page table figure section chapter results data used use between each both during through under where
""".split())
SCHEMA_VERSION = "1.1"   # 1.1: "formula" blocks; layout blocks from the model's same_block relation
# White columns after every line the recognizer reads. Training pads each line of a batch with white to the right
# (to a multiple of 64 px), so the model reads a line's end against white; without this strip a line read on its
# own, or the widest line of a batch, would end at the image edge instead. 32 px covers the CNN's receptive field
# (64 px gives identical output), and with it the output no longer depends on how lines are batched.
RECOGNIZER_PAD = 32


def _softmax(x, axis=-1):
    e = np.exp(x - x.max(axis, keepdims=True))
    return e / e.sum(axis, keepdims=True)


def line_batch(lines):
    """Recognizer input for line arrays (32, W) in [0, 1]: (batch padded with white, step mask, steps per line)."""
    wmax = max(x.shape[1] for x in lines) + RECOGNIZER_PAD
    batch = np.ones((len(lines), 1, lines[0].shape[0], wmax), np.float32)
    steps = np.array([x.shape[1] // WIDTH_STRIDE for x in lines])
    for k, x in enumerate(lines):
        batch[k, 0, :, : x.shape[1]] = x
    mask = np.arange(wmax // WIDTH_STRIDE)[None, :] >= steps[:, None]
    return batch, mask, steps


def _sigmoid(x):
    return 1 / (1 + np.exp(-np.clip(x, -60, 60)))


class TorchBackend:
    def __init__(self, weights_dir, device="cpu"):
        import torch

        from ..models import build, count_params

        self.torch, self.device, self.models, self.params = torch, torch.device(device), {}, {}
        for name in COMPONENTS:
            p = Path(weights_dir) / name / "best.pt"
            if not p.exists():
                raise FileNotFoundError(f"missing {p}; train it with `python -m tinydoc.train {name}`")
            ck = torch.load(p, map_location="cpu", weights_only=False)
            m = build(name, **(ck.get("arch") or {}))
            m.load_state_dict(ck["model"])
            self.models[name] = m.to(self.device).eval()
            self.params[name] = count_params(m)
        self.vocab = NUM_CLASSES

    def __call__(self, name, *inputs):
        with self.torch.no_grad():
            t = [self.torch.from_numpy(np.ascontiguousarray(x)).to(self.device) for x in inputs]
            out = self.models[name](*t)
            out = out if isinstance(out, tuple) else (out,)
            return [o.float().cpu().numpy() for o in out]


class OnnxBackend:
    def __init__(self, onnx_dir, threads=None):
        import onnx
        import onnxruntime as ort

        opts = ort.SessionOptions()
        if threads:
            opts.intra_op_num_threads = threads
        self.sessions, self.params = {}, {}
        for name in COMPONENTS:
            p = Path(onnx_dir) / f"{name}.onnx"
            if not p.exists():
                raise FileNotFoundError(f"missing {p}; export with `python -m tinydoc.export`")
            self.sessions[name] = ort.InferenceSession(str(p), opts, providers=["CPUExecutionProvider"])
            meta = {m.key: m.value for m in onnx.load(str(p), load_external_data=False).metadata_props}
            self.params[name] = int(meta.get("trained_parameters", 0))
        # Character vocabulary the models were trained with: releases before the inline-maths characters were
        # appended have fewer classes, and their extractor must not see the newer ids.
        width = self.sessions["recognizer"].get_outputs()[0].shape[-1]
        self.vocab = width if isinstance(width, int) else NUM_CLASSES

    def __call__(self, name, *inputs):
        s = self.sessions[name]
        feed = {i.name: np.ascontiguousarray(x) for i, x in zip(s.get_inputs(), inputs)}
        return s.run(None, feed)


class Engine:
    def __init__(self, weights="runs", backend="torch", device="cpu", recognizer_batch=None, model_blocks=None, blocks=None):
        self.backend = TorchBackend(weights, device) if backend == "torch" else OnnxBackend(weights)
        # Lines per recognizer call. On CPU (ONNX Runtime) one line per call is fastest: nothing is padded to the
        # widest line of a batch and the activations stay in cache. On 40 OmniDocBench pages (Apple M5 Pro) it cut
        # recognition from 421 to 248 ms per page at 4 threads and from 1190 to 440 ms at 1 thread, against
        # batches of 32, with the same output (see RECOGNIZER_PAD). The torch backend keeps batches, which pay
        # off on a GPU.
        self.recognizer_batch = recognizer_batch or (1 if backend == "onnx" else 32)
        # How lines are grouped into paragraphs/blocks on ordinary pages: "geometric" rules, the extractor's
        # same_block relation ("model"), or the detector's block regions ("detector"). The default comes from
        # config.json next to the weights, chosen by measuring each on validation pages.
        cfg = Path(weights) / "config.json"
        cfg = json.loads(cfg.read_text()) if cfg.exists() else {}
        if blocks is None and model_blocks is not None:
            blocks = "model" if model_blocks else "geometric"
        self.blocks = blocks or cfg.get("blocks") or ("model" if cfg.get("model_blocks") else "geometric")
        # Splitting lines merged across a column gutter: off by default (on OmniDocBench such merges were rare,
        # 5 in 380 pages, and the split hurt as often as it helped).
        self.split_gutters = cfg.get("split_gutters", False)
        # Detector input size: one pass with the long side at det_max_side. Optionally (det_fast_side) a page is
        # first detected at a smaller size and again at det_max_side only when its median line is under
        # det_min_line_px there. Off by default: on v3 it was ~15% faster and helped academic and exam pages, but
        # the official OmniDocBench evaluator showed it losing small print on newspapers (+0.017 to +0.038 text
        # edit) and magazines (+0.02), which the median-line test does not catch.
        self.det_max_side = cfg.get("det_max_side", DET_MAX_SIDE)
        self.det_fast_side = cfg.get("det_fast_side")
        self.det_min_line_px = cfg.get("det_min_line_px", 14)

    # ------------------------------------------------------------------------------------ stages
    def orientation(self, img):
        gray = orientation_page(img)
        pos = text_patches(gray, 8)
        patches = np.stack([gray[y:y + PATCH, x:x + PATCH] for x, y in pos]).astype(np.float32)[:, None] / 255
        if patches.shape[-1] != PATCH or patches.shape[-2] != PATCH:
            return 0, 1.0
        probs = _softmax(self.backend("orientation", patches)[0]).mean(0)
        k = int(probs.argmax())
        return k, float(probs[k])

    def _reading_evidence(self, img, n=16):
        """(common-word hits, mean OCR confidence) over the n widest detected lines of an image."""
        dets = sorted(self.detect(img), key=lambda d: -quad_size(d[0])[0])[:n]
        if not dets:
            return 0, 0.0
        texts = self.recognize(img, [q for q, _ in dets])
        words = [w.strip(".,:;!?()'\"").lower() for t, _ in texts for w in t.split()]
        return sum(w in COMMON_WORDS for w in words), float(np.mean([c for _, c in texts]))

    def verify_rotation(self, img, k):
        """Keep a proposed 90/180/270 correction only if the rotated page reads better than the original."""
        hits0, conf0 = self._reading_evidence(img)
        hitsk, confk = self._reading_evidence(np.ascontiguousarray(np.rot90(img, k)))
        if abs(hitsk - hits0) >= 2:
            return hitsk > hits0
        if abs(confk - conf0) > 0.05:
            return confk > conf0
        return True   # no clear evidence either way: trust the orientation classifier

    def detect_full(self, img, blocks=True):
        """Line detections plus the block-region map (None for detectors without a block channel, or when
        blocks=False because the caller does not use it)."""
        for side in ([self.det_fast_side] if self.det_fast_side else []) + [self.det_max_side]:
            x, scale = detector_input(img, max_side=side)
            out = self.backend("detector", x[None])[0][0]
            prob = _sigmoid(out[0])
            quads = kernels_to_quads(prob)
            if side == self.det_max_side or not quads or np.median([quad_size(q)[1] for q, _ in quads]) >= self.det_min_line_px:
                break
        block = _sigmoid(out[1]) if blocks and out.shape[0] > 1 else None
        return [(q / scale, s) for q, s in quads], block, scale

    def detect(self, img):
        return self.detect_full(img)[0]

    def recognize(self, img, quads):
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        crops = []
        for q in quads:
            _, h = quad_size(q)
            crop = crop_quad(gray, expand_quad(q, 0.25 * h, 0.1 * h, 0.25 * h, 0.1 * h))
            if np.median(crop[::2, ::2]) < 100:  # light text on a dark band
                crop = 255 - crop
            crops.append(line_to_input(crop))
        results = [None] * len(crops)
        order = sorted(range(len(crops)), key=lambda i: crops[i].shape[1])
        for start in range(0, len(order), self.recognizer_batch):
            idx = order[start:start + self.recognizer_batch]
            batch, mask, steps = line_batch([crops[i] for i in idx])
            probs = _softmax(self.backend("recognizer", batch, mask)[0])
            for k, i in enumerate(idx):
                text, confs = ctc_greedy_decode(probs[k, : steps[k]])
                results[i] = (text, float(np.mean(confs)) if confs else 0.0)
        return results

    def extract(self, segs, W, H):
        if not segs:
            return np.zeros((0, len(FIELD_LABELS))), np.zeros((len(REL_TYPES), 0, 0)), np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
        n = len(segs)
        labels = np.zeros((n, len(FIELD_LABELS)), np.float32)
        rel = None   # sized from the model: older exports predict fewer relations
        docs = []
        # windows of MAX_SEGMENTS in reading order, overlapping by half, for very dense pages
        step = MAX_SEGMENTS // 2
        starts = list(range(0, max(n - MAX_SEGMENTS, 0) + 1, step)) or [0]
        if starts[-1] + MAX_SEGMENTS < n:
            starts.append(n - MAX_SEGMENTS)
        for s0 in starts:
            idx = list(range(s0, min(n, s0 + MAX_SEGMENTS)))
            f = segment_features([segs[i] for i in idx], W, H)
            f["char_ids"][f["char_ids"] >= getattr(self.backend, "vocab", NUM_CLASSES)] = 0
            out = self.backend("extractor", f["char_ids"][None], f["feats"][None], f["rel_boxes"][None],
                               np.ones((1, len(idx)), np.float32))
            lp = _softmax(out[0][0])
            labels[idx, :lp.shape[1]] = np.maximum(labels[idx, :lp.shape[1]], lp)
            if rel is None:
                rel = np.zeros((out[1].shape[1], n, n), np.float32)
            rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1] = np.maximum(rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1], _sigmoid(out[1][0]))
            docs.append(_softmax(out[2][0]))
        return labels, rel, np.mean(docs, 0)

    # -------------------------------------------------------------------------------------- run
    def process(self, path, force_ocr=False):
        t_start = time.time()
        timings = {"load": 0.0, "orientation": 0.0, "deskew": 0.0, "detect": 0.0, "recognize": 0.0, "extract": 0.0, "assemble": 0.0}
        pages_out, all_segs, page_info = [], [], []
        t = time.time()
        pages = list(load_pages(path, force_ocr=force_ocr))
        timings["load"] += time.time() - t

        for page in pages:
            info = {"page": page.number, "width": page.width, "height": page.height, "text_source": page.text_source}
            if page.text_source == "pdf_text":
                segs = [{"bbox": [x0, y0, x1, y1], "quad": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], "text": t_, "conf": 1.0}
                        for x0, y0, x1, y1, t_ in page.pdf_segments]
                frame_w, frame_h = page.width, page.height
                rules, block_map, det_scale = [], None, 1.0
            else:
                img = page.image
                t = time.time()
                k, kconf = self.orientation(img)
                if k != 0 and kconf > 0.6 and not self.verify_rotation(img, k):
                    k, kconf = 0, 1.0 - kconf
                    info["rotation_rejected"] = True
                timings["orientation"] += time.time() - t
                m_total = np.eye(3)
                if k != 0 and kconf > 0.6:
                    m_total = np.vstack([rot90_matrix(k, img.shape[1], img.shape[0]), [0, 0, 1]]) @ m_total
                    img = np.ascontiguousarray(np.rot90(img, k))
                info["rotation_corrected_deg"] = ANGLES[k] if (k != 0 and kconf > 0.6) else 0
                info["orientation_confidence"] = round(kconf, 3)
                t = time.time()
                skew = estimate_skew(img)  # text is rotated `skew` degrees counter-clockwise
                if abs(skew) >= 0.5:
                    img, m = rotate_image(img, -skew)  # rotate it back
                    m_total = np.vstack([m, [0, 0, 1]]) @ m_total
                info["skew_corrected_deg"] = round(skew, 2) if abs(skew) >= 0.5 else 0.0
                timings["deskew"] += time.time() - t
                t = time.time()
                dets, block_map, det_scale = self.detect_full(img, blocks=self.blocks == "detector")
                if self.split_gutters:
                    dets = split_gutter_crossers(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), dets)
                timings["detect"] += time.time() - t
                t = time.time()
                texts = self.recognize(img, [q for q, _ in dets])
                timings["recognize"] += time.time() - t
                inv = np.linalg.inv(m_total)
                segs = []
                for (q, dscore), (text, conf) in zip(dets, texts):
                    if not text.strip():
                        continue
                    orig_q = transform_quads([q], inv)[0]   # report positions in the source image
                    segs.append({"bbox": quad_bbox(q), "quad_src": orig_q, "text": text.strip(),
                                 "conf": conf * min(1.0, dscore / 0.7)})
                frame_w, frame_h = img.shape[1], img.shape[0]
                hs = [s["bbox"][3] - s["bbox"][1] for s in segs]
                rules = find_rules(img, float(np.median(hs))) if hs else []
            for i, s in enumerate(segs):
                s["page"] = page.number
                s["id"] = f"p{page.number}-s{i}"
            t = time.time()
            labels, rel, doc = self.extract(segs, frame_w, frame_h)
            timings["extract"] += time.time() - t
            boxes = [s["bbox"] for s in segs]
            groups = None
            general = bool(segs) and DOC_TYPES[int(doc.argmax())] == "general"
            if general and self.blocks == "detector" and block_map is not None:
                groups = block_order(boxes, groups_from_block_map(boxes, block_map, det_scale), rules)
            elif general and self.blocks == "model" and rel.shape[0] > REL_TYPES.index("same_block"):
                names = [FIELD_LABELS[int(r.argmax())] for r in labels]
                groups = block_order(boxes, block_groups(boxes, names, rel[REL_TYPES.index("same_block")]), rules)
            if groups is not None:
                order = [i for g in groups for i in g]
            else:
                order = xy_cut(boxes, rules=rules) if segs else []
            page_info.append((info, segs, labels, rel, doc, order, groups))

        t = time.time()
        doc_probs = np.mean([p[4] for p in page_info if len(p[1])], 0) if any(len(p[1]) for p in page_info) else np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
        doc_type = DOC_TYPES[int(doc_probs.argmax())]

        # global segment list across pages for field assembly
        offset = 0
        glob_segs, glob_labels, glob_conf, glob_order = [], [], [], []
        n_total = sum(len(p[1]) for p in page_info)
        same_row = np.zeros((n_total, n_total), np.float32)
        key_value = np.zeros((n_total, n_total), np.float32)
        for info, segs, labels, rel, _, order, groups in page_info:
            n = len(segs)
            glob_segs += segs
            glob_labels += [FIELD_LABELS[int(r.argmax())] for r in labels]
            glob_conf += [float(r.max()) for r in labels]
            glob_order += [offset + i for i in order]
            if n:
                same_row[offset:offset + n, offset:offset + n] = rel[0]
                key_value[offset:offset + n, offset:offset + n] = rel[1]
            label_names = [FIELD_LABELS[int(r.argmax())] for r in labels]
            if groups is not None:
                blocks, tables = structure_from_groups(segs, groups, label_names, rel[0])
            else:
                blocks, tables = assemble_page_structure(segs, order, label_names, rel[0] if n else np.zeros((0, 0)))
            pages_out.append({**info, "blocks": blocks, "tables": tables, "segments": [
                {"id": s["id"], "text": s["text"], "bbox": [round(v, 1) for v in (quad_bbox(s["quad_src"]) if "quad_src" in s else s["bbox"])],
                 "quad": np.round(np.asarray(s.get("quad_src", s.get("quad"))), 1).tolist(),
                 "ocr_confidence": round(s["conf"], 3), "label": label_names[i], "label_confidence": round(float(labels[i].max()), 3)}
                for i, s in enumerate(segs)]})
            offset += n

        result = {
            "schema_version": SCHEMA_VERSION,
            "source": {"file": str(Path(path).name), "pages": len(pages)},
            "document_type": {"value": doc_type, "confidence": round(float(doc_probs.max()), 3),
                              "scores": {d: round(float(p), 3) for d, p in zip(DOC_TYPES, doc_probs)}},
        }
        if n_total == 0:
            result["document_type"] = {"value": None, "confidence": 0.0, "scores": None}
            result["fields"] = None
            result["validation"] = {"checks": [], "issues": ["no text found"], "needs_review": True}
        elif doc_type in ("receipt", "invoice"):
            fields, validation = assemble_fields(glob_segs, glob_order, glob_labels, glob_conf, same_row, key_value)
            result["fields"] = fields
            result["validation"] = validation
        else:
            result["fields"] = None
            result["validation"] = {"checks": [], "issues": [], "needs_review": float(doc_probs.max()) < 0.6}
        result["text"] = "\n\n".join("\n".join(b.get("text", "") for b in p["blocks"] if b.get("text")) for p in pages_out)
        result["pages"] = pages_out
        timings["assemble"] += time.time() - t
        result["timing_ms"] = {k: round(v * 1000, 1) for k, v in timings.items()}
        result["timing_ms"]["total"] = round((time.time() - t_start) * 1000, 1)
        result["model"] = {"parameters": dict(self.backend.params), "total_parameters": int(sum(self.backend.params.values()))}
        return result