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 singleton class that handles face recognition and manages face Doc collections for each face found.
 * Displaying an image doc anywhere will trigger this class to test if the image contains any faces.
 * If it does, each recognized face will be compared to a global set of faces (each is a face collection Doc
 * that have already been found. If the face matches a face collection Doc, then it will be added to that
 * collection along with the numerical representation of the face, its face descriptor.
 *
 * Image Doc's that are added to one or more face collection Docs will be given these metadata fields:
 *     <image data field>_Face<N> - a nunerical representation of the Nth face found in the image
 *     <image data field>_Faces  - a list of all the numerical face representations found in the image (why is this needed?)
 *
 * Face collection Doc's are created for each person identified and are stored in the Dashboard's faceDocument's list
 *
 * Each Face collection Doc represents all the images found for that person.  It has these fields:
 *     face_label - a string label for the person that was recognized (currently it's just a 'face#')
 *     face_descriptors - a list of all the face descriptors for different images of the person
 *     face_docList - a list of all image Docs that contain a face for the person
 */
export class FaceRecognitionHandler {
    static _instance: FaceRecognitionHandler;
    private _loadedModels: boolean = false;
    private _processingDocs: Set<Doc> = new Set();
    private _pendingLoadDocs: Doc[] = [];

    private static imgDocFaceField = (imgDoc: Doc, faceDoc: Doc) => `${Doc.LayoutFieldKey(imgDoc)}_${FaceRecognitionHandler.FaceDocLabel(faceDoc)}`;
    /**
     * initializes an image with an empty list of face descriptors
     * @param imgDoc image to initialize
     */
    private static initImageDocFaceDescriptors = (imgDoc: Doc) => {
        imgDoc[DocData][`${Doc.LayoutFieldKey(imgDoc)}_Faces`] = new List<List<number>>();
    };
    /**
     * returns the face descriptors for each face found on an image Doc
     * @param imgDoc
     * @returns list of face descriptors
     */
    public static ImageDocFaceDescriptors = (imgDoc: Doc) => imgDoc[DocData][`${Doc.LayoutFieldKey(imgDoc)}_Faces`] as List<List<number>>;

    /**
     * Adds metadata to an image Doc describing a face found in the image
     * @param imgDoc image Doc containing faces
     * @param faceDescriptor descriptor for the face found
     * @param faceDoc face collection Doc containing the same face
     */
    public static ImageDocAddFace = (imgDoc: Doc, faceDescriptor: List<number>, faceDoc: Doc) => {
        const faceFieldKey = FaceRecognitionHandler.imgDocFaceField(imgDoc, faceDoc);
        if (imgDoc[DocData][faceFieldKey]) {
            Cast(imgDoc[DocData][faceFieldKey], listSpec('number'), null).push(faceDescriptor as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
        } else {
            imgDoc[DocData][faceFieldKey] = new List<List<number>>([faceDescriptor]);
        }
    };

    /**
     * returns a list of all face collection Docs on the current dashboard
     * @returns face collection Doc list
     */
    public static FaceDocuments = () => DocListCast(Doc.ActiveDashboard?.[DocData].faceDocuments);

    public static DeleteFaceDoc = (faceDoc: Doc) => Doc.ActiveDashboard && Doc.RemoveDocFromList(Doc.ActiveDashboard[DocData], 'faceDocuments', faceDoc);

    /**
     * returns the labels associated with a face collection Doc
     * @param faceDoc the face collection Doc
     * @returns label string
     */
    public static FaceDocLabel = (faceDoc: Doc) => StrCast(faceDoc[DocData].face_label);
    /**
     * Returns all the face descriptors associated with a face collection Doc
     * @param faceDoc a face collection Doc
     * @returns face descriptors
     */
    public static FaceDocDescriptors = (faceDoc: Doc) => faceDoc[DocData].face_descriptors as List<List<number>>;

    /**
     * Returns a list of all face image Docs associated with the face collection
     * @param faceDoc a face collection Doc
     * @returns image Docs
     */
    public static FaceDocFaces = (faceDoc: Doc) => DocListCast(faceDoc[DocData].face_docList);

    /**
     * Adds a face image to the list of faces in a face collection Doc, and updates the face collection's list of image descriptors
     * @param img - image with faces to add to a face collection Doc
     * @param faceDescriptor - the face descriptor for the face in the image to add
     * @param faceDoc - the face collection Doc
     */
    public static FaceDocAddImageDocFace = (img: Doc, faceDescriptor: List<number>, faceDoc: Doc) => {
        Doc.AddDocToList(faceDoc, 'face_docList', img);
        Cast(faceDoc.face_descriptors, listSpec('number'), null).push(faceDescriptor as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
    };

    /**
     * Removes a face from a face Doc collection, and updates the face collection's list of image descriptors
     * @param imgDoc - image with faces to remove from the face Doc collectoin
     * @param faceDoc - the face Doc collection
     */
    public static FaceDocRemoveImageDocFace = (imgDoc: Doc, faceDoc: Doc) => {
        imgDoc[DocData][FaceRecognitionHandler.imgDocFaceField(imgDoc, faceDoc)] = new List<List<number>>();
        Doc.RemoveDocFromList(faceDoc[DocData], 'face_docList', imgDoc);
        faceDoc[DocData].face_descriptors = new List<List<number>>(FaceRecognitionHandler.FaceDocDescriptors(faceDoc).filter(fd => !(imgDoc[DocData][FaceRecognitionHandler.imgDocFaceField(imgDoc, faceDoc)] as List<List<number>>).includes(fd)));
    };

    constructor() {
        FaceRecognitionHandler._instance = this;
        this.loadModels().then(() => this._pendingLoadDocs.forEach(this.classifyFacesInImage));
        DocumentManager.Instance.AddAnyViewRenderedCB(dv => FaceRecognitionHandler.Instance.classifyFacesInImage(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 imgDoc The document being analyzed.
     */
    public classifyFacesInImage = async (imgDoc: Doc) => {
        if (!this._loadedModels || !Doc.ActiveDashboard) {
            this._pendingLoadDocs.push(imgDoc);
            return;
        }

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

        const imgUrl = ImageCast(imgDoc[Doc.LayoutFieldKey(imgDoc)]);
        // 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(imgDoc) || this._processingDocs.has(imgDoc)) {
            return;
        }

        // Mark the document as being processed.
        this._processingDocs.add(imgDoc);
        FaceRecognitionHandler.initImageDocFaceDescriptors(imgDoc);

        // 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 imgDocFaceDescriptions = await faceapi.detectAllFaces(img).withFaceLandmarks().withFaceDescriptors();

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

            if (faceDocMatch) {
                FaceRecognitionHandler.FaceDocAddImageDocFace(imgDoc, faceDescriptor, faceDocMatch);
            } 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>([imgDoc]);
                newFaceDocument.face_descriptors = new List<List<number>>([faceDescriptor]);

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

            // Assign a field in the document of the matching Face Document.
            FaceRecognitionHandler.ImageDocAddFace(imgDoc, faceDescriptor, faceDocMatch);
            Doc.AddDocToList(Doc.UserDoc(), 'examinedFaceDocs', imgDoc);
        }
        this._processingDocs.delete(imgDoc);
    };

    /**
     * Finds the most similar matching Face Document to a face descriptor
     * @param faceDescriptor face descriptor number list
     * @returns face Doc
     */
    private findMatchingFaceDoc = (faceDescriptor: Float32Array) => {
        if (!Doc.ActiveDashboard || FaceRecognitionHandler.FaceDocuments().length < 1) {
            return undefined;
        }

        const faceDescriptors = FaceRecognitionHandler.FaceDocuments().map(faceDoc => {
            const float32Array = FaceRecognitionHandler.FaceDocDescriptors(faceDoc).map(fd => new Float32Array(Array.from(fd)));
            return new faceapi.LabeledFaceDescriptors(FaceRecognitionHandler.FaceDocLabel(faceDoc), float32Array);
        });
        const faceMatcher = new FaceMatcher(faceDescriptors, 0.6);
        const match = faceMatcher.findBestMatch(faceDescriptor);
        if (match.label !== 'unknown') {
            for (const faceDoc of FaceRecognitionHandler.FaceDocuments()) {
                if (FaceRecognitionHandler.FaceDocLabel(faceDoc) === match.label) {
                    return faceDoc;
                }
            }
        }
        return undefined;
    };

    /**
     * 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;
        });
    };
}