from pathlib import Path
from typing import List

import numpy as np
from xlib import math as lib_math
from xlib.image import ImageProcessor
from xlib.onnxruntime import (InferenceSession_with_device, ORTDeviceInfo,
                              get_available_devices_info)


class S3FD:

    @staticmethod
    def get_available_devices() -> List[ORTDeviceInfo]:
        return get_available_devices_info()

    def __init__(self, device_info : ORTDeviceInfo ):
        if device_info not in S3FD.get_available_devices():
            raise Exception(f'device_info {device_info} is not in available devices for S3FD')

        path = Path(__file__).parent / 'S3FD.onnx'
        self._sess = sess = InferenceSession_with_device(str(path), device_info)
        self._input_name = sess.get_inputs()[0].name


    def extract(self, img : np.ndarray, threshold=0.95, fixed_window=0, min_face_size=40):
        """

            img     HW,HWC,NHWC     [0..255]
        """
        ip = ImageProcessor(img)

        if fixed_window != 0:
            fixed_window = max(64, max(1, fixed_window // 32) * 32 )
            img_scale = ip.fit_in(fixed_window, fixed_window, pad_to_target=True, allow_upscale=False)
        else:
            ip.pad_to_next_divisor(64, 64)
            img_scale = 1.0

        img = ip.ch(3).to_uint8().as_float32().apply( lambda img: img - [104,117,123]).get_image('NCHW')

        batches_bbox = self._sess.run(None, {self._input_name: img})

        faces_per_batch = []
        for batch in range(img.shape[0]):
            bbox = self.refine( [ x[batch] for x in batches_bbox ], threshold )

            faces = []
            for l,t,r,b,c in bbox:
                if img_scale != 1.0:
                    l,t,r,b = l/img_scale, t/img_scale, r/img_scale, b/img_scale

                bt = b-t
                if min(r-l,bt) < min_face_size:
                    continue
                b += bt*0.1

                faces.append ( (l,t,r,b) )

            faces_per_batch.append(faces)

        return faces_per_batch


    def refine(self, olist, threshold):
        bboxlist = []
        variances = [0.1, 0.2]
        for i in range(len(olist) // 2):
            ocls, oreg = olist[i * 2], olist[i * 2 + 1]

            stride = 2**(i + 2)    # 4,8,16,32,64,128
            for hindex, windex in [*zip(*np.where(ocls[1, :, :] > threshold))]:
                axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride
                score = ocls[1, hindex, windex]
                loc = np.ascontiguousarray(oreg[:, hindex, windex]).reshape((1, 4))
                priors = np.array([[axc, ayc, stride * 4, stride * 4]])
                bbox = np.concatenate((priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
                                       priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])), 1)
                bbox[:, :2] -= bbox[:, 2:] / 2
                bbox[:, 2:] += bbox[:, :2]
                x1, y1, x2, y2 = bbox[0]
                bboxlist.append([x1, y1, x2, y2, score])

        if len(bboxlist) != 0:
            bboxlist = np.array(bboxlist)
            bboxlist = bboxlist[ lib_math.nms(bboxlist[:,0], bboxlist[:,1], bboxlist[:,2], bboxlist[:,3], bboxlist[:,4], 0.3), : ]
            bboxlist = [x for x in bboxlist if x[-1] >= 0.5]

        return bboxlist