# standard imports
import typing

# external imports
import torch
from facenet_pytorch import MTCNN, InceptionResnetV1

# bsie imports
from bsie.utils import bsfs, node, ns

# inner-module imports
from ... import base

# exports
__all__: typing.Sequence[str] = (
    'FaceDetect',
    )


## code ##

bsf = ns.bsn.Face()

class FaceDetect(base.Extractor):

    CONTENT_READER = 'bsie.reader.face.FaceExtract'

    def __init__(self):
        # initialize parent with the schema
        super().__init__(bsfs.schema.from_string(base.SCHEMA_PREAMBLE + f'''
            prefix bsf: <https://schema.bsfs.io/ie/Node/Face#>

            bsn:Face rdfs:subClassOf bsfs:Node .

            <https://schema.bsfs.io/ie/Literal/Array/Feature/Face#resnet512>
                rdfs:subClassOf bsa:Feature ;
                bsfs:distance <https://schema.bsfs.io/core/distance#euclidean> ;
                bsfs:dtype <https://schema.bsfs.io/core/dtype#f32>;
                bsfs:dimension "512"^^xsd:integer .

            bse:face rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Entity ;
                rdfs:range bsn:Face .

            bsf:x rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Face ;
                rdfs:range xsd:float ;
                bsfs:unique "true"^^xsd:boolean .

            bsf:y rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Face ;
                rdfs:range xsd:float ;
                bsfs:unique "true"^^xsd:boolean .

            bsf:width rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Face ;
                rdfs:range xsd:float ;
                bsfs:unique "true"^^xsd:boolean .

            bsf:height rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Face ;
                rdfs:range xsd:float ;
                bsfs:unique "true"^^xsd:boolean .

            bsf:embedding rdfs:subClassOf bsfs:Predicate ;
                rdfs:domain bsn:Face ;
                rdfs:range <https://schema.bsfs.io/ie/Literal/Array/Feature/Face#resnet512> ;
                bsfs:unique "true"^^xsd:boolean .

            '''))

    def extract(
            self,
            subject: node.Node,
            content: dict,
            principals: typing.Iterable[bsfs.schema.Predicate],
            ) -> typing.Iterator[typing.Tuple[node.Node, bsfs.schema.Predicate, typing.Any]]:
        # check principals
        if self.schema.predicate(ns.bse.face) not in principals:
            # nothing to do; abort
            return

        for face in content:
            fnode = node.Node(ns.bsn.Face, ucid=face['ucid'])
            yield subject, ns.bse.face, fnode
            yield fnode, bsf.x, face['x']
            yield fnode, bsf.y, face['y']
            yield fnode, bsf.width, face['width']
            yield fnode, bsf.height, face['height']
            yield fnode, bsf.embedding, face['embedding'].detach().cpu().numpy()

## EOF ##