"""Requ Classify Small: reads a company's website and labels it.

Gives the industry (64 classes), a type for each menu item and heading (offering, team, activity,
audience, website, other), a customer-journey stage for each page, and sentences that show a known
business process. Needs numpy, onnxruntime, tokenizers and safetensors.

    python requ_classify.py < site.json

site.json: {"name", "description", "pages": [{"url", "title", "text", "items": [{"text", "kind"}]}]}
where an item's kind is one of: menu, footer, h1, h2, h3, button.
"""
import json, os, re, sys
from urllib.parse import urlsplit

import numpy as np
from safetensors.numpy import load_file

HERE = os.path.dirname(os.path.abspath(__file__))
KINDS = ['menu', 'footer', 'h1', 'h2', 'h3', 'button']
SENT = re.compile(r'(?<=[.!?])\s+')


class Embedder:
    """bge-small-en-v1.5 (quantised ONNX), CLS vector, normalised."""

    def __init__(self, folder):
        import onnxruntime as ort
        from tokenizers import Tokenizer
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = int(os.environ.get('REQU_THREADS', '2'))
        self.session = ort.InferenceSession(os.path.join(folder, 'model_quantized.onnx'), opts, providers=['CPUExecutionProvider'])
        self.tok = Tokenizer.from_file(os.path.join(folder, 'tokenizer.json'))
        self.tok.enable_truncation(max_length=256)
        self.inputs = {i.name for i in self.session.get_inputs()}

    def __call__(self, texts, batch=64):
        out = []
        for start in range(0, len(texts), batch):
            self.tok.enable_padding(length=None)
            enc = self.tok.encode_batch(texts[start:start + batch])
            ids = np.array([e.ids for e in enc], dtype=np.int64)
            mask = np.array([e.attention_mask for e in enc], dtype=np.int64)
            feed = {'input_ids': ids, 'attention_mask': mask}
            if 'token_type_ids' in self.inputs:
                feed['token_type_ids'] = np.zeros_like(ids)
            cls = self.session.run(None, feed)[0][:, 0]
            out.append(cls / np.linalg.norm(cls, axis=1, keepdims=True))
        return np.vstack(out) if out else np.zeros((0, 384), dtype=np.float32)


def path_of(url):
    return urlsplit(url).path.strip('/').replace('-', ' ').replace('/', ' / ') or 'home'


def site_text(site):
    items = [i for p in site['pages'] for i in p.get('items', [])]
    heads = [i['text'] for i in items if i['kind'] in ('h1', 'h2')][:25]
    menu = [i['text'] for i in items if i['kind'] == 'menu'][:30]
    home = site['pages'][0]['text'][:800] if site['pages'] else ''
    return f"{site.get('name', '')}. {site.get('description', '')} Pages: {'; '.join(p.get('title', '') for p in site['pages'])}. Menu: {', '.join(menu)}. Headings: {', '.join(heads)}. {home}"


def site_pieces(site):
    items = [i for p in site['pages'] for i in p.get('items', [])]
    out = [f"{site.get('name', '')}. {site.get('description', '')}",
           'Menu: ' + ', '.join(i['text'] for i in items if i['kind'] == 'menu')[:600],
           'Headings: ' + ', '.join(i['text'] for i in items if i['kind'] in ('h1', 'h2', 'h3'))[:600]]
    out += [f"{p.get('title', '')}. {p['text'][:500]}" for p in site['pages'][:6]]
    return [x for x in out if len(x) > 12]


def sentences(page, limit=60):
    out = [s.strip() for s in SENT.split(page['text']) if 6 <= len(s.split()) <= 45]
    return out[:limit]


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


