"""Orientation screen for 75 dpi, full logical pages supplied as RGB arrays.

Create one Detector per process and reuse it. The conservative profile is the
recommended default. Scores are screening signals, not calibrated probabilities.
"""
from pathlib import Path
import json
import cv2
import joblib
import numpy as np
import onnxruntime as ort


def ink(image, level=100):
    mask = 255 - cv2.inRange(image, (level,)*3, (255,)*3)
    mask[:15] = mask[-15:] = 0
    mask[:, :15] = mask[:, -15:] = 0
    return mask


def spatial(mask):
    result = []
    for width, height in [(8, 12), (16, 24)]:
        grid = np.sqrt(cv2.resize(mask.astype(np.float32)/255, (width, height), interpolation=cv2.INTER_AREA))
        result.append((grid/(np.linalg.norm(grid)+1e-6)).ravel())
    return np.concatenate(result)


def bounds(mask):
    ys, xs = np.where(mask)
    if len(xs) < 5:
        return 0, 0, mask.shape[1], mask.shape[0]
    x0, x1 = np.quantile(xs, [.005, .995]).astype(int)
    y0, y1 = np.quantile(ys, [.005, .995]).astype(int)
    return max(0,x0-12), max(0,y0-12), min(mask.shape[1],x1+13), min(mask.shape[0],y1+13)


class Detector:
    def __init__(self, directory=None, profile='conservative'):
        directory = Path(directory) if directory else Path(__file__).parent/'data'/'page_orientation'
        config = json.loads((directory/'settings.json').read_text())
        self.threshold = config['profiles'][profile]['threshold']
        model = joblib.load(directory/'model.joblib')
        self.heat_model, self.fusion_model = model['heat_model'], model['model']
        options = ort.SessionOptions()
        options.intra_op_num_threads = options.inter_op_num_threads = 1
        options.log_severity_level = 3
        self.session = ort.InferenceSession(str(directory/'doc.onnx'), options, providers=['CPUExecutionProvider'])
        self.input_name = self.session.get_inputs()[0].name
        self.mean = np.array([.485,.456,.406], np.float32)
        self.std = np.array([.229,.224,.225], np.float32)

    def predict(self, image_rgb):
        mask = ink(image_rgb)
        x0, y0, x1, y1 = bounds(mask)
        crop = image_rgb[y0:y1, x0:x1]
        feature = np.concatenate([spatial(mask), spatial(mask[y0:y1,x0:x1]), spatial(ink(image_rgb,140))])
        heat_score = float(self.heat_model.decision_function(feature[None])[0])
        probabilities = []
        for view in [image_rgb, crop]:
            for rotate in [False, True]:
                image = np.rot90(view,2).copy() if rotate else view
                image = cv2.resize(image,(224,224))
                tensor = ((image.astype(np.float32)/255-self.mean)/self.std).transpose(2,0,1)[None].copy()
                probabilities.extend(self.session.run(None,{self.input_name:tensor})[0][0])
        features = np.array([probabilities+[heat_score]])
        score = float(self.fusion_model.predict_proba(features)[0,1])
        return {'flag': score > self.threshold, 'score': score, 'threshold': self.threshold, 'heatmap_score': heat_score}