import * as faceapi from 'face-api.js';
import { FaceMatcher } from 'face-api.js';
import { computed } from 'mobx';
import { Doc, DocListCast } from '../../../fields/Doc';
import { DocData } from '../../../fields/DocSymbols';
import { List } from '../../../fields/List';
import { listSpec } from '../../../fields/Schema';
import { Cast, ImageCast, NumCast, StrCast } from '../../../fields/Types';
import { DocumentType } from '../../documents/DocumentTypes';
import { DocumentManager } from '../../util/DocumentManager';

/**
 * A class that handles face recognition.
 */
export class FaceRecognitionHandler {
    static _instance: FaceRecognitionHandler;
    private _loadedModels: boolean = false;
    private _processingDocs: Set<Doc> = new Set();
    private _pendingLoadDocs: Doc[] = [];

    public static FaceField = (target: Doc, doc: Doc) => `${Doc.LayoutFieldKey(target)}_${doc.face_label}`;
    public static FacesField = (target: Doc) => `${Doc.LayoutFieldKey(target)}_Faces`;

    constructor() {
        FaceRecognitionHandler._instance = this;
        this.loadModels().then(() => this._pendingLoadDocs.forEach(this.findMatches));
        DocumentManager.Instance.AddAnyViewRenderedCB(dv => FaceRecognitionHandler.Instance.findMatches(dv.Document));
    }

    @computed get examinedFaceDocs() {
        return DocListCast(Doc.UserDoc().examinedFaceDocs);
    }

    /**
     * Loads the face detection models.
     */
    loadModels = async () => {
        const MODEL_URL = `/models`;
        await faceapi.loadFaceDetectionModel(MODEL_URL);
        await faceapi.loadFaceLandmarkModel(MODEL_URL);
        await faceapi.loadFaceRecognitionModel(MODEL_URL);
        this._loadedModels = true;
    };

    public static get Instance() {
        return FaceRecognitionHandler._instance ?? new FaceRecognitionHandler();
    }

    /**
     * When a document is added, look for matching face documents.
     * @param doc The document being analyzed.
     */
    public findMatches = async (doc: Doc) => {
        if (!this._loadedModels || !Doc.ActiveDashboard) {
            this._pendingLoadDocs.push(doc);
            return;
        }

        if (doc.type === DocumentType.LOADING && !doc.loadingError) {
            setTimeout(() => this.findMatches(doc), 1000);
            return;
        }

        const imgUrl = ImageCast(doc[Doc.LayoutFieldKey(doc)]);
        // If the doc isn't an image or currently already been examined or is being processed, stop examining the document.
        if (!imgUrl || this.examinedFaceDocs.includes(doc) || this._processingDocs.has(doc)) {
            return;
        }

        // Mark the document as being processed.
        this._processingDocs.add(doc);

        // Get the image the document contains and analyze for faces.
        const [name, type] = imgUrl.url.href.split('.');
        const imageURL = `${name}_o.${type}`;

        const img = await this.loadImage(imageURL);

        const fullFaceDescriptions = await faceapi.detectAllFaces(img).withFaceLandmarks().withFaceDescriptors();

        doc[DocData][FaceRecognitionHandler.FacesField(doc)] = new List<List<number>>();

        // For each face detected, find a match.
        for (const fd of fullFaceDescriptions) {
            let match = this.findMatch(fd.descriptor);
            const converted_list = new List<number>(Array.from(fd.descriptor));

            if (match) {
                // If a matching Face Document has been found, add the document to the Face Document's associated docs and append the face
                // descriptor to the Face Document's descriptor list.
                Doc.AddDocToList(match, 'face_docList', doc);
                Cast(match.face_descriptors, listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
            } else {
                // If a matching Face Document has not been found, create a new Face Document.
                Doc.UserDoc().faceDocNum = NumCast(Doc.UserDoc().faceDocNum) + 1;

                const newFaceDocument = new Doc();
                newFaceDocument.title = `Face ${Doc.UserDoc().faceDocNum}`;
                newFaceDocument.face = ''; // just to make prettyprinting look better
                newFaceDocument.face_label = `Face${Doc.UserDoc().faceDocNum}`;
                newFaceDocument.face_docList = new List<Doc>([doc]);
                newFaceDocument.face_descriptors = new List<List<number>>([converted_list]);

                Doc.AddDocToList(Doc.ActiveDashboard[DocData], 'faceDocuments', newFaceDocument);
                match = newFaceDocument;
            }

            // Assign a field in the document of the matching Face Document.
            const faceDescripField = FaceRecognitionHandler.FaceField(doc, match);
            if (doc[DocData][faceDescripField]) {
                Cast(doc[DocData][faceDescripField], listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
            } else {
                doc[DocData][faceDescripField] = new List<List<number>>([converted_list]);
            }

            Cast(doc[DocData][FaceRecognitionHandler.FacesField(doc)], listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that

            Doc.AddDocToList(Doc.UserDoc(), 'examinedFaceDocs', doc);
        }
        this._processingDocs.delete(doc);
    };

    /**
     * Finds a matching Face Document given a descriptor
     * @param cur_descriptor The current descriptor whose match is being searched for.
     * @returns The most similar Face Document.
     */
    private findMatch(cur_descriptor: Float32Array) {
        if (!Doc.ActiveDashboard || DocListCast(Doc.ActiveDashboard[DocData].faceDocuments).length < 1) {
            return null;
        }

        const faceDescriptors: faceapi.LabeledFaceDescriptors[] = DocListCast(Doc.ActiveDashboard[DocData].faceDocuments).map(faceDocument => {
            const float32Array = (faceDocument[DocData].face_descriptors as List<List<number>>).map(faceDescriptor => new Float32Array(Array.from(faceDescriptor)));
            return new faceapi.LabeledFaceDescriptors(StrCast(faceDocument[DocData].face_label), float32Array);
        });
        const faceMatcher = new FaceMatcher(faceDescriptors, 0.6);
        const match = faceMatcher.findBestMatch(cur_descriptor);
        if (match.label !== 'unknown') {
            for (const doc of DocListCast(Doc.ActiveDashboard[DocData].faceDocuments)) {
                if (doc[DocData].face_label === match.label) {
                    return doc;
                }
            }
        }
        return null;
    }

    /**
     * Loads an image
     */
    private loadImage = (src: string): Promise<HTMLImageElement> => {
        return new Promise((resolve, reject) => {
            const img = new Image();
            img.crossOrigin = 'anonymous';
            img.onload = () => resolve(img);
            img.onerror = err => reject(err);
            img.src = src;
        });
    };
}