class RequClassify:
    def __init__(self, folder=HERE):
        self.embed = Embedder(os.path.join(folder, 'model'))
        w = load_file(os.path.join(folder, 'heads.safetensors'))
        cfg = json.load(open(os.path.join(folder, 'config.json')))
        self.heads = {n: (w[f'{n}.weight'], w[f'{n}.bias'], cfg['classes'][n]) for n in ('industry', 'item_type', 'page_stage')}
        self.proof = (w['proof.weight'], w['proof.bias'])
        self.memory = w['memory'].astype(np.float32)
        self.memory_keys = np.array(cfg['memory_keys'])
        self.vectors = w['library_vectors']
        self.t = cfg['thresholds']
        self.library = json.load(open(os.path.join(folder, 'library.json')))['industries']

    def head(self, name, X):
        coef, bias, classes = self.heads[name]
        return softmax(X @ coef.T + bias), classes

    def __call__(self, site):
        pages = site['pages']
        X = np.hstack([self.embed([site_text(site)])[0], self.embed(site_pieces(site)).mean(0)])[None, :]
        p, classes = self.head('industry', X)
        order = np.argsort(-p[0])
        industry = classes[order[0]]
        result = {'industry': {'id': industry, 'p': round(float(p[0][order[0]]), 3), 'sure': bool(p[0][order[0]] >= self.t['industry_gate']),
                               'top3': [[classes[i], round(float(p[0][i]), 3)] for i in order[:3]]}}
        flat = [(pi, it) for pi, pg in enumerate(pages) for it in pg.get('items', [])]
        result['items'] = []
        if flat:
            E = self.embed([f"{it['text']} | {it['kind']} | {path_of(pages[pi]['url'])}" for pi, it in flat])
            K = np.array([[1.0 if it['kind'] == k else 0.0 for k in KINDS] for _, it in flat])
            q, tclasses = self.head('item_type', np.hstack([E, K]))
            result['items'] = [{'page': pi, 'text': it['text'], 'kind': it['kind'], 'type': tclasses[int(r.argmax())], 'p': round(float(r.max()), 3)} for (pi, it), r in zip(flat, q)]
        if pages:
            s, sclasses = self.head('page_stage', self.embed([f"{path_of(pg['url'])}. {pg.get('title', '')}. {pg['text'][:400]}" for pg in pages]))
            result['pageStages'] = [sclasses[int(r.argmax())] for r in s]
        # Sentences that show a business process of this industry, matched by the library wording
        # and by similar sentences from training sites.
        entry = self.library[industry]
        P = self.vectors[entry['rows']]
        keys = [pid if pid.startswith('s-') else f'{industry}:{pid}' for pid in entry['ids']]
        found = []
        sents = [(pi, s) for pi, pg in enumerate(pages) for s in sentences(pg)]
        if sents:
            S = self.embed([s for _, s in sents])
            keep = (1 / (1 + np.exp(-(S @ self.proof[0].T + self.proof[1]))))[:, 0]
            chosen = np.where(keep >= self.t['match_threshold'])[0]
            if len(chosen):
                S = S[chosen]
                score = S @ P.T
                simM = S @ self.memory.T
                for j, k in enumerate(keys):
                    cols = np.where(self.memory_keys == k)[0]
                    if len(cols):
                        score[:, j] = 0.4 * score[:, j] + 0.6 * np.sort(simM[:, cols], axis=1)[:, -3:].mean(axis=1)
                for row, i in enumerate(chosen):
                    o = np.argsort(-score[row])
                    pi, s = sents[i]
                    clear = bool(score[row][o[0]] - score[row][o[1]] >= self.t['process_margin']) if len(o) > 1 else True
                    found.append({'process': entry['ids'][int(o[0])], 'clear': clear, 'page': pi, 'quote': s, 'score': round(float(keep[i]), 3)})
        found.sort(key=lambda e: -e['score'])
        best = {}
        for e in found:
            best.setdefault(e['process'], []).append(e)
        result['evidence'] = [x for v in best.values() for x in v[:2]]
        return result


if __name__ == '__main__':
    json.dump(RequClassify()(json.load(sys.stdin)), sys.stdout, indent=1)