"""Per-segment features for the extraction model. Used identically in training and inference."""

import math
import re

import numpy as np

from .charset import CHAR_TO_ID, MATH_ID

MAX_CHARS = 40
MAX_SEGMENTS = 256

_MONEY = re.compile(r"^[-(]?\s*[$€£¥₹]?\s*-?\d{1,3}([,. ]\d{3})*([.,]\d{2})\s*-?\)?$|^[-(]?[$€£¥₹]\s*\d+([.,]\d{1,2})?\)?$")
_NUMBER = re.compile(r"^[-+(]?[$€£¥₹]?\s*\d[\d,. ]*[)%]?$")
_DATE = re.compile(r"\b(\d{1,4}[-/.]\d{1,2}[-/.]\d{1,4}|\d{1,2}\s*[A-Za-z]{3,9}\.?,?\s*\d{2,4}|[A-Za-z]{3,9}\.?\s+\d{1,2},?\s+\d{4})\b")
_TIME = re.compile(r"\b\d{1,2}:\d{2}(:\d{2})?\s*([AaPp][Mm])?\b")
_PHONE = re.compile(r"(\+?\d[\d\s().-]{7,}\d)")
_CURRENCY = re.compile(r"[$€£¥₹]|\b(USD|EUR|GBP|INR|JPY|CAD|AUD|SGD|AED|CHF)\b")


def text_flags(text: str) -> list[float]:
    t = text.strip()
    n = max(len(t), 1)
    digits = sum(c.isdigit() for c in t)
    alphas = sum(c.isalpha() for c in t)
    words = t.split()
    return [
        float(bool(_NUMBER.match(t))),
        float(bool(_MONEY.match(t))),
        float(bool(_DATE.search(t))),
        float(bool(_TIME.search(t))),
        float("%" in t),
        float(bool(_CURRENCY.search(t))),
        float(":" in t),
        float(t.endswith(":")),
        float(alphas > 1 and t.upper() == t),
        float(bool(words) and all(w[:1].isupper() for w in words if w[:1].isalpha())),
        digits / n,
        alphas / n,
        min(len(t), MAX_CHARS) / MAX_CHARS,
        float("@" in t),
        float(bool(_PHONE.search(t)) and digits >= 7),
        float(len(set(t)) <= 2 and len(t) >= 3 and not t[0].isalnum()),  # ----- / ===== separators
        float(t[:1] in "•-*–" or bool(re.match(r"^\d+[.)]\s", t))),
        min(len(words), 10) / 10.0,
        float(t[:1] == "#" or "no." in t.lower()),
    ]


NUM_FLAGS = len(text_flags(""))
NUM_FEATS = 8 + 2 + 16 + NUM_FLAGS + 1


def segment_features(segments, page_w, page_h):
    """segments: list of dicts with 'bbox' [x0, y0, x1, y1] in pixels, 'text', optional 'conf'.

    Returns dict of numpy arrays: char_ids (N, MAX_CHARS) int64, feats (N, F) float32,
    rel_boxes (N, 4) float32 (box in units of median text height, for pairwise geometry).
    """
    n = len(segments)
    char_ids = np.zeros((n, MAX_CHARS), np.int64)
    feats = np.zeros((n, NUM_FEATS), np.float32)
    rel = np.zeros((n, 4), np.float32)
    if n == 0:
        return {"char_ids": char_ids, "feats": feats, "rel_boxes": rel}
    boxes = np.array([s["bbox"] for s in segments], np.float32)
    hs = np.maximum(boxes[:, 3] - boxes[:, 1], 1.0)
    h_med = float(np.median(hs))
    for i, s in enumerate(segments):
        ids = [min(CHAR_TO_ID[c], MATH_ID) for c in s["text"][:MAX_CHARS] if c in CHAR_TO_ID]
        char_ids[i, : len(ids)] = ids
        x0, y0, x1, y1 = boxes[i]
        nx0, nx1, ny0, ny1 = x0 / page_w, x1 / page_w, y0 / page_h, y1 / page_h
        cx, cy = (nx0 + nx1) / 2, (ny0 + ny1) / 2
        w, h = max(x1 - x0, 1.0), hs[i]
        f = [nx0, ny0, nx1, ny1, cx, cy, nx1 - nx0, ny1 - ny0, math.log(h / h_med), math.log(w / h) / 4]
        for freq in (1, 2, 4, 8):
            f += [math.sin(math.pi * freq * cx), math.cos(math.pi * freq * cx),
                  math.sin(math.pi * freq * cy), math.cos(math.pi * freq * cy)]
        f += text_flags(s["text"])
        f.append(float(s.get("conf", 1.0)))
        feats[i] = f
    rel[:] = boxes / h_med
    return {"char_ids": char_ids, "feats": feats, "rel_boxes": rel}