import json
import os
from pathlib import Path

import decode as D
import numpy as np
import rules


def windows(n, size, overlap):
    if n <= size:
        return [(0, n)]
    out, start = [], 0
    while True:
        end = min(n, start + size)
        out.append((start, end))
        if end == n:
            return out
        start = end - overlap


class Detector:
    def __init__(self, path, threads=None):
        import onnxruntime as ort
        from tokenizers import Tokenizer
        path = Path(path)
        self.cfg = json.loads((path / 'namecfg.json').read_text())
        self.tok = Tokenizer.from_file(str(path / 'tokenizer.json'))
        self.tok.no_padding()
        self.tok.no_truncation()
        so = ort.SessionOptions()
        so.intra_op_num_threads = threads or int(os.environ.get('NAME_THREADS', '4'))
        so.inter_op_num_threads = 1
        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.sess = ort.InferenceSession(str(path / 'model.onnx'), so, providers=['CPUExecutionProvider'])
        self.inputs = {i.name for i in self.sess.get_inputs()}
        self.window = self.cfg['window']
        self.overlap = self.cfg['overlap']
        self.tau = self.cfg['tau']
        self.cls, self.sep = self.cfg['cls_id'], self.cfg['sep_id']

    def encode(self, text):
        enc = self.tok.encode(text, add_special_tokens=False)
        offsets = []
        for a, b in enc.offsets:
            while a < b and text[a].isspace():
                a += 1
            while b > a and text[b - 1].isspace():
                b -= 1
            offsets.append((a, b))
        return enc.ids, offsets

    def scores(self, text):
        ids, offsets = self.encode(text)
        n = len(ids)
        probs = np.zeros(n, dtype=np.float32)
        best = np.full(n, -1.0)
        for a, b in windows(n, self.window, self.overlap):
            x = np.array([[self.cls] + ids[a:b] + [self.sep]], dtype=np.int64)
            feed = dict(input_ids=x, attention_mask=np.ones_like(x))
            if 'token_type_ids' in self.inputs:
                feed['token_type_ids'] = np.zeros_like(x)
            logits = self.sess.run(None, feed)[0][0, 1:-1]
            e = np.exp(logits - logits.max(-1, keepdims=True))
            p = e[:, 1] / e.sum(-1)
            j = np.arange(b - a, dtype=np.float64)
            far = np.full(b - a, 1e9)
            depth = np.minimum(j if a > 0 else far, (b - a - 1 - j) if b < n else far)
            take = depth > best[a:b]
            probs[a:b][take] = p[take]
            best[a:b][take] = depth[take]
        keep = [i for i, (a, b) in enumerate(offsets) if b > a]
        return [list(offsets[i]) for i in keep], [float(probs[i]) for i in keep]

    def findings(self, text, requester=None, authority=None):
        offsets, probs = self.scores(text)
        found = D.spans_for(text, offsets, probs, self.tau)
        found = rules.apply(text, found, dict(offsets=offsets, probs=probs), requester, authority, self.tau)
        cp = D.char_scores(text, offsets, probs)
        return [dict(start=a, end=b, score=round(max(max(cp[a:b], default=0.0), self.tau), 4)) for a, b in found]